Merge remote-tracking branch 'upstream/main'

This commit is contained in:
xlplbo
2026-07-10 20:46:38 +08:00
124 changed files with 5597 additions and 577 deletions
+1 -11
View File
@@ -1,4 +1,4 @@
.PHONY: build build-backend build-frontend build-datamanagementd test test-backend test-frontend test-frontend-critical test-datamanagementd secret-scan
.PHONY: build build-backend build-frontend test test-backend test-frontend test-frontend-critical
FRONTEND_CRITICAL_VITEST := \
src/views/auth/__tests__/LinuxDoCallbackView.spec.ts \
@@ -19,10 +19,6 @@ build-backend:
build-frontend:
@pnpm --dir frontend run build
# 编译 datamanagementd(宿主机数据管理进程)
build-datamanagementd:
@cd datamanagement && go build -o datamanagementd ./cmd/datamanagementd
# 运行测试(后端 + 前端)
test: test-backend test-frontend
@@ -36,9 +32,3 @@ test-frontend:
test-frontend-critical:
@pnpm --dir frontend exec vitest run $(FRONTEND_CRITICAL_VITEST)
test-datamanagementd:
@cd datamanagement && go test ./...
secret-scan:
@python3 tools/secret_scan.py
+1 -1
View File
@@ -1 +1 @@
0.1.149
0.1.150
@@ -265,10 +265,12 @@ func TestOpenAIFastPolicySettingsFromDTO_NormalizesServiceTier(t *testing.T) {
ServiceTier: "PRIORITY",
Action: "filter",
Scope: "all",
UserIDs: []int64{42},
}},
}
out := openaiFastPolicySettingsFromDTO(in)
require.Equal(t, service.OpenAIFastTierPriority, out.Rules[0].ServiceTier)
require.Equal(t, []int64{42}, out.Rules[0].UserIDs)
})
t.Run("non-empty values pass through (lowercased)", func(t *testing.T) {
@@ -657,11 +657,14 @@ func (h *DashboardHandler) GetUserBreakdown(c *gin.Context) {
dim.AccountID = id
}
}
if v := c.Query("request_type"); v != "" {
if rt, err := strconv.ParseInt(v, 10, 16); err == nil {
rtVal := int16(rt)
dim.RequestType = &rtVal
if v := strings.TrimSpace(c.Query("request_type")); v != "" {
parsed, err := service.ParseUsageRequestType(v)
if err != nil {
response.BadRequest(c, err.Error())
return
}
rtVal := int16(parsed)
dim.RequestType = &rtVal
}
if v := c.Query("stream"); v != "" {
if s, err := strconv.ParseBool(v); err == nil {
@@ -241,3 +241,42 @@ func TestGetUserBreakdown_NoFilters(t *testing.T) {
require.Empty(t, repo.capturedDim.Model)
require.Empty(t, repo.capturedDim.Endpoint)
}
func TestGetUserBreakdown_RequestTypeStringFilter(t *testing.T) {
cases := []struct {
name string
value string
want int16
}{
{"ws_v2", "ws_v2", int16(service.RequestTypeWSV2)},
{"stream", "stream", int16(service.RequestTypeStream)},
{"sync", "sync", int16(service.RequestTypeSync)},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
repo := &userBreakdownRepoCapture{}
router := newUserBreakdownRouter(repo)
req := httptest.NewRequest(http.MethodGet,
"/admin/dashboard/user-breakdown?start_date=2026-03-01&end_date=2026-03-16&request_type="+tc.value, nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusOK, w.Code)
require.NotNil(t, repo.capturedDim.RequestType, "request_type=%s should set filter", tc.value)
require.Equal(t, tc.want, *repo.capturedDim.RequestType)
})
}
}
func TestGetUserBreakdown_InvalidRequestType(t *testing.T) {
repo := &userBreakdownRepoCapture{}
router := newUserBreakdownRouter(repo)
req := httptest.NewRequest(http.MethodGet,
"/admin/dashboard/user-breakdown?start_date=2026-03-01&end_date=2026-03-16&request_type=bogus", nil)
w := httptest.NewRecorder()
router.ServeHTTP(w, req)
require.Equal(t, http.StatusBadRequest, w.Code)
}
+1
View File
@@ -422,6 +422,7 @@ type OpenAIFastPolicyRule struct {
ServiceTier string `json:"service_tier"`
Action string `json:"action"`
Scope string `json:"scope"`
UserIDs []int64 `json:"user_ids,omitempty"`
ErrorMessage string `json:"error_message,omitempty"`
ModelWhitelist []string `json:"model_whitelist,omitempty"`
FallbackAction string `json:"fallback_action,omitempty"`
@@ -205,6 +205,11 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
if !ok {
return
}
// body-signal compact:上游 unary 等待期间向下游发 SSE 注释行心跳,防止
// 反向代理空闲超时掐断长压缩连接(#3887)。首拍延迟一个心跳间隔,快速
// 失败仍走 JSON+状态码链路;未标记客户端流式或间隔为 0 时是 no-op。
stopCompactKeepalive := service.StartOpenAICompactSSEKeepalive(c, h.openAICompactKeepaliveInterval())
defer stopCompactKeepalive()
// 校验请求体 JSON 合法性
if !gjson.ValidBytes(body) {
@@ -402,7 +407,9 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
// Forward request
service.SetOpsLatencyMs(c, service.OpsRoutingLatencyMsKey, time.Since(routingStart).Milliseconds())
forwardStart := time.Now()
writerSizeBeforeForward := c.Writer.Size()
// 用扣除 compact 心跳字节的口径快照:心跳注释不构成语义响应,
// 不能因心跳字节变化而放弃 failover 换号(#3887)。
writerSizeBeforeForward := service.OpenAICompactKeepaliveAdjustedWrittenSize(c)
result, err := func() (*service.OpenAIForwardResult, error) {
defer func() {
if accountReleaseFunc != nil {
@@ -436,7 +443,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
} else {
var failoverErr *service.UpstreamFailoverError
if errors.As(err, &failoverErr) {
if c.Writer.Size() != writerSizeBeforeForward {
if service.OpenAICompactKeepaliveAdjustedWrittenSize(c) != writerSizeBeforeForward {
h.handleFailoverExhausted(c, failoverErr, true)
return
}
@@ -637,6 +644,13 @@ func (h *OpenAIGatewayHandler) logOpenAIRemoteCompactOutcome(c *gin.Context, sta
if status >= 200 && status < 300 {
outcome = "succeeded"
}
// compact 心跳提交后失败的 wire 状态码固化为 200,真实结局以流内错误
// 标记为准(response.failed 降级路径会 MarkOpsStreamError)。
if outcome == "succeeded" && c != nil {
if _, hasStreamErr := service.GetOpsStreamError(c); hasStreamErr {
outcome = "failed"
}
}
latencyMs := time.Since(startedAt).Milliseconds()
if latencyMs < 0 {
latencyMs = 0
@@ -1943,6 +1957,11 @@ func (h *OpenAIGatewayHandler) mapUpstreamError(statusCode int) (int, string, st
// handleStreamingAwareError handles errors that may occur after streaming has started
func (h *OpenAIGatewayHandler) handleStreamingAwareError(c *gin.Context, status int, errType, message string, streamStarted bool) {
// body-signal compact 心跳可能已把响应头提交为 200:先停心跳(建立
// happens-before,接管 ResponseWriter),并升级为流内错误处理。
if service.StopOpenAICompactSSEKeepaliveCommitted(c) {
streamStarted = true
}
if streamStarted {
// /v1/responses 的严格 SDK(Codex CLI)要求终止事件必须属于
// response.completed/failed/incomplete/cancelled 集合。
@@ -1975,6 +1994,10 @@ func (h *OpenAIGatewayHandler) ensureForwardErrorResponse(c *gin.Context, stream
if c == nil || c.Writer == nil {
return false
}
// 先停 compact 心跳再读 Writer 状态,避免与心跳 goroutine 竞争。
if service.StopOpenAICompactSSEKeepaliveCommitted(c) {
streamStarted = true
}
if service.IsResponseCommitted(c) {
return false
}
@@ -2010,7 +2033,9 @@ func openAIForwardErrorAlreadyCommunicated(c *gin.Context, writerSizeBeforeForwa
if err == nil || c == nil || c.Writer == nil {
return false
}
if c.Writer.Size() == writerSizeBeforeForward {
// 与快照同口径:排除 compact 心跳字节,避免"仅心跳写出"被误判为
// 响应已写出(#3887)。
if service.OpenAICompactKeepaliveAdjustedWrittenSize(c) == writerSizeBeforeForward {
return false
}
@@ -2036,6 +2061,14 @@ func openAIForwardErrorAlreadyCommunicated(c *gin.Context, writerSizeBeforeForwa
// errorResponse returns OpenAI API format error response
func (h *OpenAIGatewayHandler) errorResponse(c *gin.Context, status int, errType, message string) {
// body-signal compact 心跳可能已把响应头提交为 200:JSON 错误体会与已
// 提交的 SSE 流交错,必须降级为 response.failed 终止事件(#3887)。
if service.StopOpenAICompactSSEKeepaliveCommitted(c) {
service.MarkOpsStreamError(c, errType, message, status)
if writeResponsesFailedSSE(c, errType, message) {
return
}
}
c.JSON(status, gin.H{
"error": gin.H{
"type": errType,
@@ -2044,6 +2077,15 @@ func (h *OpenAIGatewayHandler) errorResponse(c *gin.Context, status int, errType
})
}
// openAICompactKeepaliveInterval 复用流式 keepalive 配置作为 compact 下游
// 心跳间隔;0 表示禁用(与流式路径语义一致)。
func (h *OpenAIGatewayHandler) openAICompactKeepaliveInterval() time.Duration {
if h.cfg == nil || h.cfg.Gateway.StreamKeepaliveInterval <= 0 {
return 0
}
return time.Duration(h.cfg.Gateway.StreamKeepaliveInterval) * time.Second
}
func setOpenAIClientTransportHTTP(c *gin.Context) {
service.SetOpenAIClientTransport(c, service.OpenAIClientTransportHTTP)
}
@@ -2306,6 +2348,16 @@ func (h *OpenAIGatewayHandler) rejectIfCyberSessionBlocked(c *gin.Context, apiKe
if !h.gatewayService.IsCyberSessionBlocked(c.Request.Context(), key) {
return false
}
// body-signal compact 心跳可能已把响应头提交为 200(cyber 检查在用户槽位
// 长等待之后执行):以 response.failed 终止事件回传;未提交时停拍后照常
// 写 JSON(#3887)。
if service.StopOpenAICompactSSEKeepaliveCommitted(c) {
service.MarkOpsStreamError(c, "permission_error", cyberSessionBlockedClientMsg, http.StatusForbidden)
if writeResponsesFailedSSE(c, "permission_error", cyberSessionBlockedClientMsg) {
h.enqueueCyberSessionBlockedOpsEntry(c, apiKey, model, key)
return true
}
}
switch format {
case cyberBlockFormatAnthropic:
c.JSON(http.StatusForbidden, gin.H{"type": "error", "error": gin.H{
@@ -30,6 +30,7 @@ func ResponsesToChatCompletionsRequest(req *ResponsesRequest) (*ChatCompletionsR
TopP: req.TopP,
Stream: req.Stream,
ServiceTier: req.ServiceTier,
ParallelToolCalls: req.ParallelToolCalls,
}
if req.Reasoning != nil {
out.ReasoningEffort = req.Reasoning.Effort
@@ -934,9 +935,17 @@ func ChatUsageToResponsesUsage(usage *ChatUsage) *ResponsesUsage {
if out.TotalTokens == 0 {
out.TotalTokens = out.InputTokens + out.OutputTokens
}
if usage.PromptTokensDetails != nil && usage.PromptTokensDetails.CachedTokens > 0 {
if usage.PromptTokensDetails != nil && (usage.PromptTokensDetails.CachedTokens > 0 ||
usage.PromptTokensDetails.CacheCreationTokens > 0 || usage.PromptTokensDetails.CacheWriteTokens > 0) {
out.InputTokensDetails = &ResponsesInputTokensDetails{
CachedTokens: usage.PromptTokensDetails.CachedTokens,
CachedTokens: usage.PromptTokensDetails.CachedTokens,
CacheCreationTokens: usage.PromptTokensDetails.CacheCreationTokens,
CacheWriteTokens: usage.PromptTokensDetails.CacheWriteTokens,
}
if usage.PromptTokensDetails.CacheWriteTokens > 0 {
out.CacheCreationInputTokens = usage.PromptTokensDetails.CacheWriteTokens
} else {
out.CacheCreationInputTokens = usage.PromptTokensDetails.CacheCreationTokens
}
}
return out
@@ -127,6 +127,26 @@ func TestResponsesToChatCompletionsRequest_TextFormatJsonSchema(t *testing.T) {
}`, string(out.ResponseFormat))
}
func TestResponsesToChatCompletionsRequest_ParallelToolCalls(t *testing.T) {
parallel := false
req := &ResponsesRequest{
Model: "gpt-4o",
Input: json.RawMessage(`[
{"role":"user","content":"Use tools"}
]`),
ParallelToolCalls: &parallel,
}
out, err := ResponsesToChatCompletionsRequest(req)
require.NoError(t, err)
require.NotNil(t, out.ParallelToolCalls)
assert.False(t, *out.ParallelToolCalls)
payload, err := json.Marshal(out)
require.NoError(t, err)
assert.Contains(t, string(payload), `"parallel_tool_calls":false`)
}
func chatMessageRoles(messages []ChatMessage) []string {
roles := make([]string, 0, len(messages))
for _, message := range messages {
@@ -32,6 +32,47 @@ func TestChatCompletionsToResponses_BasicText(t *testing.T) {
assert.Equal(t, "user", items[0].Role)
}
func TestUsageConversionsPreserveCacheWriteTokens(t *testing.T) {
var responsesUsage ResponsesUsage
require.NoError(t, json.Unmarshal([]byte(`{
"input_tokens":1000,
"output_tokens":50,
"input_tokens_details":{"cached_tokens":100,"cache_write_tokens":200}
}`), &responsesUsage))
require.NotNil(t, responsesUsage.InputTokensDetails)
require.Equal(t, 200, responsesUsage.InputTokensDetails.CacheWriteTokens)
chatUsage := chatUsageFromResponsesUsage(&responsesUsage)
require.NotNil(t, chatUsage.PromptTokensDetails)
require.Equal(t, 100, chatUsage.PromptTokensDetails.CachedTokens)
require.Equal(t, 200, chatUsage.PromptTokensDetails.CacheWriteTokens)
roundTrip := ChatUsageToResponsesUsage(chatUsage)
require.NotNil(t, roundTrip.InputTokensDetails)
require.Equal(t, 200, roundTrip.CacheCreationInputTokens)
require.Equal(t, 200, roundTrip.InputTokensDetails.CacheWriteTokens)
}
func TestResponsesUsageNestedCacheWritePresenceOverridesTopLevelAlias(t *testing.T) {
tests := []struct {
name string
nestedJSON string
want int
}{
{name: "explicit zero", nestedJSON: `{"cache_write_tokens":0}`, want: 0},
{name: "nonzero", nestedJSON: `{"cache_write_tokens":7}`, want: 7},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
var usage ResponsesUsage
payload := []byte(`{"input_tokens":20,"output_tokens":2,"cache_creation_input_tokens":19,"input_tokens_details":` + tt.nestedJSON + `}`)
require.NoError(t, json.Unmarshal(payload, &usage))
require.Equal(t, tt.want, usage.CacheCreationInputTokens)
})
}
}
func TestChatCompletionsToResponses_SystemMessage(t *testing.T) {
req := &ChatCompletionsRequest{
Model: "gpt-4o",
@@ -531,6 +572,25 @@ func TestChatCompletionsToResponses_ServiceTier(t *testing.T) {
assert.Equal(t, "flex", resp.ServiceTier)
}
func TestChatCompletionsToResponses_ParallelToolCalls(t *testing.T) {
for _, value := range []bool{false, true} {
req := &ChatCompletionsRequest{
Model: "gpt-4o",
ParallelToolCalls: &value,
Messages: []ChatMessage{{Role: "user", Content: json.RawMessage(`"Hi"`)}},
}
resp, err := ChatCompletionsToResponses(req)
require.NoError(t, err)
require.NotNil(t, resp.ParallelToolCalls)
assert.Equal(t, value, *resp.ParallelToolCalls)
payload, err := json.Marshal(resp)
require.NoError(t, err)
assert.Contains(t, string(payload), `"parallel_tool_calls":`+string(mustMarshalJSON(t, value)))
}
}
// ---------------------------------------------------------------------------
// temperature / top_p stripping for reasoning models
// ---------------------------------------------------------------------------
@@ -27,12 +27,13 @@ func ChatCompletionsToResponses(req *ChatCompletionsRequest) (*ResponsesRequest,
}
out := &ResponsesRequest{
Model: req.Model,
Instructions: req.Instructions,
Input: inputJSON,
Stream: true, // upstream always streams
Include: []string{"reasoning.encrypted_content"},
ServiceTier: req.ServiceTier,
Model: req.Model,
Instructions: req.Instructions,
Input: inputJSON,
Stream: true, // upstream always streams
Include: []string{"reasoning.encrypted_content"},
ServiceTier: req.ServiceTier,
ParallelToolCalls: req.ParallelToolCalls,
}
// Reasoning models (gpt-5.x) do not accept sampling parameters.
@@ -338,6 +338,14 @@ func chatUsageFromResponsesUsage(u *ResponsesUsage) *ChatUsage {
TotalTokens: u.InputTokens + u.OutputTokens,
}
usage.PromptTokensDetails = promptDetailsFromResponses(u.InputTokensDetails)
if u.CacheCreationInputTokens > 0 {
if usage.PromptTokensDetails == nil {
usage.PromptTokensDetails = &ChatTokenDetails{}
}
if usage.PromptTokensDetails.CacheWriteTokens == 0 && usage.PromptTokensDetails.CacheCreationTokens == 0 {
usage.PromptTokensDetails.CacheCreationTokens = u.CacheCreationInputTokens
}
}
usage.CompletionTokensDetails = completionDetailsFromResponses(u.OutputTokensDetails)
return usage
}
@@ -349,12 +357,14 @@ func promptDetailsFromResponses(src *ResponsesInputTokensDetails) *ChatTokenDeta
if src == nil {
return nil
}
if src.CachedTokens == 0 && src.AudioTokens == 0 {
if src.CachedTokens == 0 && src.AudioTokens == 0 && src.CacheCreationTokens == 0 && src.CacheWriteTokens == 0 {
return nil
}
return &ChatTokenDetails{
CachedTokens: src.CachedTokens,
AudioTokens: src.AudioTokens,
CachedTokens: src.CachedTokens,
AudioTokens: src.AudioTokens,
CacheCreationTokens: src.CacheCreationTokens,
CacheWriteTokens: src.CacheWriteTokens,
}
}
+49 -5
View File
@@ -367,9 +367,10 @@ type ResponsesSummary struct {
// ResponsesUsage holds token counts in Responses API format.
type ResponsesUsage struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
TotalTokens int `json:"total_tokens"`
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
TotalTokens int `json:"total_tokens"`
CacheCreationInputTokens int `json:"cache_creation_input_tokens,omitempty"`
// Optional detailed breakdown
InputTokensDetails *ResponsesInputTokensDetails `json:"input_tokens_details,omitempty"`
@@ -378,16 +379,30 @@ type ResponsesUsage struct {
func (u *ResponsesUsage) UnmarshalJSON(data []byte) error {
type responsesUsageAlias ResponsesUsage
type cacheTokenPresence struct {
CacheCreationTokens *int `json:"cache_creation_tokens"`
CacheWriteTokens *int `json:"cache_write_tokens"`
}
var aux struct {
responsesUsageAlias
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
CacheCreationTokens int `json:"cache_creation_tokens"`
CacheWriteInputTokens int `json:"cache_write_input_tokens"`
CacheWriteTokens int `json:"cache_write_tokens"`
PromptTokensDetails *ResponsesInputTokensDetails `json:"prompt_tokens_details,omitempty"`
CompletionTokensDetails *ResponsesOutputTokensDetails `json:"completion_tokens_details,omitempty"`
}
if err := json.Unmarshal(data, &aux); err != nil {
return err
}
var nestedPresence struct {
InputTokensDetails *cacheTokenPresence `json:"input_tokens_details"`
PromptTokensDetails *cacheTokenPresence `json:"prompt_tokens_details"`
}
if err := json.Unmarshal(data, &nestedPresence); err != nil {
return err
}
*u = ResponsesUsage(aux.responsesUsageAlias)
if u.InputTokens == 0 && aux.PromptTokens != 0 {
u.InputTokens = aux.PromptTokens
@@ -395,12 +410,36 @@ func (u *ResponsesUsage) UnmarshalJSON(data []byte) error {
if u.OutputTokens == 0 && aux.CompletionTokens != 0 {
u.OutputTokens = aux.CompletionTokens
}
if u.CacheCreationInputTokens == 0 {
switch {
case aux.CacheWriteInputTokens > 0:
u.CacheCreationInputTokens = aux.CacheWriteInputTokens
case aux.CacheCreationTokens > 0:
u.CacheCreationInputTokens = aux.CacheCreationTokens
case aux.CacheWriteTokens > 0:
u.CacheCreationInputTokens = aux.CacheWriteTokens
}
}
if u.InputTokensDetails == nil && aux.PromptTokensDetails != nil {
u.InputTokensDetails = aux.PromptTokensDetails
}
if u.OutputTokensDetails == nil && aux.CompletionTokensDetails != nil {
u.OutputTokensDetails = aux.CompletionTokensDetails
}
var canonicalCacheCreationTokens *int
switch {
case nestedPresence.InputTokensDetails != nil && nestedPresence.InputTokensDetails.CacheWriteTokens != nil:
canonicalCacheCreationTokens = nestedPresence.InputTokensDetails.CacheWriteTokens
case nestedPresence.PromptTokensDetails != nil && nestedPresence.PromptTokensDetails.CacheWriteTokens != nil:
canonicalCacheCreationTokens = nestedPresence.PromptTokensDetails.CacheWriteTokens
case nestedPresence.InputTokensDetails != nil && nestedPresence.InputTokensDetails.CacheCreationTokens != nil:
canonicalCacheCreationTokens = nestedPresence.InputTokensDetails.CacheCreationTokens
case nestedPresence.PromptTokensDetails != nil && nestedPresence.PromptTokensDetails.CacheCreationTokens != nil:
canonicalCacheCreationTokens = nestedPresence.PromptTokensDetails.CacheCreationTokens
}
if canonicalCacheCreationTokens != nil {
u.CacheCreationInputTokens = max(*canonicalCacheCreationTokens, 0)
}
if u.TotalTokens == 0 && (u.InputTokens != 0 || u.OutputTokens != 0) {
u.TotalTokens = u.InputTokens + u.OutputTokens
}
@@ -409,8 +448,10 @@ func (u *ResponsesUsage) UnmarshalJSON(data []byte) error {
// ResponsesInputTokensDetails breaks down input token usage.
type ResponsesInputTokensDetails struct {
CachedTokens int `json:"cached_tokens,omitempty"`
AudioTokens int `json:"audio_tokens,omitempty"`
CachedTokens int `json:"cached_tokens,omitempty"`
AudioTokens int `json:"audio_tokens,omitempty"`
CacheCreationTokens int `json:"cache_creation_tokens,omitempty"`
CacheWriteTokens int `json:"cache_write_tokens,omitempty"`
}
// ResponsesOutputTokensDetails breaks down output token usage.
@@ -485,6 +526,7 @@ type ChatCompletionsRequest struct {
Stream bool `json:"stream,omitempty"`
StreamOptions *ChatStreamOptions `json:"stream_options,omitempty"`
Tools []ChatTool `json:"tools,omitempty"`
ParallelToolCalls *bool `json:"parallel_tool_calls,omitempty"`
ToolChoice json.RawMessage `json:"tool_choice,omitempty"`
ReasoningEffort string `json:"reasoning_effort,omitempty"` // "low" | "medium" | "high" | "xhigh"
ServiceTier string `json:"service_tier,omitempty"`
@@ -595,6 +637,8 @@ type ChatUsage struct {
type ChatTokenDetails struct {
CachedTokens int `json:"cached_tokens,omitempty"`
AudioTokens int `json:"audio_tokens,omitempty"`
CacheCreationTokens int `json:"cache_creation_tokens,omitempty"`
CacheWriteTokens int `json:"cache_write_tokens,omitempty"`
ReasoningTokens int `json:"reasoning_tokens,omitempty"`
AcceptedPredictionTokens int `json:"accepted_prediction_tokens,omitempty"`
RejectedPredictionTokens int `json:"rejected_prediction_tokens,omitempty"`
+4
View File
@@ -41,6 +41,10 @@ const (
// Group 认证后的分组信息,由 API Key 认证中间件设置
Group Key = "ctx_group"
// UserID 认证后的 Sub2API 用户 ID,由 API Key 认证中间件设置。
// 供 service 层执行用户级策略,不能使用客户端请求体中的 user 标识替代。
UserID Key = "ctx_user_id"
// IsMaxTokensOneHaikuRequest 标识当前请求是否为 max_tokens=1 + haiku 模型的探测请求
// 用于 ClaudeCodeOnly 验证绕过(绕过 system prompt 检查,但仍需验证 User-Agent)
IsMaxTokensOneHaikuRequest Key = "ctx_is_max_tokens_one_haiku"
+1
View File
@@ -18,6 +18,7 @@ type Model struct {
// DefaultModels OpenAI models list
var DefaultModels = []Model{
{ID: "gpt-5.6", Object: "model", Created: 1780876800, OwnedBy: "openai", Type: "model", DisplayName: "GPT-5.6 (Sol)"},
{ID: "gpt-5.6-sol", Object: "model", Created: 1780876800, OwnedBy: "openai", Type: "model", DisplayName: "GPT-5.6 Sol"},
{ID: "gpt-5.6-terra", Object: "model", Created: 1780876800, OwnedBy: "openai", Type: "model", DisplayName: "GPT-5.6 Terra"},
{ID: "gpt-5.6-luna", Object: "model", Created: 1780876800, OwnedBy: "openai", Type: "model", DisplayName: "GPT-5.6 Luna"},
@@ -0,0 +1,11 @@
package openai
import (
"testing"
"github.com/stretchr/testify/require"
)
func TestDefaultModelsIncludeBareGPT56Alias(t *testing.T) {
require.Contains(t, DefaultModelIDs(), "gpt-5.6")
}
+61 -1
View File
@@ -127,7 +127,9 @@ func isCodexOfficialClientRequest(userAgent string, strict bool) bool {
// `(name; version)` 括号组——该组由 codex-rs engine 写入,保留真实 clientInfo.name。
// 故从尾部提取 name 可以恢复被 override 的真实客户端标识(例如 cccc → codex-tui)。
//
// input 应为已归一化(小写 + 去首尾空格)的 UA。
// input 应为去首尾空格的 UA;本函数本身大小写无关,大小写由调用方按需处理
// (isCodexOfficialClientRequest 传入已小写化的 UA 做匹配;PairCodexClientIdentity
// 传入原始大小写以保留 originator 的真实大小写)。
// 若无法解析则返回空字符串。
func codexUATrailerName(ua string) string {
last := strings.LastIndex(ua, "(")
@@ -195,6 +197,64 @@ func matchCodexClientHeaderStrictPrefixes(value string, prefixes []string) bool
return false
}
// PairCodexClientIdentity 由最终出站 User-Agent 推导与其配套的 originator,必要时归一化
// UA 首段,保证两者一致。上游 /backend-api/codex 会校验 originator 与 UA 首段(首个 '/'
// 之前的 client 名)是否配套,错配(如 originator=codex_cli_rs + UA=codex-tui/...)一律
// 404(issue #3901,2026-07 实测)。
//
// 推导优先级:
// 1. UA 首段是官方 originator(精确集合或 `Codex ` 家族前缀)→ 直接配对,UA 原样保留;
// 2. UA 尾部括号组 `(name; version)` 的 name 是官方 originator——CODEX_INTERNAL_ORIGINATOR_OVERRIDE
// 只改 UA 前缀不改尾部(如 cccc/0.142.0 ... (codex-tui; 0.142.0))→ 用尾部 name 重写
// UA 首段后配对,保留真实版本/OS/终端指纹;
// 3. 均不命中 → ok=false,调用方应整体回退为默认官方身份。
func PairCodexClientIdentity(userAgent string) (originator string, pairedUA string, ok bool) {
ua := strings.TrimSpace(userAgent)
slash := strings.IndexByte(ua, '/')
if slash <= 0 {
return "", "", false
}
if leading := strings.TrimSpace(ua[:slash]); isSaneCodexOriginator(leading) && IsCodexOfficialClientOriginator(leading) {
leading = canonicalizeCodexOriginator(leading)
return leading, leading + ua[slash:], true
}
// 传原始大小写 UA 提取 trailer,保留 `Codex ` 家族身份的真实大小写;含 '/' 的
// trailer 会破坏重写后 UA 首段与 originator 的一致性,直接拒绝。
if trailer := codexUATrailerName(ua); trailer != "" && !strings.ContainsRune(trailer, '/') &&
isSaneCodexOriginator(trailer) && IsCodexOfficialClientOriginator(trailer) {
trailer = canonicalizeCodexOriginator(trailer)
return trailer, trailer + ua[slash:], true
}
return "", "", false
}
// codexOriginatorMaxLen 官方 clientInfo.name 均为短 ASCII 标识,远低于此上限。
const codexOriginatorMaxLen = 64
// isSaneCodexOriginator 拒绝超长或含不可打印/非 ASCII 字节的候选 originator,
// 避免 `Codex ` 家族宽前缀把客户端可控的任意字节当作官方身份逐字转发给上游。
func isSaneCodexOriginator(name string) bool {
if name == "" || len(name) > codexOriginatorMaxLen {
return false
}
for i := 0; i < len(name); i++ {
if c := name[i]; c < 0x20 || c > 0x7e {
return false
}
}
return true
}
// canonicalizeCodexOriginator 把精确集合的官方 originator 大小写变体归一为规范小写形态
// (如 CODEX_CLI_RS → codex_cli_rs);`Codex ` 家族不在精确集合中,保留原大小写
// (其规范形态本就是混合大小写,上游按大小写敏感 starts_with("Codex ") 判定)。
func canonicalizeCodexOriginator(name string) string {
if lower := normalizeCodexClientHeader(name); codexOfficialClientOriginators[lower] {
return lower
}
return name
}
// codexEngineVersionPattern 提取版本段开头的三段数字 X.Y.Z(忽略 -alpha 等后缀)。
var codexEngineVersionPattern = regexp.MustCompile(`^(\d+\.\d+\.\d+)`)
@@ -0,0 +1,86 @@
package openai
import (
"strings"
"testing"
"github.com/stretchr/testify/require"
)
func TestPairCodexClientIdentity(t *testing.T) {
tests := []struct {
name string
ua string
wantOriginator string
wantUA string
wantOK bool
}{
{
name: "cli 首段直接配对",
ua: "codex_cli_rs/0.144.1 (Ubuntu 22.4.0; x86_64) xterm-256color",
wantOriginator: "codex_cli_rs",
wantUA: "codex_cli_rs/0.144.1 (Ubuntu 22.4.0; x86_64) xterm-256color",
wantOK: true,
},
{
name: "tui 首段直接配对",
ua: "codex-tui/0.140.2 (Mac OS X 14.0; arm64) iTerm (codex-tui; 0.140.2)",
wantOriginator: "codex-tui",
wantUA: "codex-tui/0.140.2 (Mac OS X 14.0; arm64) iTerm (codex-tui; 0.140.2)",
wantOK: true,
},
{
name: "Codex 家族前缀配对保留原大小写",
ua: "Codex Desktop/1.2.3",
wantOriginator: "Codex Desktop",
wantUA: "Codex Desktop/1.2.3",
wantOK: true,
},
{
name: "originator override 用尾部 name 重写首段",
ua: "cccc/0.142.0 (Ubuntu 22.4.0; x86_64) screen (codex-tui; 0.142.0)",
wantOriginator: "codex-tui",
wantUA: "codex-tui/0.142.0 (Ubuntu 22.4.0; x86_64) screen (codex-tui; 0.142.0)",
wantOK: true,
},
{
name: "override 尾部恢复保留 Codex 家族真实大小写",
ua: "cccc/1.2.3 (Ubuntu 22.4.0; x86_64) term (Codex Desktop; 1.2.3)",
wantOriginator: "Codex Desktop",
wantUA: "Codex Desktop/1.2.3 (Ubuntu 22.4.0; x86_64) term (Codex Desktop; 1.2.3)",
wantOK: true,
},
{name: "含斜杠的尾部 name 拒绝配对(防自不一致身份)", ua: "foo/1.0 (Codex Desktop/2; 1.0)", wantOK: false},
{
name: "精确集合大小写变体归一为规范小写",
ua: "CODEX_CLI_RS/1.0.0",
wantOriginator: "codex_cli_rs",
wantUA: "codex_cli_rs/1.0.0",
wantOK: true,
},
{
name: "首段尾随空格重建为规范 UA",
ua: "codex-tui /1.0.0",
wantOriginator: "codex-tui",
wantUA: "codex-tui/1.0.0",
wantOK: true,
},
{name: "家族前缀夹带不可打印字节拒绝", ua: "Codex \x01evil/1.0.0", wantOK: false},
{name: "家族前缀夹带非 ASCII 字节拒绝", ua: "Codex \xc3\xa9vil/1.0.0", wantOK: false},
{name: "超长首段拒绝", ua: "Codex " + strings.Repeat("a", 80) + "/1.0.0", wantOK: false},
{name: "第三方 UA 不可配对", ua: "luna/1.0.0", wantOK: false},
{name: "伪造前缀不可配对", ua: "codex_cli_rs_evil/1.0.0", wantOK: false},
{name: "浏览器 UA 不可配对", ua: "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0 Safari/537.36", wantOK: false},
{name: "无斜杠不可配对", ua: "curl", wantOK: false},
{name: "空 UA 不可配对", ua: "", wantOK: false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
originator, pairedUA, ok := PairCodexClientIdentity(tt.ua)
require.Equal(t, tt.wantOK, ok)
require.Equal(t, tt.wantOriginator, originator)
require.Equal(t, tt.wantUA, pairedUA)
})
}
}
+1 -1
View File
@@ -729,7 +729,7 @@ func (r *accountRepository) ListOAuthRefreshCandidates(ctx context.Context) ([]s
FROM accounts
WHERE deleted_at IS NULL
AND status = 'active'
AND type = 'oauth'
AND type IN ('oauth', 'setup-token')
AND platform IN ('anthropic', 'openai', 'gemini', 'antigravity')
AND credentials ? 'refresh_token'
AND btrim(credentials->>'refresh_token') <> ''
@@ -43,7 +43,8 @@ func TestAccountRepository_ListOAuthRefreshCandidates_SQLFilter(t *testing.T) {
normalized := normalizeSQLWhitespace(capturedSQL)
require.Contains(t, normalized, "deleted_at IS NULL")
require.Contains(t, normalized, "status = 'active'")
require.Contains(t, normalized, "type = 'oauth'")
// setup-token 的 access_token 同为 8h 短期令牌,必须与 oauth 一起纳入后台刷新候选
require.Contains(t, normalized, "type IN ('oauth', 'setup-token')")
require.Contains(t, normalized, "platform IN ('anthropic', 'openai', 'gemini', 'antigravity')")
require.Contains(t, normalized, "credentials ? 'refresh_token'")
require.Contains(t, normalized, "btrim(credentials->>'refresh_token') <> ''")
+12 -4
View File
@@ -199,16 +199,24 @@ func appendRequestTypeOrStreamQueryFilter(query string, args []any, requestType
// buildRequestTypeFilterCondition 在 request_type 过滤时兼容 legacy 字段,避免历史数据漏查。
func buildRequestTypeFilterCondition(startArgIndex int, requestType int16) (string, []any) {
return buildRequestTypeFilterConditionWithAlias(startArgIndex, requestType, "")
}
func buildRequestTypeFilterConditionWithAlias(startArgIndex int, requestType int16, alias string) (string, []any) {
normalized := service.RequestTypeFromInt16(requestType)
requestTypeArg := int16(normalized)
prefix := ""
if alias != "" {
prefix = alias + "."
}
switch normalized {
case service.RequestTypeSync:
return fmt.Sprintf("(request_type = $%d OR (request_type = %d AND stream = FALSE AND openai_ws_mode = FALSE))", startArgIndex, int16(service.RequestTypeUnknown)), []any{requestTypeArg}
return fmt.Sprintf("(%srequest_type = $%d OR (%srequest_type = %d AND %sstream = FALSE AND %sopenai_ws_mode = FALSE))", prefix, startArgIndex, prefix, int16(service.RequestTypeUnknown), prefix, prefix), []any{requestTypeArg}
case service.RequestTypeStream:
return fmt.Sprintf("(request_type = $%d OR (request_type = %d AND stream = TRUE AND openai_ws_mode = FALSE))", startArgIndex, int16(service.RequestTypeUnknown)), []any{requestTypeArg}
return fmt.Sprintf("(%srequest_type = $%d OR (%srequest_type = %d AND %sstream = TRUE AND %sopenai_ws_mode = FALSE))", prefix, startArgIndex, prefix, int16(service.RequestTypeUnknown), prefix, prefix), []any{requestTypeArg}
case service.RequestTypeWSV2:
return fmt.Sprintf("(request_type = $%d OR (request_type = %d AND openai_ws_mode = TRUE))", startArgIndex, int16(service.RequestTypeUnknown)), []any{requestTypeArg}
return fmt.Sprintf("(%srequest_type = $%d OR (%srequest_type = %d AND %sopenai_ws_mode = TRUE))", prefix, startArgIndex, prefix, int16(service.RequestTypeUnknown), prefix), []any{requestTypeArg}
default:
return fmt.Sprintf("request_type = $%d", startArgIndex), []any{requestTypeArg}
return fmt.Sprintf("%srequest_type = $%d", prefix, startArgIndex), []any{requestTypeArg}
}
}
@@ -3,9 +3,14 @@
package repository
import (
"context"
"regexp"
"testing"
"time"
"github.com/DATA-DOG/go-sqlmock"
"github.com/Wei-Shaw/sub2api/internal/pkg/usagestats"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/stretchr/testify/require"
)
@@ -48,3 +53,27 @@ func TestResolveModelDimensionExpression(t *testing.T) {
})
}
}
func TestGetUserBreakdownStatsRequestTypeIncludesLegacyFallback(t *testing.T) {
db, mock := newSQLMock(t)
repo := &usageLogRepository{sql: db}
start := time.Date(2026, 7, 1, 0, 0, 0, 0, time.UTC)
end := start.Add(24 * time.Hour)
requestType := int16(service.RequestTypeStream)
legacyFilter := `(ul.request_type = $3 OR (ul.request_type = 0 AND ul.stream = TRUE AND ul.openai_ws_mode = FALSE))`
mock.ExpectQuery(regexp.QuoteMeta(legacyFilter)).
WithArgs(start, end, requestType).
WillReturnRows(sqlmock.NewRows([]string{
"user_id", "email", "requests", "input_tokens", "output_tokens",
"cache_tokens", "total_tokens", "cost", "actual_cost", "account_cost",
}))
rows, err := repo.GetUserBreakdownStats(context.Background(), start, end, usagestats.UserBreakdownDimension{
RequestType: &requestType,
}, 0)
require.NoError(t, err)
require.Empty(t, rows)
require.NoError(t, mock.ExpectationsWereMet())
}
@@ -642,8 +642,9 @@ func (r *usageLogRepository) GetUserBreakdownStats(ctx context.Context, startTim
args = append(args, dim.AccountID)
}
if dim.RequestType != nil {
query += fmt.Sprintf(" AND ul.request_type = $%d", len(args)+1)
args = append(args, *dim.RequestType)
condition, conditionArgs := buildRequestTypeFilterConditionWithAlias(len(args)+1, *dim.RequestType, "ul")
query += " AND " + condition
args = append(args, conditionArgs...)
}
if dim.Stream != nil {
query += fmt.Sprintf(" AND ul.stream = $%d", len(args)+1)
@@ -53,11 +53,11 @@ func (s *UserMsgQueueCacheSuite) TestReconcileExpiredLockCandidatesRemovesNatura
require.NoError(s.T(), err)
require.True(s.T(), acquired)
score, err := s.rdb.ZScore(s.ctx, umqLockIndexKey, "702").Result()
_, err = s.rdb.ZScore(s.ctx, umqLockIndexKey, "702").Result()
require.NoError(s.T(), err)
require.Eventually(s.T(), func() bool {
nowMs, err := s.cache.GetCurrentTimeMs(s.ctx)
return err == nil && nowMs >= int64(score)
_, err := s.rdb.Get(s.ctx, umqLockKey(accountID)).Result()
return errors.Is(err, redis.Nil)
}, time.Second, 10*time.Millisecond)
cleaned, err := s.cache.ReconcileExpiredLockCandidates(s.ctx, 1000)
+44
View File
@@ -32,6 +32,8 @@ type userRepository struct {
sql sqlExecutor
}
var _ service.RedeemUserAdjustmentRepository = (*userRepository)(nil)
func NewUserRepository(client *dbent.Client, sqlDB *sql.DB) service.UserRepository {
return newUserRepositoryWithSQL(client, sqlDB)
}
@@ -751,6 +753,27 @@ func (r *userRepository) UpdateBalance(ctx context.Context, id int64, amount flo
return nil
}
func (r *userRepository) ApplyRedeemBalanceAdjustment(ctx context.Context, id int64, delta float64) error {
const updateSQL = `
UPDATE users
SET balance = GREATEST(balance + $1, 0), updated_at = NOW()
WHERE id = $2 AND deleted_at IS NULL
`
client := clientFromContext(ctx, r.client)
result, err := client.ExecContext(ctx, updateSQL, delta, id)
if err != nil {
return err
}
affected, err := result.RowsAffected()
if err != nil {
return err
}
if affected == 0 {
return service.ErrUserNotFound
}
return nil
}
// DeductBalance 扣除用户余额
// 透支策略:允许余额变为负数,确保当前请求能够完成
// 中间件会阻止余额 <= 0 的用户发起后续请求
@@ -792,6 +815,27 @@ func (r *userRepository) UpdateConcurrency(ctx context.Context, id int64, amount
return nil
}
func (r *userRepository) ApplyRedeemConcurrencyAdjustment(ctx context.Context, id int64, delta int) error {
const updateSQL = `
UPDATE users
SET concurrency = GREATEST(concurrency + $1, 0), updated_at = NOW()
WHERE id = $2 AND deleted_at IS NULL
`
client := clientFromContext(ctx, r.client)
result, err := client.ExecContext(ctx, updateSQL, delta, id)
if err != nil {
return err
}
affected, err := result.RowsAffected()
if err != nil {
return err
}
if affected == 0 {
return service.ErrUserNotFound
}
return nil
}
func (r *userRepository) BatchSetConcurrency(ctx context.Context, userIDs []int64, value int) (int, error) {
if len(userIDs) == 0 {
return 0, nil
@@ -4,6 +4,7 @@ package repository
import (
"context"
"sync"
"testing"
"time"
@@ -353,6 +354,29 @@ func (s *UserRepoSuite) TestUpdateBalance_Negative() {
s.Require().InDelta(7.0, got.Balance, 1e-6)
}
func (s *UserRepoSuite) TestApplyRedeemBalanceAdjustment_ConcurrentNeverNegative() {
user := s.mustCreateUser(&service.User{Email: "redeem-bal-concurrent@test.com", Balance: 10})
var wg sync.WaitGroup
errs := make(chan error, 2)
for i := 0; i < 2; i++ {
wg.Add(1)
go func() {
defer wg.Done()
errs <- s.repo.ApplyRedeemBalanceAdjustment(context.Background(), user.ID, -7)
}()
}
wg.Wait()
close(errs)
for err := range errs {
s.Require().NoError(err)
}
got, err := s.repo.GetByID(s.ctx, user.ID)
s.Require().NoError(err)
s.Require().InDelta(0, got.Balance, 1e-6)
}
func (s *UserRepoSuite) TestDeductBalance() {
user := s.mustCreateUser(&service.User{Email: "deduct@test.com", Balance: 10})
@@ -425,6 +449,29 @@ func (s *UserRepoSuite) TestUpdateConcurrency_Negative() {
s.Require().Equal(3, got.Concurrency)
}
func (s *UserRepoSuite) TestApplyRedeemConcurrencyAdjustment_ConcurrentNeverNegative() {
user := s.mustCreateUser(&service.User{Email: "redeem-concurrency-concurrent@test.com", Concurrency: 10})
var wg sync.WaitGroup
errs := make(chan error, 2)
for i := 0; i < 2; i++ {
wg.Add(1)
go func() {
defer wg.Done()
errs <- s.repo.ApplyRedeemConcurrencyAdjustment(context.Background(), user.ID, -7)
}()
}
wg.Wait()
close(errs)
for err := range errs {
s.Require().NoError(err)
}
got, err := s.repo.GetByID(s.ctx, user.ID)
s.Require().NoError(err)
s.Require().Equal(0, got.Concurrency)
}
// --- ExistsByEmail ---
func (s *UserRepoSuite) TestExistsByEmail() {
@@ -0,0 +1,56 @@
package repository
import (
"context"
"testing"
"github.com/DATA-DOG/go-sqlmock"
dbent "github.com/Wei-Shaw/sub2api/ent"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/stretchr/testify/require"
"entgo.io/ent/dialect"
entsql "entgo.io/ent/dialect/sql"
)
func newRedeemAdjustmentRepoMock(t *testing.T) (*userRepository, sqlmock.Sqlmock) {
t.Helper()
db, mock, err := sqlmock.New()
require.NoError(t, err)
t.Cleanup(func() { _ = db.Close() })
driver := entsql.OpenDB(dialect.Postgres, db)
client := dbent.NewClient(dbent.Driver(driver))
t.Cleanup(func() { _ = client.Close() })
return newUserRepositoryWithSQL(client, db), mock
}
func TestApplyRedeemBalanceAdjustment_UsesAtomicFloor(t *testing.T) {
repo, mock := newRedeemAdjustmentRepoMock(t)
mock.ExpectExec(`UPDATE users SET balance = GREATEST\(balance \+ \$1, 0\), updated_at = NOW\(\) WHERE id = \$2 AND deleted_at IS NULL`).
WithArgs(-7.0, int64(42)).
WillReturnResult(sqlmock.NewResult(0, 1))
require.NoError(t, repo.ApplyRedeemBalanceAdjustment(context.Background(), 42, -7))
require.NoError(t, mock.ExpectationsWereMet())
}
func TestApplyRedeemConcurrencyAdjustment_UsesAtomicFloor(t *testing.T) {
repo, mock := newRedeemAdjustmentRepoMock(t)
mock.ExpectExec(`UPDATE users SET concurrency = GREATEST\(concurrency \+ \$1, 0\), updated_at = NOW\(\) WHERE id = \$2 AND deleted_at IS NULL`).
WithArgs(-7, int64(42)).
WillReturnResult(sqlmock.NewResult(0, 1))
require.NoError(t, repo.ApplyRedeemConcurrencyAdjustment(context.Background(), 42, -7))
require.NoError(t, mock.ExpectationsWereMet())
}
func TestApplyRedeemAdjustment_MissingUser(t *testing.T) {
repo, mock := newRedeemAdjustmentRepoMock(t)
mock.ExpectExec(`UPDATE users SET balance = GREATEST\(balance \+ \$1, 0\), updated_at = NOW\(\) WHERE id = \$2 AND deleted_at IS NULL`).
WithArgs(-1.0, int64(404)).
WillReturnResult(sqlmock.NewResult(0, 0))
err := repo.ApplyRedeemBalanceAdjustment(context.Background(), 404, -1)
require.ErrorIs(t, err, service.ErrUserNotFound)
require.NoError(t, mock.ExpectationsWereMet())
}
@@ -366,31 +366,85 @@ func (r *userSubscriptionRepository) ActivateWindows(ctx context.Context, id int
return translatePersistenceError(err, service.ErrSubscriptionNotFound, nil)
}
func (r *userSubscriptionRepository) ResetDailyUsage(ctx context.Context, id int64, newWindowStart time.Time) error {
func (r *userSubscriptionRepository) ResetUsageWindows(ctx context.Context, id int64, resetDaily, resetWeekly, resetMonthly bool, newWindowStart time.Time) error {
client := clientFromContext(ctx, r.client)
_, err := client.UserSubscription.UpdateOneID(id).
update := client.UserSubscription.UpdateOneID(id)
if resetDaily {
update.SetDailyUsageUsd(0).SetDailyWindowStart(newWindowStart)
}
if resetWeekly {
update.SetWeeklyUsageUsd(0).SetWeeklyWindowStart(newWindowStart)
}
if resetMonthly {
update.SetMonthlyUsageUsd(0).SetMonthlyWindowStart(newWindowStart)
}
_, err := update.Save(ctx)
return translatePersistenceError(err, service.ErrSubscriptionNotFound, nil)
}
func (r *userSubscriptionRepository) ResetDailyUsage(ctx context.Context, id int64, expectedWindowStart *time.Time, newWindowStart time.Time) error {
client := clientFromContext(ctx, r.client)
query := client.UserSubscription.Update().Where(usersubscription.IDEQ(id))
if expectedWindowStart == nil {
query = query.Where(usersubscription.DailyWindowStartIsNil())
} else {
query = query.Where(usersubscription.DailyWindowStartEQ(*expectedWindowStart))
}
n, err := query.
SetDailyUsageUsd(0).
SetDailyWindowStart(newWindowStart).
Save(ctx)
return translatePersistenceError(err, service.ErrSubscriptionNotFound, nil)
return r.translateConditionalWindowReset(ctx, client, id, n, err)
}
func (r *userSubscriptionRepository) ResetWeeklyUsage(ctx context.Context, id int64, newWindowStart time.Time) error {
func (r *userSubscriptionRepository) ResetWeeklyUsage(ctx context.Context, id int64, expectedWindowStart *time.Time, newWindowStart time.Time) error {
client := clientFromContext(ctx, r.client)
_, err := client.UserSubscription.UpdateOneID(id).
query := client.UserSubscription.Update().Where(usersubscription.IDEQ(id))
if expectedWindowStart == nil {
query = query.Where(usersubscription.WeeklyWindowStartIsNil())
} else {
query = query.Where(usersubscription.WeeklyWindowStartEQ(*expectedWindowStart))
}
n, err := query.
SetWeeklyUsageUsd(0).
SetWeeklyWindowStart(newWindowStart).
Save(ctx)
return translatePersistenceError(err, service.ErrSubscriptionNotFound, nil)
return r.translateConditionalWindowReset(ctx, client, id, n, err)
}
func (r *userSubscriptionRepository) ResetMonthlyUsage(ctx context.Context, id int64, newWindowStart time.Time) error {
func (r *userSubscriptionRepository) ResetMonthlyUsage(ctx context.Context, id int64, expectedWindowStart *time.Time, newWindowStart time.Time) error {
client := clientFromContext(ctx, r.client)
_, err := client.UserSubscription.UpdateOneID(id).
query := client.UserSubscription.Update().Where(usersubscription.IDEQ(id))
if expectedWindowStart == nil {
query = query.Where(usersubscription.MonthlyWindowStartIsNil())
} else {
query = query.Where(usersubscription.MonthlyWindowStartEQ(*expectedWindowStart))
}
n, err := query.
SetMonthlyUsageUsd(0).
SetMonthlyWindowStart(newWindowStart).
Save(ctx)
return translatePersistenceError(err, service.ErrSubscriptionNotFound, nil)
return r.translateConditionalWindowReset(ctx, client, id, n, err)
}
func (r *userSubscriptionRepository) translateConditionalWindowReset(ctx context.Context, client *dbent.Client, id int64, affected int, err error) error {
if err != nil {
return translatePersistenceError(err, service.ErrSubscriptionNotFound, nil)
}
if affected > 0 {
return nil
}
// A stale reset is an expected no-op: another request already advanced the
// window. Preserve not-found semantics for callers that target a missing row.
exists, err := client.UserSubscription.Query().Where(usersubscription.IDEQ(id)).Exist(ctx)
if err != nil {
return translatePersistenceError(err, service.ErrSubscriptionNotFound, nil)
}
if !exists {
return service.ErrSubscriptionNotFound
}
return nil
}
// IncrementUsage 原子性地累加订阅用量。
@@ -472,7 +472,7 @@ func (s *UserSubscriptionRepoSuite) TestResetDailyUsage() {
})
resetAt := time.Date(2025, 1, 2, 0, 0, 0, 0, time.UTC)
err := s.repo.ResetDailyUsage(s.ctx, sub.ID, resetAt)
err := s.repo.ResetDailyUsage(s.ctx, sub.ID, sub.DailyWindowStart, resetAt)
s.Require().NoError(err, "ResetDailyUsage")
got, err := s.repo.GetByID(s.ctx, sub.ID)
@@ -483,6 +483,47 @@ func (s *UserSubscriptionRepoSuite) TestResetDailyUsage() {
s.Require().WithinDuration(resetAt, *got.DailyWindowStart, time.Microsecond)
}
func (s *UserSubscriptionRepoSuite) TestResetDailyUsage_StaleResetDoesNotClearNewWindowUsage() {
user := s.mustCreateUser("resetd-cas@test.com", service.RoleUser)
group := s.mustCreateGroup("g-resetd-cas")
oldWindowStart := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC)
sub := s.mustCreateSubscription(user.ID, group.ID, func(c *dbent.UserSubscriptionCreate) {
c.SetDailyWindowStart(oldWindowStart)
c.SetDailyUsageUsd(10)
})
newWindowStart := oldWindowStart.Add(24 * time.Hour)
s.Require().NoError(s.repo.ResetDailyUsage(s.ctx, sub.ID, &oldWindowStart, newWindowStart))
s.Require().NoError(s.repo.IncrementUsage(s.ctx, sub.ID, 3))
// Simulate a second request carrying the stale old-window snapshot.
s.Require().NoError(s.repo.ResetDailyUsage(s.ctx, sub.ID, &oldWindowStart, newWindowStart))
got, err := s.repo.GetByID(s.ctx, sub.ID)
s.Require().NoError(err)
s.Require().InDelta(3, got.DailyUsageUSD, 1e-6)
s.Require().WithinDuration(newWindowStart, *got.DailyWindowStart, time.Microsecond)
}
func (s *UserSubscriptionRepoSuite) TestResetUsageWindows_ClearsUsageAfterAutomaticWindowAdvance() {
user := s.mustCreateUser("admin-reset-current@test.com", service.RoleUser)
group := s.mustCreateGroup("g-admin-reset-current")
oldWindowStart := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC)
sub := s.mustCreateSubscription(user.ID, group.ID, func(c *dbent.UserSubscriptionCreate) {
c.SetDailyWindowStart(oldWindowStart)
c.SetDailyUsageUsd(10)
})
newWindowStart := oldWindowStart.Add(24 * time.Hour)
s.Require().NoError(s.repo.ResetDailyUsage(s.ctx, sub.ID, &oldWindowStart, newWindowStart))
s.Require().NoError(s.repo.IncrementUsage(s.ctx, sub.ID, 3))
s.Require().NoError(s.repo.ResetUsageWindows(s.ctx, sub.ID, true, false, false, newWindowStart))
got, err := s.repo.GetByID(s.ctx, sub.ID)
s.Require().NoError(err)
s.Require().InDelta(0, got.DailyUsageUSD, 1e-6)
s.Require().WithinDuration(newWindowStart, *got.DailyWindowStart, time.Microsecond)
}
func (s *UserSubscriptionRepoSuite) TestResetWeeklyUsage() {
user := s.mustCreateUser("resetw@test.com", service.RoleUser)
group := s.mustCreateGroup("g-resetw")
@@ -492,7 +533,7 @@ func (s *UserSubscriptionRepoSuite) TestResetWeeklyUsage() {
})
resetAt := time.Date(2025, 1, 6, 0, 0, 0, 0, time.UTC)
err := s.repo.ResetWeeklyUsage(s.ctx, sub.ID, resetAt)
err := s.repo.ResetWeeklyUsage(s.ctx, sub.ID, sub.WeeklyWindowStart, resetAt)
s.Require().NoError(err, "ResetWeeklyUsage")
got, err := s.repo.GetByID(s.ctx, sub.ID)
@@ -511,7 +552,7 @@ func (s *UserSubscriptionRepoSuite) TestResetMonthlyUsage() {
})
resetAt := time.Date(2025, 2, 1, 0, 0, 0, 0, time.UTC)
err := s.repo.ResetMonthlyUsage(s.ctx, sub.ID, resetAt)
err := s.repo.ResetMonthlyUsage(s.ctx, sub.ID, sub.MonthlyWindowStart, resetAt)
s.Require().NoError(err, "ResetMonthlyUsage")
got, err := s.repo.GetByID(s.ctx, sub.ID)
@@ -723,7 +764,7 @@ func (s *UserSubscriptionRepoSuite) TestActiveExpiredBoundaries_UsageAndReset_Ba
s.Require().NotNil(after.MonthlyWindowStart, "expected MonthlyWindowStart activated")
resetAt := time.Now().Truncate(time.Microsecond) // truncate to microsecond for DB precision
s.Require().NoError(s.repo.ResetDailyUsage(s.ctx, active.ID, resetAt), "ResetDailyUsage")
s.Require().NoError(s.repo.ResetDailyUsage(s.ctx, active.ID, after.DailyWindowStart, resetAt), "ResetDailyUsage")
afterReset, err := s.repo.GetByID(s.ctx, active.ID)
s.Require().NoError(err, "GetByID after reset")
s.Require().InDelta(0.0, afterReset.DailyUsageUSD, 1e-6)
+6 -3
View File
@@ -2123,13 +2123,16 @@ func (stubUserSubscriptionRepo) UpdateNotes(ctx context.Context, subscriptionID
func (stubUserSubscriptionRepo) ActivateWindows(ctx context.Context, id int64, start time.Time) error {
return errors.New("not implemented")
}
func (stubUserSubscriptionRepo) ResetDailyUsage(ctx context.Context, id int64, newWindowStart time.Time) error {
func (stubUserSubscriptionRepo) ResetUsageWindows(ctx context.Context, id int64, resetDaily, resetWeekly, resetMonthly bool, newWindowStart time.Time) error {
return errors.New("not implemented")
}
func (stubUserSubscriptionRepo) ResetWeeklyUsage(ctx context.Context, id int64, newWindowStart time.Time) error {
func (stubUserSubscriptionRepo) ResetDailyUsage(ctx context.Context, id int64, expectedWindowStart *time.Time, newWindowStart time.Time) error {
return errors.New("not implemented")
}
func (stubUserSubscriptionRepo) ResetMonthlyUsage(ctx context.Context, id int64, newWindowStart time.Time) error {
func (stubUserSubscriptionRepo) ResetWeeklyUsage(ctx context.Context, id int64, expectedWindowStart *time.Time, newWindowStart time.Time) error {
return errors.New("not implemented")
}
func (stubUserSubscriptionRepo) ResetMonthlyUsage(ctx context.Context, id int64, expectedWindowStart *time.Time, newWindowStart time.Time) error {
return errors.New("not implemented")
}
func (stubUserSubscriptionRepo) IncrementUsage(ctx context.Context, id int64, costUSD float64) error {
@@ -126,6 +126,8 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti
if abortIfAPIKeyGroupNotAllowed(c, apiKey) {
return
}
ctx := context.WithValue(c.Request.Context(), ctxkey.UserID, apiKey.User.ID)
c.Request = c.Request.WithContext(ctx)
// ── 4. SimpleMode → early return ─────────────────────────────
@@ -193,6 +195,15 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti
// 订阅模式:验证订阅限额
if subscription != nil {
needsMaintenance, validateErr := subscriptionService.ValidateAndCheckLimits(subscription, apiKey.Group)
if needsMaintenance {
refreshed, maintenanceErr := subscriptionService.EnsureWindowMaintenance(c.Request.Context(), subscription)
if maintenanceErr != nil {
AbortWithError(c, 500, "SUBSCRIPTION_MAINTENANCE_FAILED", "Failed to maintain subscription usage windows")
return
}
subscription = refreshed
_, validateErr = subscriptionService.ValidateAndCheckLimits(subscription, apiKey.Group)
}
if validateErr != nil {
code := "SUBSCRIPTION_INVALID"
status := 403
@@ -205,12 +216,6 @@ func apiKeyAuthWithSubscription(apiKeyService *service.APIKeyService, subscripti
AbortWithError(c, status, code, validateErr.Error())
return
}
// 窗口维护异步化(不阻塞请求)
if needsMaintenance {
maintenanceCopy := *subscription
subscriptionService.DoWindowMaintenance(&maintenanceCopy)
}
} else {
// 非订阅模式 或 订阅模式但 subscriptionService 未注入:回退到余额检查
if apiKeyBalanceBelowAuthThreshold(apiKey.User.Balance, cfg) {
@@ -141,6 +141,15 @@ func APIKeyAuthWithSubscriptionGoogle(apiKeyService *service.APIKeyService, subs
}
needsMaintenance, err := subscriptionService.ValidateAndCheckLimits(subscription, apiKey.Group)
if needsMaintenance {
refreshed, maintenanceErr := subscriptionService.EnsureWindowMaintenance(c.Request.Context(), subscription)
if maintenanceErr != nil {
abortWithGoogleError(c, 500, "Failed to maintain subscription usage windows")
return
}
subscription = refreshed
_, err = subscriptionService.ValidateAndCheckLimits(subscription, apiKey.Group)
}
if err != nil {
status := 403
if errors.Is(err, service.ErrDailyLimitExceeded) ||
@@ -153,11 +162,6 @@ func APIKeyAuthWithSubscriptionGoogle(apiKeyService *service.APIKeyService, subs
}
c.Set(string(ContextKeySubscription), subscription)
if needsMaintenance {
maintenanceCopy := *subscription
subscriptionService.DoWindowMaintenance(&maintenanceCopy)
}
} else {
if apiKeyBalanceBelowAuthThreshold(apiKey.User.Balance, cfg) {
abortWithGoogleError(c, 403, "Insufficient account balance")
@@ -24,6 +24,7 @@ type fakeAPIKeyRepo struct {
}
type fakeGoogleSubscriptionRepo struct {
getByID func(ctx context.Context, id int64) (*service.UserSubscription, error)
getActive func(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error)
updateStatus func(ctx context.Context, subscriptionID int64, status string) error
activateWindow func(ctx context.Context, id int64, start time.Time) error
@@ -115,6 +116,9 @@ func (f fakeGoogleSubscriptionRepo) Create(ctx context.Context, sub *service.Use
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) GetByID(ctx context.Context, id int64) (*service.UserSubscription, error) {
if f.getByID != nil {
return f.getByID(ctx, id)
}
return nil, errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) GetByIDIncludeDeleted(ctx context.Context, id int64) (*service.UserSubscription, error) {
@@ -174,19 +178,22 @@ func (f fakeGoogleSubscriptionRepo) ActivateWindows(ctx context.Context, id int6
}
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ResetDailyUsage(ctx context.Context, id int64, start time.Time) error {
func (f fakeGoogleSubscriptionRepo) ResetUsageWindows(context.Context, int64, bool, bool, bool, time.Time) error {
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ResetDailyUsage(ctx context.Context, id int64, _ *time.Time, start time.Time) error {
if f.resetDaily != nil {
return f.resetDaily(ctx, id, start)
}
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ResetWeeklyUsage(ctx context.Context, id int64, start time.Time) error {
func (f fakeGoogleSubscriptionRepo) ResetWeeklyUsage(ctx context.Context, id int64, _ *time.Time, start time.Time) error {
if f.resetWeekly != nil {
return f.resetWeekly(ctx, id, start)
}
return errors.New("not implemented")
}
func (f fakeGoogleSubscriptionRepo) ResetMonthlyUsage(ctx context.Context, id int64, start time.Time) error {
func (f fakeGoogleSubscriptionRepo) ResetMonthlyUsage(ctx context.Context, id int64, _ *time.Time, start time.Time) error {
if f.resetMonthly != nil {
return f.resetMonthly(ctx, id, start)
}
@@ -58,7 +58,7 @@ func TestSimpleModeBypassesQuotaCheck(t *testing.T) {
},
}
t.Run("standard_mode_needs_maintenance_does_not_block_request", func(t *testing.T) {
t.Run("standard_mode_completes_maintenance_before_request", func(t *testing.T) {
cfg := &config.Config{RunMode: config.RunModeStandard}
cfg.SubscriptionMaintenance.WorkerCount = 1
cfg.SubscriptionMaintenance.QueueSize = 1
@@ -67,16 +67,22 @@ func TestSimpleModeBypassesQuotaCheck(t *testing.T) {
past := time.Now().Add(-48 * time.Hour)
sub := &service.UserSubscription{
ID: 55,
UserID: user.ID,
GroupID: group.ID,
Status: service.SubscriptionStatusActive,
ExpiresAt: time.Now().Add(24 * time.Hour),
DailyWindowStart: &past,
DailyUsageUSD: 0,
ID: 55,
UserID: user.ID,
GroupID: group.ID,
Status: service.SubscriptionStatusActive,
ExpiresAt: time.Now().Add(24 * time.Hour),
DailyWindowStart: &past,
WeeklyWindowStart: &past,
MonthlyWindowStart: &past,
DailyUsageUSD: 0,
}
maintenanceCalled := make(chan struct{}, 1)
subscriptionRepo := &stubUserSubscriptionRepo{
getByID: func(ctx context.Context, id int64) (*service.UserSubscription, error) {
clone := *sub
return &clone, nil
},
getActive: func(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error) {
clone := *sub
return &clone, nil
@@ -84,11 +90,19 @@ func TestSimpleModeBypassesQuotaCheck(t *testing.T) {
updateStatus: func(ctx context.Context, subscriptionID int64, status string) error { return nil },
activateWindow: func(ctx context.Context, id int64, start time.Time) error { return nil },
resetDaily: func(ctx context.Context, id int64, start time.Time) error {
sub.DailyWindowStart = &start
sub.DailyUsageUSD = 0
maintenanceCalled <- struct{}{}
return nil
},
resetWeekly: func(ctx context.Context, id int64, start time.Time) error { return nil },
resetMonthly: func(ctx context.Context, id int64, start time.Time) error { return nil },
resetWeekly: func(ctx context.Context, id int64, start time.Time) error {
sub.WeeklyWindowStart = &start
return nil
},
resetMonthly: func(ctx context.Context, id int64, start time.Time) error {
sub.MonthlyWindowStart = &start
return nil
},
}
subscriptionService := service.NewSubscriptionService(nil, subscriptionRepo, nil, nil, cfg)
t.Cleanup(subscriptionService.Stop)
@@ -105,10 +119,57 @@ func TestSimpleModeBypassesQuotaCheck(t *testing.T) {
case <-maintenanceCalled:
// ok
case <-time.After(time.Second):
t.Fatalf("expected maintenance to be scheduled")
t.Fatalf("expected maintenance to complete before response")
}
})
t.Run("standard_mode_revalidates_cas_loser_from_database", func(t *testing.T) {
cfg := &config.Config{RunMode: config.RunModeStandard}
apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
past := time.Now().Add(-48 * time.Hour)
current := time.Now()
stale := &service.UserSubscription{
ID: 56,
UserID: user.ID,
GroupID: group.ID,
Status: service.SubscriptionStatusActive,
ExpiresAt: current.Add(24 * time.Hour),
DailyWindowStart: &past,
WeeklyWindowStart: &past,
MonthlyWindowStart: &past,
DailyUsageUSD: 10,
}
fresh := *stale
fresh.DailyWindowStart = &current
fresh.WeeklyWindowStart = &current
fresh.MonthlyWindowStart = &current
fresh.DailyUsageUSD = 2
subscriptionRepo := &stubUserSubscriptionRepo{
getActive: func(context.Context, int64, int64) (*service.UserSubscription, error) {
clone := *stale
return &clone, nil
},
getByID: func(context.Context, int64) (*service.UserSubscription, error) {
clone := fresh
return &clone, nil
},
resetDaily: func(context.Context, int64, time.Time) error { return nil },
resetWeekly: func(context.Context, int64, time.Time) error { return nil },
resetMonthly: func(context.Context, int64, time.Time) error { return nil },
}
subscriptionService := service.NewSubscriptionService(nil, subscriptionRepo, nil, nil, cfg)
router := newAuthTestRouter(apiKeyService, subscriptionService, cfg)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/t", nil)
req.Header.Set("x-api-key", apiKey.Key)
router.ServeHTTP(w, req)
require.Equal(t, http.StatusTooManyRequests, w.Code)
})
t.Run("simple_mode_bypasses_quota_check", func(t *testing.T) {
cfg := &config.Config{RunMode: config.RunModeSimple}
apiKeyService := service.NewAPIKeyService(apiKeyRepo, nil, nil, nil, nil, nil, cfg)
@@ -225,6 +286,11 @@ func TestAPIKeyAuthSetsGroupContext(t *testing.T) {
c.JSON(http.StatusInternalServerError, gin.H{"ok": false})
return
}
userIDFromCtx, ok := c.Request.Context().Value(ctxkey.UserID).(int64)
if !ok || userIDFromCtx != user.ID {
c.JSON(http.StatusInternalServerError, gin.H{"ok": false})
return
}
c.JSON(http.StatusOK, gin.H{"ok": true})
})
@@ -1210,6 +1276,7 @@ func (r *stubApiKeyRepo) GetRateLimitData(ctx context.Context, id int64) (*servi
}
type stubUserSubscriptionRepo struct {
getByID func(ctx context.Context, id int64) (*service.UserSubscription, error)
getActive func(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error)
updateStatus func(ctx context.Context, subscriptionID int64, status string) error
activateWindow func(ctx context.Context, id int64, start time.Time) error
@@ -1258,6 +1325,9 @@ func (r *stubUserSubscriptionRepo) Create(ctx context.Context, sub *service.User
}
func (r *stubUserSubscriptionRepo) GetByID(ctx context.Context, id int64) (*service.UserSubscription, error) {
if r.getByID != nil {
return r.getByID(ctx, id)
}
return nil, errors.New("not implemented")
}
@@ -1334,21 +1404,25 @@ func (r *stubUserSubscriptionRepo) ActivateWindows(ctx context.Context, id int64
return errors.New("not implemented")
}
func (r *stubUserSubscriptionRepo) ResetDailyUsage(ctx context.Context, id int64, newWindowStart time.Time) error {
func (r *stubUserSubscriptionRepo) ResetUsageWindows(context.Context, int64, bool, bool, bool, time.Time) error {
return errors.New("not implemented")
}
func (r *stubUserSubscriptionRepo) ResetDailyUsage(ctx context.Context, id int64, _ *time.Time, newWindowStart time.Time) error {
if r.resetDaily != nil {
return r.resetDaily(ctx, id, newWindowStart)
}
return errors.New("not implemented")
}
func (r *stubUserSubscriptionRepo) ResetWeeklyUsage(ctx context.Context, id int64, newWindowStart time.Time) error {
func (r *stubUserSubscriptionRepo) ResetWeeklyUsage(ctx context.Context, id int64, _ *time.Time, newWindowStart time.Time) error {
if r.resetWeekly != nil {
return r.resetWeekly(ctx, id, newWindowStart)
}
return errors.New("not implemented")
}
func (r *stubUserSubscriptionRepo) ResetMonthlyUsage(ctx context.Context, id int64, newWindowStart time.Time) error {
func (r *stubUserSubscriptionRepo) ResetMonthlyUsage(ctx context.Context, id int64, _ *time.Time, newWindowStart time.Time) error {
if r.resetMonthly != nil {
return r.resetMonthly(ctx, id, newWindowStart)
}
@@ -0,0 +1,189 @@
package middleware
import (
"bytes"
"context"
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
func TestAPIKeyAuthForwardsUserScopedOpenAIFastPolicyToUpstream(t *testing.T) {
gin.SetMode(gin.TestMode)
upstreamBodies := make(chan []byte, 2)
upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
body, err := io.ReadAll(r.Body)
if err != nil {
http.Error(w, "read request body", http.StatusInternalServerError)
return
}
upstreamBodies <- body
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"id":"resp_test","object":"response","model":"gpt-5","status":"completed","usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}`))
}))
defer upstreamServer.Close()
settings := &service.OpenAIFastPolicySettings{
Rules: []service.OpenAIFastPolicyRule{
{
ServiceTier: service.OpenAIFastTierPriority,
Action: service.BetaPolicyActionFilter,
Scope: service.BetaPolicyScopeAll,
},
{
ServiceTier: service.OpenAIFastTierPriority,
Action: service.BetaPolicyActionPass,
Scope: service.BetaPolicyScopeAll,
UserIDs: []int64{42},
},
},
}
settingsJSON, err := json.Marshal(settings)
require.NoError(t, err)
cfg := &config.Config{RunMode: config.RunModeSimple}
cfg.Security.URLAllowlist.Enabled = false
cfg.Security.URLAllowlist.AllowInsecureHTTP = true
settingService := service.NewSettingService(&openAIFastPolicyForwardingSettingRepo{
value: string(settingsJSON),
}, cfg)
gatewayService := service.NewOpenAIGatewayService(
nil, nil, nil, nil, nil, nil, nil, cfg,
nil, nil, nil, nil, nil, &openAIFastPolicyForwardingHTTPUpstream{client: upstreamServer.Client()},
nil, nil, nil, nil, nil, nil, settingService, nil,
)
groupID := int64(101)
group := &service.Group{
ID: groupID,
Name: "openai",
Status: service.StatusActive,
Platform: service.PlatformOpenAI,
Hydrated: true,
}
apiKeys := map[string]*service.APIKey{
"key-user-42": newOpenAIFastPolicyForwardingAPIKey(1, "key-user-42", 42, groupID, group),
"key-user-43": newOpenAIFastPolicyForwardingAPIKey(2, "key-user-43", 43, groupID, group),
}
apiKeyService := service.NewAPIKeyService(&openAIFastPolicyForwardingAPIKeyRepo{apiKeys: apiKeys}, nil, nil, nil, nil, nil, cfg)
account := &service.Account{
ID: 900,
Name: "openai-upstream",
Platform: service.PlatformOpenAI,
Type: service.AccountTypeAPIKey,
Status: service.StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": upstreamServer.URL,
},
Extra: map[string]any{"use_responses_api": true},
}
router := gin.New()
router.Use(gin.HandlerFunc(NewAPIKeyAuthMiddleware(apiKeyService, nil, cfg)))
router.POST("/v1/responses", func(c *gin.Context) {
body, readErr := io.ReadAll(c.Request.Body)
if readErr != nil {
c.Status(http.StatusBadRequest)
return
}
service.SetOpenAIClientTransport(c, service.OpenAIClientTransportHTTP)
if _, forwardErr := gatewayService.Forward(c.Request.Context(), c, account, body); forwardErr != nil {
c.Status(http.StatusBadGateway)
return
}
c.Status(http.StatusOK)
})
send := func(apiKey string) {
request := httptest.NewRequest(
http.MethodPost,
"/v1/responses",
bytes.NewBufferString(`{"model":"gpt-5","stream":false,"service_tier":"priority","input":"hi"}`),
)
request.Header.Set("Content-Type", "application/json")
request.Header.Set("x-api-key", apiKey)
response := httptest.NewRecorder()
router.ServeHTTP(response, request)
require.Equal(t, http.StatusOK, response.Code)
}
send("key-user-42")
send("key-user-43")
allowedUserBody := <-upstreamBodies
otherUserBody := <-upstreamBodies
require.Equal(t, service.OpenAIFastTierPriority, gjson.GetBytes(allowedUserBody, "service_tier").String())
require.False(t, gjson.GetBytes(otherUserBody, "service_tier").Exists())
}
func newOpenAIFastPolicyForwardingAPIKey(id int64, key string, userID, groupID int64, group *service.Group) *service.APIKey {
return &service.APIKey{
ID: id,
UserID: userID,
Key: key,
Status: service.StatusActive,
GroupID: &groupID,
User: &service.User{
ID: userID,
Role: service.RoleUser,
Status: service.StatusActive,
Balance: 10,
Concurrency: 1,
},
Group: group,
}
}
type openAIFastPolicyForwardingAPIKeyRepo struct {
service.APIKeyRepository
apiKeys map[string]*service.APIKey
}
func (r *openAIFastPolicyForwardingAPIKeyRepo) GetByKeyForAuth(_ context.Context, key string) (*service.APIKey, error) {
apiKey, ok := r.apiKeys[key]
if !ok {
return nil, service.ErrAPIKeyNotFound
}
clone := *apiKey
return &clone, nil
}
func (r *openAIFastPolicyForwardingAPIKeyRepo) UpdateLastUsed(context.Context, int64, time.Time) error {
return nil
}
type openAIFastPolicyForwardingSettingRepo struct {
service.SettingRepository
value string
}
func (r *openAIFastPolicyForwardingSettingRepo) GetValue(context.Context, string) (string, error) {
return r.value, nil
}
type openAIFastPolicyForwardingHTTPUpstream struct {
client *http.Client
}
func (u *openAIFastPolicyForwardingHTTPUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
return u.client.Do(req)
}
func (u *openAIFastPolicyForwardingHTTPUpstream) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, _ *tlsfingerprint.Profile) (*http.Response, error) {
return u.Do(req, proxyURL, accountID, accountConcurrency)
}
@@ -611,6 +611,8 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account
req.Header.Set("User-Agent", codexCLIUserAgent)
}
setOpenAIChatGPTAccountHeaders(req.Header, credentialAccount)
// 与真实转发一致:originator 与最终 User-Agent 首段配套,否则上游 404(issue #3901)。
enforceCodexIdentityHeaders(req.Header)
}
// 账号级请求头覆写:测试请求与真实转发保持一致的最终头
@@ -1712,13 +1714,15 @@ func (s *AccountTestService) testOpenAIImageOAuth(c *gin.Context, ctx context.Co
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "text/event-stream")
req.Header.Set("OpenAI-Beta", "responses=experimental")
req.Header.Set("originator", "opencode")
req.Header.Set("originator", "codex_cli_rs")
if customUA := strings.TrimSpace(account.GetOpenAIUserAgent()); customUA != "" {
req.Header.Set("User-Agent", customUA)
} else {
req.Header.Set("User-Agent", codexCLIUserAgent)
}
setOpenAIChatGPTAccountHeaders(req.Header, account)
// 与真实转发一致:originator 与最终 User-Agent 首段配套(原 opencode 与 Codex UA 错配会 404,issue #3901)。
enforceCodexIdentityHeaders(req.Header)
proxyURL := ""
if account.ProxyID != nil && account.Proxy != nil {
@@ -111,7 +111,7 @@ const (
apiQueryMaxJitter = 800 * time.Millisecond // 用量查询最大随机延迟
windowStatsCacheTTL = 1 * time.Minute
openAIProbeCacheTTL = 10 * time.Minute
openAICodexProbeVersion = "0.125.0"
openAICodexProbeVersion = "0.144.1"
)
// UsageCache 封装账户使用量相关的缓存
@@ -712,6 +712,9 @@ func (s *AccountUsageService) probeOpenAICodexSnapshot(ctx context.Context, acco
req.Header.Set("User-Agent", strings.TrimSpace(fp.UserAgent))
}
}
// 与真实转发一致:originator 与最终 User-Agent(可能来自指纹缓存,如 codex-tui)首段配套,
// 否则探针被上游 404(issue #3901)。
enforceCodexIdentityHeaders(req.Header)
setOpenAIChatGPTAccountHeaders(req.Header, account)
proxyURL := ""
+115 -57
View File
@@ -90,22 +90,24 @@ type BillingCache interface {
// ModelPricing 模型价格配置(per-token价格,与LiteLLM格式一致)
type ModelPricing struct {
InputPricePerToken float64 // 每token输入价格 (USD)
InputPricePerTokenPriority float64 // priority service tier 下每token输入价格 (USD)
ImageInputPricePerToken float64 // 图片输入 token 价格 (USD),用于多模态 embedding 等图文不同价场景;为 0 时回退到 InputPricePerToken
OutputPricePerToken float64 // 每token输出价格 (USD)
OutputPricePerTokenPriority float64 // priority service tier 下每token输出价格 (USD)
CacheCreationPricePerToken float64 // 缓存创建每token价格 (USD)
CacheReadPricePerToken float64 // 缓存读取每token价格 (USD)
CacheReadPricePerTokenPriority float64 // priority service tier 下缓存读取每token价格 (USD)
CacheCreation5mPrice float64 // 5分钟缓存创建每token价格 (USD)
CacheCreation1hPrice float64 // 1小时缓存创建每token价格 (USD)
SupportsCacheBreakdown bool // 是否支持详细的缓存分类
LongContextInputThreshold int // 超过阈值后按整次会话提升输入价格
LongContextInputMultiplier float64 // 长上下文整次会话输入倍率
LongContextOutputMultiplier float64 // 长上下文整次会话输出倍率
ImageOutputPricePerToken float64 // 图片输出 token 价格 (USD)
ImageOutputPriceExplicit bool // 是否由渠道定价显式设定(为 true 时即使 == 0 也不回退)
InputPricePerToken float64 // 每token输入价格 (USD)
InputPricePerTokenPriority float64 // priority service tier 下每token输入价格 (USD)
ImageInputPricePerToken float64 // 图片输入 token 价格 (USD),用于多模态 embedding 等图文不同价场景;为 0 时回退到 InputPricePerToken
OutputPricePerToken float64 // 每token输出价格 (USD)
OutputPricePerTokenPriority float64 // priority service tier 下每token输出价格 (USD)
CacheCreationPricePerToken float64 // 缓存创建每token价格 (USD)
CacheCreationPricePerTokenPriority float64 // priority service tier 下缓存创建每token价格 (USD)
CacheCreationPriceExplicit bool // 是否由渠道/区间定价显式设定(为 true 时即使 == 0 也不回退)
CacheReadPricePerToken float64 // 缓存读取每token价格 (USD)
CacheReadPricePerTokenPriority float64 // priority service tier 下缓存读取每token价格 (USD)
CacheCreation5mPrice float64 // 5分钟缓存创建每token价格 (USD)
CacheCreation1hPrice float64 // 1小时缓存创建每token价格 (USD)
SupportsCacheBreakdown bool // 是否支持详细的缓存分类
LongContextInputThreshold int // 超过阈值后按整次会话提升输入价格
LongContextInputMultiplier float64 // 长上下文整次会话输入倍率
LongContextOutputMultiplier float64 // 长上下文整次会话输出倍率
ImageOutputPricePerToken float64 // 图片输出 token 价格 (USD)
ImageOutputPriceExplicit bool // 是否由渠道定价显式设定(为 true 时即使 == 0 也不回退)
}
const (
@@ -122,7 +124,8 @@ func usePriorityServiceTierPricing(serviceTier string, pricing *ModelPricing) bo
if pricing == nil || normalizeBillingServiceTier(serviceTier) != "priority" {
return false
}
return pricing.InputPricePerTokenPriority > 0 || pricing.OutputPricePerTokenPriority > 0 || pricing.CacheReadPricePerTokenPriority > 0
return pricing.InputPricePerTokenPriority > 0 || pricing.OutputPricePerTokenPriority > 0 ||
pricing.CacheCreationPricePerTokenPriority > 0 || pricing.CacheReadPricePerTokenPriority > 0
}
func serviceTierCostMultiplier(serviceTier string) float64 {
@@ -280,10 +283,46 @@ func (s *BillingService) initFallbackPricing() {
s.fallbackPrices["gpt-5.5"] = s.fallbackPrices["gpt-5.4"]
s.fallbackPrices["gpt-5.5-pro"] = s.fallbackPrices["gpt-5.4"]
// GPT-5.6(sol / terra / luna)暂无独立定价,回退到 GPT-5.4。
s.fallbackPrices["gpt-5.6-sol"] = s.fallbackPrices["gpt-5.4"]
s.fallbackPrices["gpt-5.6-terra"] = s.fallbackPrices["gpt-5.4"]
s.fallbackPrices["gpt-5.6-luna"] = s.fallbackPrices["gpt-5.4"]
// OpenAI GPT-5.6 官方价格(USD/token)。缓存写入为输入价的 1.25 倍。
s.fallbackPrices["gpt-5.6-sol"] = &ModelPricing{
InputPricePerToken: 5e-6,
InputPricePerTokenPriority: 10e-6,
OutputPricePerToken: 30e-6,
OutputPricePerTokenPriority: 60e-6,
CacheCreationPricePerToken: 6.25e-6,
CacheCreationPricePerTokenPriority: 12.5e-6,
CacheReadPricePerToken: 0.5e-6,
CacheReadPricePerTokenPriority: 1e-6,
LongContextInputThreshold: openAIGPT54LongContextInputThreshold,
LongContextInputMultiplier: openAIGPT54LongContextInputMultiplier,
LongContextOutputMultiplier: openAIGPT54LongContextOutputMultiplier,
}
s.fallbackPrices["gpt-5.6-terra"] = &ModelPricing{
InputPricePerToken: 2.5e-6,
InputPricePerTokenPriority: 5e-6,
OutputPricePerToken: 15e-6,
OutputPricePerTokenPriority: 30e-6,
CacheCreationPricePerToken: 3.125e-6,
CacheCreationPricePerTokenPriority: 6.25e-6,
CacheReadPricePerToken: 0.25e-6,
CacheReadPricePerTokenPriority: 0.5e-6,
LongContextInputThreshold: openAIGPT54LongContextInputThreshold,
LongContextInputMultiplier: openAIGPT54LongContextInputMultiplier,
LongContextOutputMultiplier: openAIGPT54LongContextOutputMultiplier,
}
s.fallbackPrices["gpt-5.6-luna"] = &ModelPricing{
InputPricePerToken: 1e-6,
InputPricePerTokenPriority: 2e-6,
OutputPricePerToken: 6e-6,
OutputPricePerTokenPriority: 12e-6,
CacheCreationPricePerToken: 1.25e-6,
CacheCreationPricePerTokenPriority: 2.5e-6,
CacheReadPricePerToken: 0.1e-6,
CacheReadPricePerTokenPriority: 0.2e-6,
LongContextInputThreshold: openAIGPT54LongContextInputThreshold,
LongContextInputMultiplier: openAIGPT54LongContextInputMultiplier,
LongContextOutputMultiplier: openAIGPT54LongContextOutputMultiplier,
}
s.fallbackPrices["gpt-5.4-mini"] = &ModelPricing{
InputPricePerToken: 7.5e-7,
@@ -739,20 +778,21 @@ func (s *BillingService) GetModelPricing(model string) (*ModelPricing, error) {
price1h := litellmPricing.CacheCreationInputTokenCostAbove1hr
enableBreakdown := price1h > 0 && price1h > price5m
return s.applyModelSpecificPricingPolicy(model, &ModelPricing{
InputPricePerToken: litellmPricing.InputCostPerToken,
InputPricePerTokenPriority: litellmPricing.InputCostPerTokenPriority,
OutputPricePerToken: litellmPricing.OutputCostPerToken,
OutputPricePerTokenPriority: litellmPricing.OutputCostPerTokenPriority,
CacheCreationPricePerToken: litellmPricing.CacheCreationInputTokenCost,
CacheReadPricePerToken: litellmPricing.CacheReadInputTokenCost,
CacheReadPricePerTokenPriority: litellmPricing.CacheReadInputTokenCostPriority,
CacheCreation5mPrice: price5m,
CacheCreation1hPrice: price1h,
SupportsCacheBreakdown: enableBreakdown,
LongContextInputThreshold: litellmPricing.LongContextInputTokenThreshold,
LongContextInputMultiplier: litellmPricing.LongContextInputCostMultiplier,
LongContextOutputMultiplier: litellmPricing.LongContextOutputCostMultiplier,
ImageOutputPricePerToken: litellmPricing.OutputCostPerImageToken,
InputPricePerToken: litellmPricing.InputCostPerToken,
InputPricePerTokenPriority: litellmPricing.InputCostPerTokenPriority,
OutputPricePerToken: litellmPricing.OutputCostPerToken,
OutputPricePerTokenPriority: litellmPricing.OutputCostPerTokenPriority,
CacheCreationPricePerToken: litellmPricing.CacheCreationInputTokenCost,
CacheCreationPricePerTokenPriority: litellmPricing.CacheCreationInputTokenCostPriority,
CacheReadPricePerToken: litellmPricing.CacheReadInputTokenCost,
CacheReadPricePerTokenPriority: litellmPricing.CacheReadInputTokenCostPriority,
CacheCreation5mPrice: price5m,
CacheCreation1hPrice: price1h,
SupportsCacheBreakdown: enableBreakdown,
LongContextInputThreshold: litellmPricing.LongContextInputTokenThreshold,
LongContextInputMultiplier: litellmPricing.LongContextInputCostMultiplier,
LongContextOutputMultiplier: litellmPricing.LongContextOutputCostMultiplier,
ImageOutputPricePerToken: litellmPricing.OutputCostPerImageToken,
}), nil
}
}
@@ -794,6 +834,8 @@ func (s *BillingService) GetModelPricingWithChannel(model string, channelPricing
}
if channelPricing.CacheWritePrice != nil {
pricing.CacheCreationPricePerToken = *channelPricing.CacheWritePrice
pricing.CacheCreationPricePerTokenPriority = *channelPricing.CacheWritePrice
pricing.CacheCreationPriceExplicit = true
pricing.CacheCreation5mPrice = *channelPricing.CacheWritePrice
pricing.CacheCreation1hPrice = *channelPricing.CacheWritePrice
}
@@ -867,7 +909,7 @@ func (s *BillingService) CalculateCostUnified(input CostInput) (*CostBreakdown,
// calculateTokenCost 按 token 区间计费
func (s *BillingService) calculateTokenCost(resolved *ResolvedPricing, input CostInput) (*CostBreakdown, error) {
totalContext := input.Tokens.InputTokens + input.Tokens.CacheReadTokens
totalContext := input.Tokens.InputTokens + input.Tokens.CacheCreationTokens + input.Tokens.CacheReadTokens
pricing := input.Resolver.GetIntervalPricing(resolved, totalContext)
if pricing == nil {
@@ -897,6 +939,7 @@ func (s *BillingService) computeTokenBreakdown(
inputPrice := pricing.InputPricePerToken
outputPrice := pricing.OutputPricePerToken
cacheReadPrice := pricing.CacheReadPricePerToken
cacheCreationPrice := pricing.CacheCreationPricePerToken
cacheCreationMultiplier := 1.0
tierMultiplier := 1.0
@@ -910,6 +953,9 @@ func (s *BillingService) computeTokenBreakdown(
if pricing.CacheReadPricePerTokenPriority > 0 {
cacheReadPrice = pricing.CacheReadPricePerTokenPriority
}
if pricing.CacheCreationPricePerTokenPriority > 0 {
cacheCreationPrice = pricing.CacheCreationPricePerTokenPriority
}
} else {
tierMultiplier = serviceTierCostMultiplier(serviceTier)
}
@@ -963,7 +1009,7 @@ func (s *BillingService) computeTokenBreakdown(
}
// 缓存创建费用
bd.CacheCreationCost = s.computeCacheCreationCost(pricing, tokens, cacheCreationMultiplier)
bd.CacheCreationCost = s.computeCacheCreationCost(pricing, tokens, cacheCreationPrice, cacheCreationMultiplier)
bd.CacheReadCost = float64(tokens.CacheReadTokens) * cacheReadPrice
@@ -984,7 +1030,7 @@ func (s *BillingService) computeTokenBreakdown(
// computeCacheCreationCost 计算缓存创建费用(支持 5m/1h 分类或标准计费)。
// multiplier 用于长上下文等场景下的整体价格缩放(普通调用传 1.0 即可)。
func (s *BillingService) computeCacheCreationCost(pricing *ModelPricing, tokens UsageTokens, multiplier float64) float64 {
func (s *BillingService) computeCacheCreationCost(pricing *ModelPricing, tokens UsageTokens, price, multiplier float64) float64 {
if pricing.SupportsCacheBreakdown && (pricing.CacheCreation5mPrice > 0 || pricing.CacheCreation1hPrice > 0) {
if tokens.CacheCreation5mTokens == 0 && tokens.CacheCreation1hTokens == 0 && tokens.CacheCreationTokens > 0 {
// API 未返回 ephemeral 明细,回退到全部按 5m 单价计费
@@ -993,7 +1039,7 @@ func (s *BillingService) computeCacheCreationCost(pricing *ModelPricing, tokens
return float64(tokens.CacheCreation5mTokens)*pricing.CacheCreation5mPrice*multiplier +
float64(tokens.CacheCreation1hTokens)*pricing.CacheCreation1hPrice*multiplier
}
return float64(tokens.CacheCreationTokens) * pricing.CacheCreationPricePerToken * multiplier
return float64(tokens.CacheCreationTokens) * price * multiplier
}
// calculatePerRequestCost 按次/图片计费
@@ -1010,7 +1056,7 @@ func (s *BillingService) calculatePerRequestCost(resolved *ResolvedPricing, inpu
}
if unitPrice == 0 {
totalContext := input.Tokens.InputTokens + input.Tokens.CacheReadTokens
totalContext := input.Tokens.InputTokens + input.Tokens.CacheCreationTokens + input.Tokens.CacheReadTokens
unitPrice = input.Resolver.GetRequestTierPriceByContext(resolved, totalContext)
}
@@ -1057,21 +1103,38 @@ func (s *BillingService) applyModelSpecificPricingPolicy(model string, pricing *
if pricing == nil {
return nil
}
if !isOpenAIGPT54Model(model) {
normalized := normalizeKnownOpenAICodexModel(model)
isGPT56 := isOpenAIGPT56Model(normalized)
usesLegacyLongContextPricing := usesOpenAILegacyLongContextPricing(normalized)
if !isGPT56 && !usesLegacyLongContextPricing {
return pricing
}
if pricing.LongContextInputThreshold > 0 && pricing.LongContextInputMultiplier > 0 && pricing.LongContextOutputMultiplier > 0 {
needsLongContextPolicy := (isGPT56 || usesLegacyLongContextPricing) &&
(pricing.LongContextInputThreshold <= 0 || pricing.LongContextInputMultiplier <= 0 || pricing.LongContextOutputMultiplier <= 0)
needsCacheCreationPolicy := isGPT56 && !pricing.CacheCreationPriceExplicit && (pricing.CacheCreationPricePerToken <= 0 ||
(pricing.InputPricePerTokenPriority > 0 && pricing.CacheCreationPricePerTokenPriority <= 0))
if !needsLongContextPolicy && !needsCacheCreationPolicy {
return pricing
}
cloned := *pricing
if cloned.LongContextInputThreshold <= 0 {
cloned.LongContextInputThreshold = openAIGPT54LongContextInputThreshold
if isGPT56 && !cloned.CacheCreationPriceExplicit {
if cloned.CacheCreationPricePerToken <= 0 {
cloned.CacheCreationPricePerToken = cloned.InputPricePerToken * 1.25
}
if cloned.CacheCreationPricePerTokenPriority <= 0 {
cloned.CacheCreationPricePerTokenPriority = cloned.InputPricePerTokenPriority * 1.25
}
}
if cloned.LongContextInputMultiplier <= 0 {
cloned.LongContextInputMultiplier = openAIGPT54LongContextInputMultiplier
}
if cloned.LongContextOutputMultiplier <= 0 {
cloned.LongContextOutputMultiplier = openAIGPT54LongContextOutputMultiplier
if isGPT56 || usesLegacyLongContextPricing {
if cloned.LongContextInputThreshold <= 0 {
cloned.LongContextInputThreshold = openAIGPT54LongContextInputThreshold
}
if cloned.LongContextInputMultiplier <= 0 {
cloned.LongContextInputMultiplier = openAIGPT54LongContextInputMultiplier
}
if cloned.LongContextOutputMultiplier <= 0 {
cloned.LongContextOutputMultiplier = openAIGPT54LongContextOutputMultiplier
}
}
return &cloned
}
@@ -1083,17 +1146,12 @@ func (s *BillingService) shouldApplySessionLongContextPricing(tokens UsageTokens
if pricing.LongContextInputMultiplier <= 1 && pricing.LongContextOutputMultiplier <= 1 {
return false
}
totalInputTokens := tokens.InputTokens + tokens.CacheReadTokens
totalInputTokens := tokens.InputTokens + tokens.CacheCreationTokens + tokens.CacheReadTokens
return totalInputTokens > pricing.LongContextInputThreshold
}
func isOpenAIGPT54Model(model string) bool {
// 仅当模型字符串实际属于已知 GPT-5/Codex 族时才做归一判定,避免
// normalizeCodexModel 的默认兜底把非 OpenAI 模型(claude-*、gemini-*、gpt-4o)
// 误识别为 gpt-5.4。
normalized := normalizeKnownOpenAICodexModel(model)
return normalized == "gpt-5.4" || normalized == "gpt-5.5" || normalized == "gpt-5.5-pro" ||
normalized == "gpt-5.6-sol" || normalized == "gpt-5.6-terra" || normalized == "gpt-5.6-luna"
func usesOpenAILegacyLongContextPricing(normalized string) bool {
return normalized == "gpt-5.4" || normalized == "gpt-5.5" || normalized == "gpt-5.5-pro"
}
// CalculateCostWithConfig 使用配置中的默认倍率计算费用
@@ -93,7 +93,7 @@ func openAIJSONToolsContainImageGeneration(tools gjson.Result) bool {
}
found := false
tools.ForEach(func(_, item gjson.Result) bool {
if openAIJSONString(item.Get("type")) == "image_generation" {
if isOpenAIImageGenerationType(openAIJSONString(item.Get("type"))) {
found = true
return false
}
@@ -106,12 +106,20 @@ func openAIJSONToolsContainImageGeneration(tools gjson.Result) bool {
return found
}
func isOpenAIImageGenerationType(value string) bool {
return strings.TrimSpace(value) == "image_generation"
}
func isOpenAIImageGenNamespaceName(value string) bool {
return strings.TrimSpace(value) == "image_gen"
}
// isImageGenNamespaceTool detects the Codex namespace-style image generation
// tool declaration: { "type": "namespace", "name": "image_gen", ... }.
// Codex /image uses this instead of the flat { "type": "image_generation" }.
func isImageGenNamespaceTool(tool gjson.Result) bool {
return openAIJSONString(tool.Get("type")) == "namespace" &&
openAIJSONString(tool.Get("name")) == "image_gen"
isOpenAIImageGenNamespaceName(openAIJSONString(tool.Get("name")))
}
// openAIJSONInputContainsImageGenTool scans Responses input items for
@@ -127,27 +135,19 @@ func openAIJSONInputContainsImageGenTool(input gjson.Result) bool {
if openAIJSONString(item.Get("type")) != "additional_tools" {
return true
}
tools := item.Get("tools")
if !tools.IsArray() {
return true
}
tools.ForEach(func(_, tool gjson.Result) bool {
if isImageGenNamespaceTool(tool) {
found = true
return false
}
return true
})
found = openAIJSONToolsContainImageGeneration(item.Get("tools"))
return !found
})
return found
}
func openAIRequestBodyHasImageGenerationTool(body []byte) bool {
func openAIRequestBodyHasImageGenerationDeclaration(body []byte) bool {
if len(body) == 0 || !gjson.ValidBytes(body) {
return false
}
return openAIJSONToolsContainImageGeneration(gjson.GetBytes(body, "tools"))
return openAIJSONToolsContainImageGeneration(gjson.GetBytes(body, "tools")) ||
openAIJSONInputContainsImageGenTool(gjson.GetBytes(body, "input")) ||
openAIJSONToolChoiceSelectsImageGeneration(gjson.GetBytes(body, "tool_choice"))
}
func openAIRequestBodyImageGenerationToolNeedsNormalization(body []byte) bool {
@@ -178,18 +178,24 @@ func openAIJSONToolChoiceSelectsImageGeneration(choice gjson.Result) bool {
return false
}
if choice.Type == gjson.String {
return strings.TrimSpace(choice.String()) == "image_generation"
return isOpenAIImageGenerationType(choice.String())
}
if !choice.IsObject() {
return false
}
if strings.TrimSpace(choice.Get("type").String()) == "image_generation" {
choiceType := openAIJSONString(choice.Get("type"))
if isOpenAIImageGenerationType(choiceType) {
return true
}
if strings.TrimSpace(choice.Get("tool.type").String()) == "image_generation" {
if choiceType == "namespace" &&
(isOpenAIImageGenNamespaceName(openAIJSONString(choice.Get("name"))) ||
isOpenAIImageGenNamespaceName(openAIJSONString(choice.Get("namespace")))) {
return true
}
if strings.TrimSpace(choice.Get("function.name").String()) == "image_generation" {
if tool := choice.Get("tool"); tool.IsObject() && openAIJSONToolChoiceSelectsImageGeneration(tool) {
return true
}
if isOpenAIImageGenerationType(openAIJSONString(choice.Get("function.name"))) {
return true
}
return false
@@ -198,15 +204,21 @@ func openAIJSONToolChoiceSelectsImageGeneration(choice gjson.Result) bool {
func openAIAnyToolChoiceSelectsImageGeneration(choice any) bool {
switch v := choice.(type) {
case string:
return strings.TrimSpace(v) == "image_generation"
return isOpenAIImageGenerationType(v)
case map[string]any:
if strings.TrimSpace(firstNonEmptyString(v["type"])) == "image_generation" {
choiceType := strings.TrimSpace(firstNonEmptyString(v["type"]))
if isOpenAIImageGenerationType(choiceType) {
return true
}
if tool, ok := v["tool"].(map[string]any); ok && strings.TrimSpace(firstNonEmptyString(tool["type"])) == "image_generation" {
if choiceType == "namespace" &&
(isOpenAIImageGenNamespaceName(firstNonEmptyString(v["name"])) ||
isOpenAIImageGenNamespaceName(firstNonEmptyString(v["namespace"]))) {
return true
}
if fn, ok := v["function"].(map[string]any); ok && strings.TrimSpace(firstNonEmptyString(fn["name"])) == "image_generation" {
if tool, ok := v["tool"].(map[string]any); ok && openAIAnyToolChoiceSelectsImageGeneration(tool) {
return true
}
if fn, ok := v["function"].(map[string]any); ok && isOpenAIImageGenerationType(firstNonEmptyString(fn["name"])) {
return true
}
}
@@ -41,6 +41,20 @@ func TestIsImageGenerationIntent(t *testing.T) {
body: []byte(`{"model":"gpt-5.4","tool_choice":{"type":"image_generation"}}`),
want: true,
},
{
name: "namespace image_gen tool choice",
endpoint: "/v1/responses",
model: "gpt-5.5",
body: []byte(`{"model":"gpt-5.5","tool_choice":{"type":"namespace","name":"image_gen"}}`),
want: true,
},
{
name: "custom imagegen function tool choice is not image intent",
endpoint: "/v1/responses",
model: "gpt-5.5",
body: []byte(`{"model":"gpt-5.5","tool_choice":{"function":{"name":"imagegen"}}}`),
want: false,
},
{
name: "required tool choice alone is text",
endpoint: "/v1/responses",
@@ -62,6 +76,13 @@ func TestIsImageGenerationIntent(t *testing.T) {
body: []byte(`{"model":"gpt-5.5","tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}]}`),
want: true,
},
{
name: "custom namespace with nested imagegen function is not image intent",
endpoint: "/v1/responses",
model: "gpt-5.5",
body: []byte(`{"model":"gpt-5.5","tools":[{"type":"namespace","name":"media_tools","tools":[{"type":"function","name":"imagegen"}]}]}`),
want: false,
},
{
name: "namespace image_gen in input additional_tools (Responses Lite)",
endpoint: "/v1/responses",
@@ -118,6 +139,40 @@ func TestIsImageGenerationIntentMap_NamespaceImageGen(t *testing.T) {
},
want: true,
},
{
name: "custom namespace with nested imagegen function is not image intent",
reqBody: map[string]any{
"model": "gpt-5.5",
"tools": []any{
map[string]any{
"type": "namespace",
"name": "media_tools",
"tools": []any{
map[string]any{"type": "function", "name": "imagegen"},
},
},
},
},
want: false,
},
{
name: "namespace image_gen tool choice",
reqBody: map[string]any{
"model": "gpt-5.5",
"tool_choice": map[string]any{"type": "namespace", "name": "image_gen"},
},
want: true,
},
{
name: "custom imagegen function tool choice is not image intent",
reqBody: map[string]any{
"model": "gpt-5.5",
"tool_choice": map[string]any{
"function": map[string]any{"name": "imagegen"},
},
},
want: false,
},
{
name: "non-image namespace not flagged",
reqBody: map[string]any{
@@ -183,6 +183,8 @@ func (r *ModelPricingResolver) applyTokenOverrides(chPricing *ChannelModelPricin
}
if chPricing.CacheWritePrice != nil {
resolved.BasePricing.CacheCreationPricePerToken = *chPricing.CacheWritePrice
resolved.BasePricing.CacheCreationPricePerTokenPriority = *chPricing.CacheWritePrice
resolved.BasePricing.CacheCreationPriceExplicit = true
resolved.BasePricing.CacheCreation5mPrice = *chPricing.CacheWritePrice
resolved.BasePricing.CacheCreation1hPrice = *chPricing.CacheWritePrice
}
@@ -251,6 +253,8 @@ func intervalToModelPricing(iv *PricingInterval, supportsCacheBreakdown bool, ch
}
if iv.CacheWritePrice != nil {
pricing.CacheCreationPricePerToken = *iv.CacheWritePrice
pricing.CacheCreationPricePerTokenPriority = *iv.CacheWritePrice
pricing.CacheCreationPriceExplicit = true
pricing.CacheCreation5mPrice = *iv.CacheWritePrice
pricing.CacheCreation1hPrice = *iv.CacheWritePrice
}
@@ -114,6 +114,52 @@ func TestGetIntervalPricing_NoMatch_FallsBackToBase(t *testing.T) {
require.Equal(t, basePricing, result)
}
func TestGPT56ExplicitZeroCacheWritePriceIsPreserved(t *testing.T) {
bs := &BillingService{}
resolver := NewModelPricingResolver(nil, bs)
zero := 0.0
t.Run("flat channel price", func(t *testing.T) {
resolved := &ResolvedPricing{
Mode: BillingModeToken,
BasePricing: &ModelPricing{
InputPricePerToken: 5e-6,
OutputPricePerToken: 30e-6,
},
}
resolver.applyTokenOverrides(&ChannelModelPricing{CacheWritePrice: &zero}, resolved)
require.True(t, resolved.BasePricing.CacheCreationPriceExplicit)
cost, err := bs.CalculateCostUnified(CostInput{
Model: "gpt-5.6-sol",
Tokens: UsageTokens{CacheCreationTokens: 100},
RateMultiplier: 1,
Resolver: resolver,
Resolved: resolved,
})
require.NoError(t, err)
require.Zero(t, cost.CacheCreationCost)
})
t.Run("interval price", func(t *testing.T) {
pricing := intervalToModelPricing(&PricingInterval{CacheWritePrice: &zero}, false, nil)
require.True(t, pricing.CacheCreationPriceExplicit)
cost, err := bs.CalculateCostUnified(CostInput{
Model: "gpt-5.6-sol",
Tokens: UsageTokens{CacheCreationTokens: 100},
RateMultiplier: 1,
Resolver: resolver,
Resolved: &ResolvedPricing{
Mode: BillingModeToken,
BasePricing: pricing,
},
})
require.NoError(t, err)
require.Zero(t, cost.CacheCreationCost)
})
}
func TestGetRequestTierPrice(t *testing.T) {
bs := newTestBillingServiceForResolver()
r := NewModelPricingResolver(&ChannelService{}, bs)
@@ -0,0 +1,34 @@
package service
import (
"net/http"
"strings"
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
)
// codexUpstreamMinVersion 上游 /backend-api/codex 接受的最低 version 头:
// 若请求携带 version 且低于该值,上游直接 404(issue #3901,2026-07 实测)。
const codexUpstreamMinVersion = "0.144.0"
// enforceCodexIdentityHeaders 收口 OAuth(ChatGPT 内部接口)出站请求的客户端身份头。
// 上游要求 originator 与 User-Agent 首段配套且为官方客户端标识,version 头(若携带)
// 不低于 0.144.0,任一不满足即 404(issue #3901)。以最终 User-Agent 为准推导配套
// originator;推导不出官方身份(第三方 UA / UA 缺失)时整体回退为默认 Codex CLI 身份。
//
// 仅对携带 originator 的请求生效——compat messages bridge 故意不带 originator,保持原样。
// 必须在所有 User-Agent 改写(自定义 UA / ForceCodexCLI / 浏览器 UA 兜底)之后调用。
func enforceCodexIdentityHeaders(h http.Header) {
if h == nil || h.Get("originator") == "" {
return
}
originator, pairedUA, ok := openai.PairCodexClientIdentity(h.Get("user-agent"))
if !ok {
originator, pairedUA = "codex_cli_rs", codexCLIUserAgent
}
h.Set("user-agent", pairedUA)
h.Set("originator", originator)
if v := strings.TrimSpace(h.Get("version")); v != "" && CompareVersions(v, codexUpstreamMinVersion) < 0 {
h.Set("version", codexCLIVersion)
}
}
@@ -0,0 +1,114 @@
package service
import (
"net/http"
"testing"
"github.com/stretchr/testify/require"
)
func TestEnforceCodexIdentityHeaders(t *testing.T) {
const tuiUA = "codex-tui/0.140.2 (Mac OS X 14.0; arm64) iTerm (codex-tui; 0.140.2)"
tests := []struct {
name string
originator string
userAgent string
version string
wantOriginator string
wantUA string
wantVersion string
}{
{
name: "错配 originator 按最终 UA 重配",
originator: "codex_cli_rs",
userAgent: tuiUA,
wantOriginator: "codex-tui",
wantUA: tuiUA,
},
{
name: "官方配套身份原样保留",
originator: "codex-tui",
userAgent: tuiUA,
wantOriginator: "codex-tui",
wantUA: tuiUA,
},
{
name: "第三方 UA 整体回退默认身份",
originator: "opencode",
userAgent: "luna/1.0.0",
wantOriginator: "codex_cli_rs",
wantUA: codexCLIUserAgent,
},
{
name: "UA 缺失回退默认身份",
originator: "codex_vscode",
wantOriginator: "codex_cli_rs",
wantUA: codexCLIUserAgent,
},
{
name: "originator override UA 首段被尾部真实身份重写",
originator: "cccc",
userAgent: "cccc/0.142.0 (Ubuntu 22.4.0; x86_64) screen (codex-tui; 0.142.0)",
wantOriginator: "codex-tui",
wantUA: "codex-tui/0.142.0 (Ubuntu 22.4.0; x86_64) screen (codex-tui; 0.142.0)",
},
{
name: "低于门槛的 version 提升为内置版本",
originator: "codex_cli_rs",
userAgent: "codex_cli_rs/0.125.0",
version: "0.125.0",
wantOriginator: "codex_cli_rs",
wantUA: "codex_cli_rs/0.125.0",
wantVersion: codexCLIVersion,
},
{
name: "达标 version 原样保留",
originator: "codex_cli_rs",
userAgent: "codex_cli_rs/0.145.0",
version: "0.145.0",
wantOriginator: "codex_cli_rs",
wantUA: "codex_cli_rs/0.145.0",
wantVersion: "0.145.0",
},
{
name: "未携带 version 不注入",
originator: "codex_cli_rs",
userAgent: "codex_cli_rs/0.98.0",
wantOriginator: "codex_cli_rs",
wantUA: "codex_cli_rs/0.98.0",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
h := make(http.Header)
if tt.originator != "" {
h.Set("originator", tt.originator)
}
if tt.userAgent != "" {
h.Set("user-agent", tt.userAgent)
}
if tt.version != "" {
h.Set("version", tt.version)
}
enforceCodexIdentityHeaders(h)
require.Equal(t, tt.wantOriginator, h.Get("originator"))
require.Equal(t, tt.wantUA, h.Get("user-agent"))
require.Equal(t, tt.wantVersion, h.Get("version"))
})
}
}
// compat messages bridge 故意不带 originator:收口必须保持 no-op,不得注入身份头。
func TestEnforceCodexIdentityHeaders_NoOriginatorIsNoop(t *testing.T) {
h := make(http.Header)
h.Set("user-agent", "luna/1.0.0")
enforceCodexIdentityHeaders(h)
require.Empty(t, h.Get("originator"))
require.Equal(t, "luna/1.0.0", h.Get("user-agent"))
}
@@ -596,7 +596,7 @@ func hasOpenAIImageGenerationTool(reqBody map[string]any) bool {
if toolsContainImageGeneration(reqBody["tools"]) {
return true
}
return inputContainsImageGenNamespace(reqBody["input"])
return inputContainsImageGenerationTool(reqBody["input"])
}
func toolsContainImageGeneration(rawTools any) bool {
@@ -612,22 +612,24 @@ func toolsContainImageGeneration(rawTools any) bool {
if !ok {
continue
}
if strings.TrimSpace(firstNonEmptyString(toolMap["type"])) == "image_generation" {
return true
}
if isImageGenNamespaceToolMap(toolMap) {
if isOpenAIImageGenerationToolMap(toolMap) {
return true
}
}
return false
}
func isImageGenNamespaceToolMap(tool map[string]any) bool {
return strings.TrimSpace(firstNonEmptyString(tool["type"])) == "namespace" &&
strings.TrimSpace(firstNonEmptyString(tool["name"])) == "image_gen"
func isOpenAIImageGenerationToolMap(tool map[string]any) bool {
return isOpenAIImageGenerationType(firstNonEmptyString(tool["type"])) ||
isImageGenNamespaceToolMap(tool)
}
func inputContainsImageGenNamespace(rawInput any) bool {
func isImageGenNamespaceToolMap(tool map[string]any) bool {
return strings.TrimSpace(firstNonEmptyString(tool["type"])) == "namespace" &&
isOpenAIImageGenNamespaceName(firstNonEmptyString(tool["name"]))
}
func inputContainsImageGenerationTool(rawInput any) bool {
input, ok := rawInput.([]any)
if !ok {
return false
@@ -647,54 +649,110 @@ func inputContainsImageGenNamespace(rawInput any) bool {
return false
}
// stripOpenAIImageGenerationTools keeps account-level strip policy symmetric
// across standard Responses tools, Responses Lite additional_tools, and tool_choice.
func stripOpenAIImageGenerationTools(reqBody map[string]any) bool {
rawTools, ok := reqBody["tools"]
if reqBody == nil {
return false
}
modified := stripOpenAIImageGenerationToolList(reqBody, "tools")
if stripOpenAIImageGenerationToolsFromInput(reqBody) {
modified = true
}
if openAIAnyToolChoiceSelectsImageGeneration(reqBody["tool_choice"]) {
delete(reqBody, "tool_choice")
modified = true
}
return modified
}
func stripOpenAIImageGenerationToolList(container map[string]any, key string) bool {
rawTools, ok := container[key]
if !ok || rawTools == nil {
if openAIAnyToolChoiceSelectsImageGeneration(reqBody["tool_choice"]) {
delete(reqBody, "tool_choice")
return true
}
return false
}
tools, ok := rawTools.([]any)
if !ok {
if openAIAnyToolChoiceSelectsImageGeneration(reqBody["tool_choice"]) {
delete(reqBody, "tool_choice")
return true
}
return false
}
filtered := make([]any, 0, len(tools))
removed := false
for _, rawTool := range tools {
if toolMap, ok := rawTool.(map[string]any); ok &&
strings.TrimSpace(firstNonEmptyString(toolMap["type"])) == "image_generation" {
if toolMap, ok := rawTool.(map[string]any); ok && isOpenAIImageGenerationToolMap(toolMap) {
removed = true
continue
}
filtered = append(filtered, rawTool)
}
if !removed && !openAIAnyToolChoiceSelectsImageGeneration(reqBody["tool_choice"]) {
if !removed {
return false
}
if removed {
if len(filtered) == 0 {
delete(reqBody, "tools")
} else {
reqBody["tools"] = filtered
}
}
if openAIAnyToolChoiceSelectsImageGeneration(reqBody["tool_choice"]) {
delete(reqBody, "tool_choice")
if len(filtered) == 0 {
delete(container, key)
} else {
container[key] = filtered
}
return true
}
// stripCodexSparkImageGenerationTools removes image_generation tool entries from
// reqBody["tools"]. gpt-5.3-codex-spark rejects that tool upstream with HTTP 400
// (invalid_request_error, param=tools), and Codex CLI advertises it by default, so
// it must be dropped for spark. When the tools list becomes empty the key is removed.
// Returns true when the body was modified.
func stripOpenAIImageGenerationToolsFromInput(reqBody map[string]any) bool {
input, ok := reqBody["input"].([]any)
if !ok {
return false
}
filteredInput := make([]any, 0, len(input))
modified := false
for _, rawItem := range input {
item, ok := rawItem.(map[string]any)
if !ok || strings.TrimSpace(firstNonEmptyString(item["type"])) != "additional_tools" {
filteredInput = append(filteredInput, rawItem)
continue
}
if !stripOpenAIImageGenerationToolList(item, "tools") {
filteredInput = append(filteredInput, rawItem)
continue
}
modified = true
if _, hasTools := item["tools"]; hasTools {
filteredInput = append(filteredInput, rawItem)
}
// An empty additional_tools carrier is not useful upstream; drop the item
// after its only declared capability has been removed.
}
if modified {
reqBody["input"] = filteredInput
}
return modified
}
// stripOpenAIImageGenerationToolsFromRawPayload is the shared adapter for paths
// that forward raw HTTP or WebSocket payloads without the normal request map.
func stripOpenAIImageGenerationToolsFromRawPayload(payload []byte) ([]byte, bool, error) {
if !openAIRequestBodyHasImageGenerationDeclaration(payload) {
if json.Valid(payload) {
return payload, false, nil
}
var invalidPayload map[string]any
return payload, false, json.Unmarshal(payload, &invalidPayload)
}
payloadMap := make(map[string]any)
if err := json.Unmarshal(payload, &payloadMap); err != nil {
return payload, false, err
}
if !stripOpenAIImageGenerationTools(payloadMap) {
return payload, false, nil
}
rebuilt, err := json.Marshal(payloadMap)
if err != nil {
return payload, false, err
}
return rebuilt, true, nil
}
// stripCodexSparkImageGenerationTools removes image tool declarations and choices.
// gpt-5.3-codex-spark rejects those capabilities upstream, while Codex clients may
// advertise them by default.
func stripCodexSparkImageGenerationTools(reqBody map[string]any) bool {
return stripOpenAIImageGenerationTools(reqBody)
}
@@ -796,6 +796,110 @@ func TestApplyCodexOAuthTransform_StripsImageGenerationToolForSparkAlias(t *test
require.False(t, hasTools)
}
func TestStripOpenAIImageGenerationTools_StripsNamespaceFormats(t *testing.T) {
imageNamespace := func() map[string]any {
return map[string]any{
"type": "namespace",
"name": "image_gen",
"tools": []any{
map[string]any{"type": "function", "name": "imagegen"},
},
}
}
codeNamespace := func() map[string]any {
return map[string]any{
"type": "namespace",
"name": "code_tools",
"tools": []any{
map[string]any{"type": "function", "name": "run"},
},
}
}
reqBody := map[string]any{
"model": "gpt-5.5",
"tools": []any{
map[string]any{"type": "function", "name": "shell"},
imageNamespace(),
codeNamespace(),
},
"input": []any{
map[string]any{"type": "message", "role": "user", "content": "hello"},
map[string]any{
"type": "additional_tools",
"tools": []any{imageNamespace(), codeNamespace()},
},
map[string]any{
"type": "additional_tools",
"tools": []any{imageNamespace()},
},
},
"tool_choice": map[string]any{"type": "namespace", "name": "image_gen"},
}
require.True(t, stripOpenAIImageGenerationTools(reqBody))
require.False(t, hasOpenAIImageGenerationTool(reqBody))
require.NotContains(t, reqBody, "tool_choice")
tools, ok := reqBody["tools"].([]any)
require.True(t, ok)
require.Len(t, tools, 2)
firstTool, ok := tools[0].(map[string]any)
require.True(t, ok)
secondTool, ok := tools[1].(map[string]any)
require.True(t, ok)
require.Equal(t, "shell", firstTool["name"])
require.Equal(t, "code_tools", secondTool["name"])
input, ok := reqBody["input"].([]any)
require.True(t, ok)
require.Len(t, input, 2)
message, ok := input[0].(map[string]any)
require.True(t, ok)
require.Equal(t, "message", message["type"])
additionalToolsItem, ok := input[1].(map[string]any)
require.True(t, ok)
additionalTools, ok := additionalToolsItem["tools"].([]any)
require.True(t, ok)
require.Len(t, additionalTools, 1)
additionalTool, ok := additionalTools[0].(map[string]any)
require.True(t, ok)
require.Equal(t, "code_tools", additionalTool["name"])
require.False(t, stripOpenAIImageGenerationTools(reqBody), "stripping should be idempotent")
}
func TestStripOpenAIImageGenerationTools_KeepsNonImageNamespaces(t *testing.T) {
reqBody := map[string]any{
"tools": []any{
map[string]any{"type": "namespace", "name": "code_tools"},
},
"input": []any{
map[string]any{
"type": "additional_tools",
"tools": []any{
map[string]any{"type": "namespace", "name": "browser_tools"},
},
},
},
"tool_choice": "auto",
}
require.False(t, stripOpenAIImageGenerationTools(reqBody))
require.Equal(t, "auto", reqBody["tool_choice"])
require.False(t, hasOpenAIImageGenerationTool(reqBody))
}
func TestStripOpenAIImageGenerationTools_KeepsCustomImagegenFunctionChoice(t *testing.T) {
reqBody := map[string]any{
"tool_choice": map[string]any{
"function": map[string]any{"name": "imagegen"},
},
}
require.False(t, stripOpenAIImageGenerationTools(reqBody))
require.Contains(t, reqBody, "tool_choice")
}
// Non-spark Codex models support image_generation; the tool must be preserved.
func TestApplyCodexOAuthTransform_KeepsImageGenerationToolForNonSpark(t *testing.T) {
reqBody := map[string]any{
@@ -0,0 +1,21 @@
//go:build unit
package service
import (
"strings"
"testing"
"github.com/stretchr/testify/require"
)
func TestCodexVersionConstants_Consistency(t *testing.T) {
require.Equal(t, codexCLIVersion, openAICodexProbeVersion,
"codexCLIVersion and openAICodexProbeVersion must stay in sync")
require.True(t, strings.Contains(codexCLIUserAgent, "codex_cli_rs/"+codexCLIVersion),
"codexCLIUserAgent must embed codexCLIVersion")
require.True(t, strings.Contains(DefaultOpenAICodexUserAgent, codexCLIVersion),
"DefaultOpenAICodexUserAgent must embed codexCLIVersion")
}
@@ -0,0 +1,233 @@
package service
import (
"net/http"
"sync"
"time"
"github.com/gin-gonic/gin"
)
// openAICompactSSEKeepaliveKey 存放 body-signal compact 请求的下游 SSE 心跳器。
const openAICompactSSEKeepaliveKey = "openai_compact_sse_keepalive"
// openAICompactSSEKeepalive 在 compact 上游 unary 等待期间向下游写 SSE 注释行
// 心跳。上游 /responses/compact 在模型处理期间不发送任何字节(大上下文可长达
// 数分钟),下游若经过反向代理(Nginx/Cloudflare Tunnel 等),零字节静默会触发
// 代理的空闲/读超时并掐断连接,Codex 只会盲目重连并重复消耗上游 compact
// 配额(#3887)。SSE 注释行在 eventsource 解析层被直接忽略,不会进入客户端
// 事件流。
//
// 首拍延迟一个 interval:绝大多数硬错误(鉴权/参数/限流)在此之前返回,仍走
// 原 JSON+状态码链路(Codex 按 HTTP 状态码重试);首拍之后状态码固化为 200,
// 后续错误由写回方降级为 response.failed 流内终止事件。
type openAICompactSSEKeepalive struct {
mu sync.Mutex
writer gin.ResponseWriter
started bool
stopped bool
// bytes 是心跳已写出的注释字节数。心跳不构成语义响应,handler 的
// "Forward 期间是否已写响应"判定(failover 放弃换号的依据)必须扣除
// 这部分字节,见 OpenAICompactKeepaliveAdjustedWrittenSize。
bytes int
stop chan struct{}
}
// StartOpenAICompactSSEKeepalive 为已标记 body-signal 客户端流式的 compact
// 请求启动下游心跳,返回幂等的停止函数。interval<=0 或请求未标记时为 no-op。
//
// 同时把 c.Writer 替换为 openAICompactKeepaliveWriter:请求 goroutine 的任何
// 响应构造都会先在心跳互斥锁下停拍,未被显式拦截的写回路径(如 Forward
// 内部的本地拒绝)也不会与心跳 goroutine 产生数据竞争或字节交错。
func StartOpenAICompactSSEKeepalive(c *gin.Context, interval time.Duration) func() {
if c == nil || c.Writer == nil || interval <= 0 || !openAICompactClientWantsStream(c) {
return func() {}
}
k := &openAICompactSSEKeepalive{
writer: c.Writer,
stop: make(chan struct{}),
}
c.Set(openAICompactSSEKeepaliveKey, k)
c.Writer = &openAICompactKeepaliveWriter{ResponseWriter: c.Writer, k: k}
var reqDone <-chan struct{}
if c.Request != nil {
reqDone = c.Request.Context().Done()
}
go func() {
timer := time.NewTimer(interval)
defer timer.Stop()
for {
select {
case <-k.stop:
return
case <-reqDone:
return
case <-timer.C:
}
if !k.beat() {
return
}
timer.Reset(interval)
}
}()
return k.Stop
}
// beat 在锁内提交(首次)响应头并写出一条 SSE 注释行;返回 false 表示心跳已
// 停止或下游写入失败,goroutine 应退出。
func (k *openAICompactSSEKeepalive) beat() bool {
k.mu.Lock()
defer k.mu.Unlock()
if k.stopped {
return false
}
if !k.started {
header := k.writer.Header()
header.Set("Content-Type", "text/event-stream")
header.Set("Cache-Control", "no-cache")
header.Set("Connection", "keep-alive")
header.Set("X-Accel-Buffering", "no")
k.writer.WriteHeader(http.StatusOK)
k.started = true
}
n, err := k.writer.Write([]byte(": keepalive\n\n"))
k.bytes += n
if err != nil {
k.stopped = true
return false
}
k.writer.Flush()
return true
}
// Stop 停止心跳;幂等,可与写回路径并发调用。
func (k *openAICompactSSEKeepalive) Stop() {
k.mu.Lock()
k.markStoppedLocked()
k.mu.Unlock()
}
func (k *openAICompactSSEKeepalive) markStoppedLocked() {
if k.stopped {
return
}
k.stopped = true
close(k.stop)
}
// StopOpenAICompactSSEKeepaliveCommitted 停止当前请求的 compact 心跳(若有)
// 并报告响应头是否已被心跳提交为 200。写回方以此决定继续走原 JSON/状态码
// 链路,还是降级为流内终止事件。调用后不会再有心跳字节写出,且经由互斥锁
// 与心跳 goroutine 建立 happens-before,调用方可安全接管 ResponseWriter。
func StopOpenAICompactSSEKeepaliveCommitted(c *gin.Context) bool {
if c == nil {
return false
}
value, ok := c.Get(openAICompactSSEKeepaliveKey)
if !ok {
return false
}
k, ok := value.(*openAICompactSSEKeepalive)
if !ok || k == nil {
return false
}
k.mu.Lock()
k.markStoppedLocked()
committed := k.started
k.mu.Unlock()
return committed
}
// OpenAICompactKeepaliveAdjustedWrittenSize 返回排除 compact 心跳注释字节后
// 的响应已写字节数;无心跳的请求等价于 c.Writer.Size()。心跳字节不构成语义
// 响应——handler 以"Forward 前后 Size 是否变化"判定是否已向客户端写出响应
// (变化则放弃 failover 换号),该判定不得被心跳污染,否则 compact 请求
// 一旦在上游等待期间发过心跳,上游 429/5xx 就不再换号(#3887 加固审计)。
// 仅心跳字节时归一化为 -1(gin 的"未写出"哨兵值),与提交前的快照可比。
func OpenAICompactKeepaliveAdjustedWrittenSize(c *gin.Context) int {
if c == nil || c.Writer == nil {
return -1
}
value, ok := c.Get(openAICompactSSEKeepaliveKey)
if !ok {
return c.Writer.Size()
}
k, ok := value.(*openAICompactSSEKeepalive)
if !ok || k == nil {
return c.Writer.Size()
}
k.mu.Lock()
defer k.mu.Unlock()
size := k.writer.Size()
if size < 0 {
return size
}
if real := size - k.bytes; real > 0 {
return real
}
return -1
}
// openAICompactKeepaliveWriter 包装 gin.ResponseWriter:写侧方法先停拍心跳
// (互斥锁下建立 happens-before),读侧方法仅加锁不停拍——热路径的状态读取
// (如 Forward 前的 Size 快照)不能误杀心跳。心跳 goroutine 直接写内层
// writer(k.writer),不经过本包装器,不会递归。
type openAICompactKeepaliveWriter struct {
gin.ResponseWriter
k *openAICompactSSEKeepalive
}
// suspend 停拍心跳;幂等。任何响应构造(含 Header 访问——写响应必先操作
// 响应头)都视为请求侧接管 ResponseWriter。
func (w *openAICompactKeepaliveWriter) suspend() {
w.k.Stop()
}
func (w *openAICompactKeepaliveWriter) Header() http.Header {
w.suspend()
return w.ResponseWriter.Header()
}
func (w *openAICompactKeepaliveWriter) Write(data []byte) (int, error) {
w.suspend()
return w.ResponseWriter.Write(data)
}
func (w *openAICompactKeepaliveWriter) WriteString(s string) (int, error) {
w.suspend()
return w.ResponseWriter.WriteString(s)
}
func (w *openAICompactKeepaliveWriter) WriteHeader(code int) {
w.suspend()
w.ResponseWriter.WriteHeader(code)
}
func (w *openAICompactKeepaliveWriter) WriteHeaderNow() {
w.suspend()
w.ResponseWriter.WriteHeaderNow()
}
func (w *openAICompactKeepaliveWriter) Flush() {
w.suspend()
w.ResponseWriter.Flush()
}
func (w *openAICompactKeepaliveWriter) Status() int {
w.k.mu.Lock()
defer w.k.mu.Unlock()
return w.ResponseWriter.Status()
}
func (w *openAICompactKeepaliveWriter) Size() int {
w.k.mu.Lock()
defer w.k.mu.Unlock()
return w.ResponseWriter.Size()
}
func (w *openAICompactKeepaliveWriter) Written() bool {
w.k.mu.Lock()
defer w.k.mu.Unlock()
return w.ResponseWriter.Written()
}
@@ -0,0 +1,192 @@
package service
import (
"net/http"
"strings"
"testing"
"time"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
const keepaliveTestInterval = 10 * time.Millisecond
// waitForKeepaliveBeats 等待至少一次心跳写出。读取 recorder 前必须先经
// StopOpenAICompactSSEKeepaliveCommitted 停拍建立 happens-before。
func waitForKeepaliveBeats() {
time.Sleep(20 * keepaliveTestInterval)
}
// stripKeepaliveComments 去掉 SSE 注释块,返回真实事件文本。
func stripKeepaliveComments(body string) string {
var blocks []string
for _, block := range strings.Split(strings.TrimSpace(body), "\n\n") {
if strings.HasPrefix(strings.TrimSpace(block), ":") {
continue
}
blocks = append(blocks, block)
}
return strings.Join(blocks, "\n\n")
}
func TestStartOpenAICompactSSEKeepalive_NoopWhenUnmarkedOrDisabled(t *testing.T) {
// 未标记 client stream:不启动。
c, rec := newCompactBridgeTestContext(t, false)
stop := StartOpenAICompactSSEKeepalive(c, keepaliveTestInterval)
waitForKeepaliveBeats()
stop()
require.Zero(t, rec.Body.Len())
require.False(t, StopOpenAICompactSSEKeepaliveCommitted(c))
// interval=0(配置禁用):不启动。
c, rec = newCompactBridgeTestContext(t, true)
stop = StartOpenAICompactSSEKeepalive(c, 0)
waitForKeepaliveBeats()
stop()
require.Zero(t, rec.Body.Len())
require.False(t, StopOpenAICompactSSEKeepaliveCommitted(c))
}
func TestOpenAICompactSSEKeepalive_CommitsHeadersAndComments(t *testing.T) {
c, rec := newCompactBridgeTestContext(t, true)
stop := StartOpenAICompactSSEKeepalive(c, keepaliveTestInterval)
defer stop()
waitForKeepaliveBeats()
require.True(t, StopOpenAICompactSSEKeepaliveCommitted(c))
require.Equal(t, http.StatusOK, rec.Code)
require.Equal(t, "text/event-stream", rec.Header().Get("Content-Type"))
require.Equal(t, "no", rec.Header().Get("X-Accel-Buffering"))
require.Contains(t, rec.Body.String(), ": keepalive\n\n")
}
func TestOpenAICompactSSEKeepalive_StopBeforeFirstBeatKeepsWriterUntouched(t *testing.T) {
c, rec := newCompactBridgeTestContext(t, true)
stop := StartOpenAICompactSSEKeepalive(c, time.Hour)
stop()
waitForKeepaliveBeats()
require.Zero(t, rec.Body.Len())
require.False(t, StopOpenAICompactSSEKeepaliveCommitted(c))
}
// 心跳已提交后,2xx 桥接续写事件而不重复提交响应头。
func TestWriteOpenAICompactSSEBridge_AfterKeepaliveCommitAppendsEvents(t *testing.T) {
c, rec := newCompactBridgeTestContext(t, true)
stop := StartOpenAICompactSSEKeepalive(c, keepaliveTestInterval)
defer stop()
waitForKeepaliveBeats()
finalResponse := []byte(`{"id":"resp_ka_1","output":[{"id":"cmp_ka","type":"compaction","encrypted_content":"x"}],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}`)
require.True(t, writeOpenAICompactSSEBridge(c, http.StatusOK, finalResponse))
require.Equal(t, http.StatusOK, rec.Code)
events := parseCompactBridgeSSE(t, stripKeepaliveComments(rec.Body.String()))
require.Len(t, events, 2)
require.Equal(t, "response.output_item.done", events[0][0])
require.Equal(t, "compaction", gjson.Get(events[0][1], "item.type").String())
require.Equal(t, "response.completed", events[1][0])
require.Equal(t, "resp_ka_1", gjson.Get(events[1][1], "response.id").String())
}
// 心跳已提交后上游非 2xx:状态码无法回传,必须以 response.failed 终止事件
// 收尾(Codex 将其作为终止事件处理),并标记流内错误供 ops 采集。
func TestWriteOpenAICompactSSEBridge_AfterKeepaliveCommitFailureEmitsFailedEvent(t *testing.T) {
c, rec := newCompactBridgeTestContext(t, true)
stop := StartOpenAICompactSSEKeepalive(c, keepaliveTestInterval)
defer stop()
waitForKeepaliveBeats()
require.True(t, writeOpenAICompactSSEBridge(c, http.StatusBadGateway, []byte(`{"error":{"message":"upstream exploded"}}`)))
events := parseCompactBridgeSSE(t, stripKeepaliveComments(rec.Body.String()))
require.Len(t, events, 1)
require.Equal(t, "response.failed", events[0][0])
require.Equal(t, "failed", gjson.Get(events[0][1], "response.status").String())
require.Contains(t, gjson.Get(events[0][1], "response.error.message").String(), "upstream exploded")
require.NotEmpty(t, gjson.Get(events[0][1], "response.id").String())
streamErr, ok := GetOpsStreamError(c)
require.True(t, ok)
require.Equal(t, http.StatusBadGateway, streamErr.IntendedStatus)
}
// 心跳未提交时非 2xx 行为不变:返回 false,调用方按原 JSON+状态码写回。
func TestWriteOpenAICompactSSEBridge_BeforeKeepaliveCommitFailureKeepsJSONPath(t *testing.T) {
c, rec := newCompactBridgeTestContext(t, true)
stop := StartOpenAICompactSSEKeepalive(c, time.Hour)
stop()
require.False(t, writeOpenAICompactSSEBridge(c, http.StatusBadGateway, []byte(`{"error":{"message":"fast fail"}}`)))
require.Zero(t, rec.Body.Len())
}
// 未被显式拦截的写回路径(直接操作 c.Writer)也必须与心跳互斥:包装器在
// 请求侧任何响应构造时停拍。-race 下验证无数据竞争,且停拍后不再有心跳
// 字节写出。
func TestOpenAICompactKeepaliveWriter_RequestSideWriteSuspendsBeats(t *testing.T) {
c, rec := newCompactBridgeTestContext(t, true)
stop := StartOpenAICompactSSEKeepalive(c, keepaliveTestInterval)
defer stop()
waitForKeepaliveBeats()
// 模拟未拦截路径的直接写回(如 Forward 内部本地拒绝的 c.JSON)。
_, err := c.Writer.Write([]byte(`{"error":"local reject"}`))
require.NoError(t, err)
lenAfterWrite := rec.Body.Len()
waitForKeepaliveBeats()
require.Equal(t, lenAfterWrite, rec.Body.Len(), "请求侧写回后心跳必须停止")
require.Contains(t, rec.Body.String(), ": keepalive\n\n")
require.Contains(t, rec.Body.String(), `{"error":"local reject"}`)
}
// fast policy block 在心跳提交后必须降级为 response.failed 终止事件。
func TestWriteOpenAIFastPolicyBlockedResponse_AfterKeepaliveCommit(t *testing.T) {
c, rec := newCompactBridgeTestContext(t, true)
stop := StartOpenAICompactSSEKeepalive(c, keepaliveTestInterval)
defer stop()
waitForKeepaliveBeats()
writeOpenAIFastPolicyBlockedResponse(c, &OpenAIFastBlockedError{Message: "tier blocked"})
require.Equal(t, http.StatusOK, rec.Code)
events := parseCompactBridgeSSE(t, stripKeepaliveComments(rec.Body.String()))
require.Len(t, events, 1)
require.Equal(t, "response.failed", events[0][0])
require.Equal(t, "permission_error", gjson.Get(events[0][1], "response.error.code").String())
require.Contains(t, gjson.Get(events[0][1], "response.error.message").String(), "tier blocked")
}
// failover"是否已写响应"判定的口径:心跳字节必须被排除,否则 compact 在
// 上游等待期间发过心跳后,可换号的 failover 会被误判放弃;真实响应字节
// 写出后口径必须变化。
func TestOpenAICompactKeepaliveAdjustedWrittenSize_ExcludesHeartbeatBytes(t *testing.T) {
c, rec := newCompactBridgeTestContext(t, true)
// 无心跳的请求:等价于 c.Writer.Size()。
require.Equal(t, c.Writer.Size(), OpenAICompactKeepaliveAdjustedWrittenSize(c))
stop := StartOpenAICompactSSEKeepalive(c, keepaliveTestInterval)
defer stop()
before := OpenAICompactKeepaliveAdjustedWrittenSize(c)
waitForKeepaliveBeats()
require.Equal(t, before, OpenAICompactKeepaliveAdjustedWrittenSize(c), "仅心跳字节不得改变判定口径")
// 真实响应字节写出(经包装器,先停拍再写)后口径必须变化。
_, err := c.Writer.Write([]byte("real-bytes"))
require.NoError(t, err)
require.Equal(t, len("real-bytes"), OpenAICompactKeepaliveAdjustedWrittenSize(c))
require.Contains(t, rec.Body.String(), ": keepalive\n\n")
}
// fast policy block 在心跳未提交时保持 403 JSON 原语义。
func TestWriteOpenAIFastPolicyBlockedResponse_BeforeKeepaliveCommit(t *testing.T) {
c, rec := newCompactBridgeTestContext(t, true)
stop := StartOpenAICompactSSEKeepalive(c, time.Hour)
defer stop()
writeOpenAIFastPolicyBlockedResponse(c, &OpenAIFastBlockedError{Message: "tier blocked"})
require.Equal(t, http.StatusForbidden, rec.Code)
require.Equal(t, "permission_error", gjson.Get(rec.Body.String(), "error.type").String())
}
@@ -3,6 +3,8 @@ package service
import (
"bytes"
"encoding/json"
"net/http"
"strconv"
"strings"
"github.com/gin-gonic/gin"
@@ -48,25 +50,90 @@ func openAICompactClientWantsStream(c *gin.Context) bool {
// compact v2 的消费协议合成为最小 Responses SSE 流写回客户端。仅当请求被标记
// 为 body-signal 客户端流式、状态码为 2xx 且 body 是合法 JSON 对象时生效;
// 返回 false 表示未写出任何内容,调用方应按原路径写回。
//
// 若下游心跳已把响应头提交为 200(见 openAICompactSSEKeepalive),则本函数
// 必须接管一切写回:非 2xx 或不可合成的响应降级为 response.failed 终止事件,
// 不能再返回 false(否则调用方的 JSON 写回会与已提交的 SSE 流交错)。
func writeOpenAICompactSSEBridge(c *gin.Context, statusCode int, finalResponse []byte) bool {
if c == nil || statusCode < 200 || statusCode >= 300 || !openAICompactClientWantsStream(c) {
if c == nil || !openAICompactClientWantsStream(c) {
return false
}
// 先停心跳再写回,避免注释行与最终事件交错;停止后经互斥锁与心跳
// goroutine 建立 happens-before,可安全接管 ResponseWriter。
committed := StopOpenAICompactSSEKeepaliveCommitted(c)
if statusCode < 200 || statusCode >= 300 {
if committed {
writeOpenAICompactSSEFailure(c, statusCode, finalResponse)
return true
}
return false
}
payload, ok := buildOpenAICompactSSEPayload(finalResponse)
if !ok {
if committed {
writeOpenAICompactSSEFailure(c, http.StatusBadGateway, finalResponse)
return true
}
return false
}
header := c.Writer.Header()
header.Set("Content-Type", "text/event-stream")
header.Set("Cache-Control", "no-cache")
header.Set("Connection", "keep-alive")
header.Set("X-Accel-Buffering", "no")
c.Writer.WriteHeader(statusCode)
if !committed {
header := c.Writer.Header()
header.Set("Content-Type", "text/event-stream")
header.Set("Cache-Control", "no-cache")
header.Set("Connection", "keep-alive")
header.Set("X-Accel-Buffering", "no")
c.Writer.WriteHeader(statusCode)
}
_, _ = c.Writer.Write(payload)
c.Writer.Flush()
return true
}
// writeOpenAICompactSSEFailure 从上游错误 body 提取错误消息后,以
// response.failed 终止事件回传。仅用于心跳已提交 200、无法再按 HTTP 状态码
// 回传错误的场景。
func writeOpenAICompactSSEFailure(c *gin.Context, statusCode int, errorBody []byte) {
message := ""
if len(errorBody) > 0 {
message = sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(errorBody)))
}
if message == "" {
message = "Upstream compact request failed with HTTP " + strconv.Itoa(statusCode)
}
writeOpenAICompactSSEFailureMessage(c, statusCode, "upstream_error", message)
}
// writeOpenAICompactSSEFailureMessage 写出 response.failed 终止事件。Codex 对
// 流式 Responses 请求把 response.failed 作为合法终止事件处理(普通 error 帧
// 不被识别,会退化为 "stream closed before response.completed" 盲重连)。
// 同时标记流内错误,保证挂在 200 流上的失败仍进入 ops 错误看板。
func writeOpenAICompactSSEFailureMessage(c *gin.Context, statusCode int, errType, message string) {
if c == nil {
return
}
MarkOpsStreamError(c, errType, message, statusCode)
payload, err := json.Marshal(map[string]any{
"type": "response.failed",
"response": map[string]any{
"id": "resp_" + strings.ReplaceAll(uuid.NewString(), "-", ""),
"object": "response",
"status": "failed",
"output": []any{},
"error": map[string]any{
"code": errType,
"message": message,
},
},
})
if err != nil {
return
}
_, _ = c.Writer.Write([]byte("event: response.failed\ndata: "))
_, _ = c.Writer.Write(payload)
_, _ = c.Writer.Write([]byte("\n\n"))
c.Writer.Flush()
}
// buildOpenAICompactSSEPayload 把 compact 的 Response JSON 转成 SSE 事件序列:
// 每个 output[] item 一条 response.output_item.done,最后一条 response.completed
// 携带完整 response 对象。Codex 的 SSE 解析只从 output_item.done 收集 item,
@@ -258,6 +258,246 @@ func TestHandleSSEToJSON_CompactClientStreamBridgesToSSE(t *testing.T) {
require.Equal(t, "resp_compact_sse", gjson.Get(events[1][1], "response.id").String())
}
// 回归 #3887(#3777 问题 2):上游对 compact 返回 SSE,compaction item 只在
// raw output_item.done 中、终态 response.completed 的 output 为空。SSE→JSON
// 提取必须保留 raw item 修补终态 output,否则桥接合成 0 个 output_item.done,
// Codex 报 "expected exactly one compaction output item, got 0" 并盲目重试,
// 每次重试都重新计费。fixture 取自 #3777 的上游实录形态。
func TestHandleSSEToJSON_CompactRawOutputItemDoneRepairsEmptyTerminalOutput(t *testing.T) {
svc := newCompactBridgeTestService()
c, rec := newCompactBridgeTestContext(t, true)
upstreamSSE := strings.Join([]string{
`data: {"type":"response.output_item.done","output_index":0,"item":{"id":"cmp_1","type":"compaction_summary","status":"completed","summary":[{"type":"summary_text","text":"compact summary"}],"encrypted_content":"compact-payload","opaque":{"kept":true}}}`,
``,
`data: {"type":"response.completed","response":{"id":"resp_compact","object":"response","model":"gpt-5.1-codex","status":"completed","output":[],"usage":{"input_tokens":9,"output_tokens":4,"total_tokens":13}}}`,
``,
}, "\n")
resp := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(upstreamSSE)),
}
result, err := svc.handleNonStreamingResponse(context.Background(), resp, c, &Account{ID: 1, Type: AccountTypeOAuth}, "gpt-5.5", "gpt-5.5")
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, "text/event-stream", rec.Header().Get("Content-Type"))
events := parseCompactBridgeSSE(t, rec.Body.String())
require.Len(t, events, 2)
require.Equal(t, "response.output_item.done", events[0][0])
item := gjson.Get(events[0][1], "item")
require.Equal(t, "compaction_summary", item.Get("type").String())
require.Equal(t, "cmp_1", item.Get("id").String())
require.Equal(t, "compact-payload", item.Get("encrypted_content").String())
require.Equal(t, "compact summary", item.Get("summary.0.text").String())
require.True(t, item.Get("opaque.kept").Bool(), "raw item 字段必须逐字节保留")
require.Equal(t, "response.completed", events[1][0])
require.Equal(t, "resp_compact", gjson.Get(events[1][1], "response.id").String())
require.Len(t, gjson.Get(events[1][1], "response.output").Array(), 1)
require.Equal(t, int64(13), gjson.Get(events[1][1], "response.usage.total_tokens").Int())
require.NotNil(t, result.usage)
require.Equal(t, 9, result.usage.InputTokens)
require.Equal(t, 4, result.usage.OutputTokens)
}
// 同一形态经透传分支(handlePassthroughSSEToJSON)也必须修补。
func TestHandlePassthroughSSEToJSON_CompactRawOutputItemDoneRepairsEmptyTerminalOutput(t *testing.T) {
svc := newCompactBridgeTestService()
c, rec := newCompactBridgeTestContext(t, true)
upstreamSSE := strings.Join([]string{
`data: {"type":"response.output_item.done","output_index":0,"item":{"id":"cmp_pt_1","type":"compaction","status":"completed","encrypted_content":"compact-pt-raw"}}`,
``,
`data: {"type":"response.completed","response":{"id":"resp_compact_pt_raw","object":"response","status":"completed","output":[],"usage":{"input_tokens":6,"output_tokens":2,"total_tokens":8}}}`,
``,
}, "\n")
resp := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(upstreamSSE)),
}
result, err := svc.handleNonStreamingResponsePassthrough(context.Background(), resp, c, "gpt-5.5", "")
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, "text/event-stream", rec.Header().Get("Content-Type"))
events := parseCompactBridgeSSE(t, rec.Body.String())
require.Len(t, events, 2)
require.Equal(t, "compaction", gjson.Get(events[0][1], "item.type").String())
require.Equal(t, "compact-pt-raw", gjson.Get(events[0][1], "item.encrypted_content").String())
require.Len(t, gjson.Get(events[1][1], "response.output").Array(), 1)
}
// path-based(Codex v1 unary、链式 sub2api)未标记 client stream:同一上游
// 形态修补后仍按 JSON 写回,output 中必须包含 compaction item。
func TestHandleSSEToJSON_PathBasedCompactRawOutputItemDoneRepairsJSON(t *testing.T) {
svc := newCompactBridgeTestService()
c, rec := newCompactBridgeTestContext(t, false)
upstreamSSE := strings.Join([]string{
`data: {"type":"response.output_item.done","output_index":0,"item":{"id":"cmp_v1","type":"compaction_summary","encrypted_content":"compact-v1-raw"}}`,
``,
`data: {"type":"response.completed","response":{"id":"resp_compact_v1","object":"response","status":"completed","output":[],"usage":{"input_tokens":5,"output_tokens":1,"total_tokens":6}}}`,
``,
}, "\n")
resp := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(upstreamSSE)),
}
result, err := svc.handleNonStreamingResponse(context.Background(), resp, c, &Account{ID: 1, Type: AccountTypeOAuth}, "gpt-5.5", "gpt-5.5")
require.NoError(t, err)
require.NotNil(t, result)
// 写回 body 必须是修补后的 JSON 文档(非 SSE 事件流)。
body := rec.Body.String()
require.NotContains(t, body, "event:")
require.NotContains(t, body, "data:")
require.Equal(t, "resp_compact_v1", gjson.Get(body, "id").String())
require.Equal(t, "compaction_summary", gjson.Get(body, "output.0.type").String())
require.Equal(t, "compact-v1-raw", gjson.Get(body, "output.0.encrypted_content").String())
}
// raw done item 是协议上的最终完整形态,优先于 delta 重建且不得重复计入。
func TestReconstructResponseOutputFromSSE_PrefersRawDoneItems(t *testing.T) {
bodyText := strings.Join([]string{
`data: {"type":"response.output_text.delta","delta":"hel"}`,
`data: {"type":"response.output_text.delta","delta":"lo"}`,
`data: {"type":"response.output_item.done","output_index":0,"item":{"id":"msg_1","type":"message","role":"assistant","status":"completed","content":[{"type":"output_text","text":"hello"}]}}`,
`data: {"type":"response.completed","response":{"id":"resp_1","output":[]}}`,
}, "\n")
outputJSON, ok := reconstructResponseOutputFromSSE(bodyText)
require.True(t, ok)
items := gjson.ParseBytes(outputJSON).Array()
require.Len(t, items, 1, "raw done item 与 delta 重建不得重复")
require.Equal(t, "msg_1", items[0].Get("id").String())
require.Equal(t, "hello", items[0].Get("content.0.text").String())
}
// 无任何 done 事件时,退回收集 output_item.added 中的 compaction 类 item。
func TestReconstructResponseOutputFromSSE_CompactionAddedFallback(t *testing.T) {
bodyText := strings.Join([]string{
`data: {"type":"response.output_item.added","output_index":0,"item":{"id":"cmp_add","type":"compaction","encrypted_content":"added-only"}}`,
`data: {"type":"response.completed","response":{"id":"resp_1","output":[]}}`,
}, "\n")
outputJSON, ok := reconstructResponseOutputFromSSE(bodyText)
require.True(t, ok)
items := gjson.ParseBytes(outputJSON).Array()
require.Len(t, items, 1)
require.Equal(t, "compaction", items[0].Get("type").String())
require.Equal(t, "added-only", items[0].Get("encrypted_content").String())
}
// 混合形态:其他 item 有 done、compaction 只在 added 中——compaction 必须
// 被补入;done 已含 compaction 时 added 不得重复计入。
func TestReconstructResponseOutputFromSSE_MixedDoneAndCompactionAdded(t *testing.T) {
bodyText := strings.Join([]string{
`data: {"type":"response.output_item.added","output_index":0,"item":{"id":"cmp_mixed","type":"compaction","encrypted_content":"mixed"}}`,
`data: {"type":"response.output_item.done","output_index":1,"item":{"id":"msg_1","type":"message","content":[{"type":"output_text","text":"hi"}]}}`,
`data: {"type":"response.completed","response":{"id":"resp_1","output":[]}}`,
}, "\n")
outputJSON, ok := reconstructResponseOutputFromSSE(bodyText)
require.True(t, ok)
items := gjson.ParseBytes(outputJSON).Array()
require.Len(t, items, 2)
require.Equal(t, "msg_1", items[0].Get("id").String())
require.Equal(t, "cmp_mixed", items[1].Get("id").String())
// done 已含 compaction:added 中的同一 item(无 id 可去重的最坏情况用
// 不同 raw 表达)不得再收集,Codex 要求恰好一个 compaction item。
bodyText = strings.Join([]string{
`data: {"type":"response.output_item.added","output_index":0,"item":{"type":"compaction","status":"in_progress"}}`,
`data: {"type":"response.output_item.done","output_index":0,"item":{"type":"compaction","status":"completed","encrypted_content":"final"}}`,
`data: {"type":"response.completed","response":{"id":"resp_1","output":[]}}`,
}, "\n")
outputJSON, ok = reconstructResponseOutputFromSSE(bodyText)
require.True(t, ok)
items = gjson.ParseBytes(outputJSON).Array()
require.Len(t, items, 1)
require.Equal(t, "final", items[0].Get("encrypted_content").String())
}
// 上游不一致形态:终态 output 非空(含 message)但 compaction 只在 raw
// output_item.done 中。146 纯流式透传下 Codex 直接读事件流能拿到 compaction,
// SSE→JSON 提取必须补入等价结果。
func TestHandleSSEToJSON_CompactSupplementsMissingCompactionIntoNonEmptyOutput(t *testing.T) {
svc := newCompactBridgeTestService()
c, rec := newCompactBridgeTestContext(t, true)
upstreamSSE := strings.Join([]string{
`data: {"type":"response.output_item.done","output_index":0,"item":{"id":"cmp_sup","type":"compaction","encrypted_content":"supplement"}}`,
``,
`data: {"type":"response.completed","response":{"id":"resp_sup","object":"response","status":"completed","output":[{"id":"msg_sup","type":"message","role":"assistant","content":[{"type":"output_text","text":"note"}]}],"usage":{"input_tokens":2,"output_tokens":1,"total_tokens":3}}}`,
``,
}, "\n")
resp := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(upstreamSSE)),
}
result, err := svc.handleNonStreamingResponse(context.Background(), resp, c, &Account{ID: 1, Type: AccountTypeOAuth}, "gpt-5.5", "gpt-5.5")
require.NoError(t, err)
require.NotNil(t, result)
events := parseCompactBridgeSSE(t, rec.Body.String())
require.Len(t, events, 3)
itemTypes := []string{
gjson.Get(events[0][1], "item.type").String(),
gjson.Get(events[1][1], "item.type").String(),
}
require.Contains(t, itemTypes, "compaction")
require.Contains(t, itemTypes, "message")
require.Equal(t, "response.completed", events[2][0])
require.Len(t, gjson.Get(events[2][1], "response.output").Array(), 2)
}
// 补全逻辑的门控:非 compact 请求原样返回;终态已含 compaction 不重复补入。
func TestSupplementCompactionItemFromSSE_Gating(t *testing.T) {
bodyText := `data: {"type":"response.output_item.done","item":{"id":"cmp_g","type":"compaction","encrypted_content":"g"}}` + "\n"
// 非 compact 路径:不补入。
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
finalResponse := []byte(`{"id":"r1","output":[{"type":"message"}]}`)
require.Equal(t, string(finalResponse), string(supplementCompactionItemFromSSE(c, finalResponse, bodyText)))
// compact 路径 + 终态已含 compaction:不重复补入。
c2, _ := newCompactBridgeTestContext(t, false)
already := []byte(`{"id":"r2","output":[{"type":"compaction","encrypted_content":"x"}]}`)
require.Equal(t, string(already), string(supplementCompactionItemFromSSE(c2, already, bodyText)))
// compact 路径 + 终态非空缺 compaction:补入到末尾。
missing := []byte(`{"id":"r3","output":[{"type":"message"}]}`)
patched := supplementCompactionItemFromSSE(c2, missing, bodyText)
items := gjson.GetBytes(patched, "output").Array()
require.Len(t, items, 2)
require.Equal(t, "compaction", items[1].Get("type").String())
require.Equal(t, "g", items[1].Get("encrypted_content").String())
}
// 非 compaction 的 output_item.added 不参与回退收集(added 阶段的 message
// 通常是空壳),仍走 delta 重建。
func TestReconstructResponseOutputFromSSE_NonCompactionAddedStillUsesDeltas(t *testing.T) {
bodyText := strings.Join([]string{
`data: {"type":"response.output_item.added","output_index":0,"item":{"id":"msg_1","type":"message","content":[]}}`,
`data: {"type":"response.output_text.delta","delta":"hi"}`,
`data: {"type":"response.completed","response":{"id":"resp_1","output":[]}}`,
}, "\n")
outputJSON, ok := reconstructResponseOutputFromSSE(bodyText)
require.True(t, ok)
items := gjson.ParseBytes(outputJSON).Array()
require.Len(t, items, 1)
require.Equal(t, "hi", items[0].Get("content.0.text").String())
}
// 透传分支(OAuth passthrough)同样命中桥接。
func TestHandleNonStreamingResponsePassthrough_CompactClientStreamBridgesToSSE(t *testing.T) {
svc := newCompactBridgeTestService()
+2 -11
View File
@@ -206,17 +206,8 @@ func extractOpenAIEmbeddingsUsage(body []byte) OpenAIUsage {
usage.Get("completion_tokens"),
usage.Get("output_tokens"),
)
cacheReadTokens := firstPositiveGJSONInt(
usage.Get("prompt_tokens_details.cached_tokens"),
usage.Get("input_tokens_details.cached_tokens"),
usage.Get("cache_read_tokens"),
usage.Get("cache_read_input_tokens"),
)
cacheCreationTokens := firstPositiveGJSONInt(
usage.Get("cache_creation_tokens"),
usage.Get("cache_creation_input_tokens"),
usage.Get("input_tokens_details.cache_creation_tokens"),
)
cacheReadTokens := openAICacheReadTokensFromUsage(usage)
cacheCreationTokens := openAICacheCreationTokensFromUsage(usage)
// 多模态 embedding(如 doubao-embedding-vision)回传图文 token 拆分,
// 用于图文不同价计费;纯文本 embedding 该字段为 0,行为不变。
imageInputTokens := firstPositiveGJSONInt(
@@ -7,6 +7,7 @@ import (
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
@@ -138,6 +139,37 @@ func TestEvaluateOpenAIFastPolicy_ScopeFiltersOAuth(t *testing.T) {
require.Equal(t, BetaPolicyActionPass, action)
}
func TestEvaluateOpenAIFastPolicy_UserScopedRuleOverridesGlobalRule(t *testing.T) {
settings := &OpenAIFastPolicySettings{
Rules: []OpenAIFastPolicyRule{
{
ServiceTier: OpenAIFastTierPriority,
Action: BetaPolicyActionFilter,
Scope: BetaPolicyScopeAll,
},
{
ServiceTier: OpenAIFastTierPriority,
Action: BetaPolicyActionPass,
Scope: BetaPolicyScopeAll,
UserIDs: []int64{42},
},
},
}
svc := newOpenAIGatewayServiceWithSettings(t, settings)
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
allowedUserCtx := context.WithValue(context.Background(), ctxkey.UserID, int64(42))
action, _ := svc.evaluateOpenAIFastPolicy(allowedUserCtx, account, "gpt-5.5", OpenAIFastTierPriority)
require.Equal(t, BetaPolicyActionPass, action)
otherUserCtx := context.WithValue(context.Background(), ctxkey.UserID, int64(43))
action, _ = svc.evaluateOpenAIFastPolicy(otherUserCtx, account, "gpt-5.5", OpenAIFastTierPriority)
require.Equal(t, BetaPolicyActionFilter, action)
action, _ = svc.evaluateOpenAIFastPolicy(context.Background(), account, "gpt-5.5", OpenAIFastTierPriority)
require.Equal(t, BetaPolicyActionFilter, action)
}
func TestApplyOpenAIFastPolicyToBody_DefaultPassesPriorityAndFast(t *testing.T) {
svc := newOpenAIGatewayServiceWithSettings(t, DefaultOpenAIFastPolicySettings())
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
@@ -179,6 +211,37 @@ func TestApplyOpenAIFastPolicyToBody_ExplicitFilterRemovesField(t *testing.T) {
require.NotContains(t, string(updated), `"service_tier"`)
}
func TestApplyOpenAIFastPolicyToBody_UserScopedRuleOverridesGlobalRule(t *testing.T) {
settings := &OpenAIFastPolicySettings{
Rules: []OpenAIFastPolicyRule{
{
ServiceTier: OpenAIFastTierPriority,
Action: BetaPolicyActionFilter,
Scope: BetaPolicyScopeAll,
},
{
ServiceTier: OpenAIFastTierPriority,
Action: BetaPolicyActionPass,
Scope: BetaPolicyScopeAll,
UserIDs: []int64{42},
},
},
}
svc := newOpenAIGatewayServiceWithSettings(t, settings)
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
body := []byte(`{"model":"gpt-5.5","service_tier":"priority"}`)
allowedUserCtx := context.WithValue(context.Background(), ctxkey.UserID, int64(42))
updated, err := svc.applyOpenAIFastPolicyToBody(allowedUserCtx, account, "gpt-5.5", body)
require.NoError(t, err)
require.Equal(t, "priority", gjson.GetBytes(updated, "service_tier").String())
otherUserCtx := context.WithValue(context.Background(), ctxkey.UserID, int64(43))
updated, err = svc.applyOpenAIFastPolicyToBody(otherUserCtx, account, "gpt-5.5", body)
require.NoError(t, err)
require.NotContains(t, string(updated), `"service_tier"`)
}
func TestApplyOpenAIFastPolicyToBody_ForcePriorityRewritesKnownTier(t *testing.T) {
settings := &OpenAIFastPolicySettings{
Rules: []OpenAIFastPolicyRule{{
@@ -309,12 +372,34 @@ func TestSetOpenAIFastPolicySettings_Validation(t *testing.T) {
})
require.Error(t, err)
// Non-positive and duplicate user IDs are rejected.
err = svc.SetOpenAIFastPolicySettings(context.Background(), &OpenAIFastPolicySettings{
Rules: []OpenAIFastPolicyRule{{
ServiceTier: OpenAIFastTierPriority,
Action: BetaPolicyActionPass,
Scope: BetaPolicyScopeAll,
UserIDs: []int64{0},
}},
})
require.Error(t, err)
err = svc.SetOpenAIFastPolicySettings(context.Background(), &OpenAIFastPolicySettings{
Rules: []OpenAIFastPolicyRule{{
ServiceTier: OpenAIFastTierPriority,
Action: BetaPolicyActionPass,
Scope: BetaPolicyScopeAll,
UserIDs: []int64{42, 42},
}},
})
require.Error(t, err)
// Valid settings persisted
err = svc.SetOpenAIFastPolicySettings(context.Background(), &OpenAIFastPolicySettings{
Rules: []OpenAIFastPolicyRule{{
ServiceTier: OpenAIFastTierPriority,
Action: OpenAIFastPolicyActionForcePriority,
Scope: BetaPolicyScopeAll,
UserIDs: []int64{42, 43},
}},
})
require.NoError(t, err)
@@ -324,4 +409,5 @@ func TestSetOpenAIFastPolicySettings_Validation(t *testing.T) {
require.Len(t, got.Rules, 1)
require.Equal(t, OpenAIFastTierPriority, got.Rules[0].ServiceTier)
require.Equal(t, OpenAIFastPolicyActionForcePriority, got.Rules[0].Action)
require.Equal(t, []int64{42, 43}, got.Rules[0].UserIDs)
}
@@ -14,6 +14,7 @@ import (
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
coderws "github.com/coder/websocket"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
@@ -67,6 +68,39 @@ func TestWSResponseCreate_ExplicitFilterStripsServiceTier(t *testing.T) {
require.NotContains(t, string(updated), `"service_tier"`)
}
func TestWSResponseCreate_UserScopedRuleOverridesGlobalRule(t *testing.T) {
settings := &OpenAIFastPolicySettings{
Rules: []OpenAIFastPolicyRule{
{
ServiceTier: OpenAIFastTierPriority,
Action: BetaPolicyActionFilter,
Scope: BetaPolicyScopeAll,
},
{
ServiceTier: OpenAIFastTierPriority,
Action: BetaPolicyActionPass,
Scope: BetaPolicyScopeAll,
UserIDs: []int64{42},
},
},
}
svc := newOpenAIGatewayServiceWithSettings(t, settings)
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
frame := []byte(`{"type":"response.create","model":"gpt-5.5","service_tier":"priority"}`)
allowedUserCtx := context.WithValue(context.Background(), ctxkey.UserID, int64(42))
updated, blocked, err := svc.applyOpenAIFastPolicyToWSResponseCreate(allowedUserCtx, account, "gpt-5.5", frame)
require.NoError(t, err)
require.Nil(t, blocked)
require.Equal(t, "priority", gjson.GetBytes(updated, "service_tier").String())
otherUserCtx := context.WithValue(context.Background(), ctxkey.UserID, int64(43))
updated, blocked, err = svc.applyOpenAIFastPolicyToWSResponseCreate(otherUserCtx, account, "gpt-5.5", frame)
require.NoError(t, err)
require.Nil(t, blocked)
require.NotContains(t, string(updated), `"service_tier"`)
}
func TestWSResponseCreate_ForcePriorityRewritesKnownTier(t *testing.T) {
settings := &OpenAIFastPolicySettings{
Rules: []OpenAIFastPolicyRule{{
@@ -1015,7 +1049,7 @@ func TestPassthroughUsageMeta_TracksReasoningEffortAcrossTurns(t *testing.T) {
firstOut, firstBlocked, firstErr := svc.applyOpenAIFastPolicyToWSResponseCreate(context.Background(), account, capturedSessionModel, firstFrame)
require.NoError(t, firstErr)
require.Nil(t, firstBlocked)
meta.initFromFirstFrame(firstOut)
meta.initFromFirstFrame(firstOut, capturedSessionModel)
require.NotNil(t, meta.reasoningEffort.Load())
require.Equal(t, "medium", *meta.reasoningEffort.Load())
@@ -1032,7 +1066,7 @@ func TestPassthroughUsageMeta_TracksReasoningEffortAcrossTurns(t *testing.T) {
out, blocked, policyErr := svc.applyOpenAIFastPolicyToWSResponseCreate(context.Background(), account, model, payload)
if policyErr == nil && blocked == nil &&
strings.TrimSpace(gjson.GetBytes(payload, "type").String()) == "response.create" {
meta.updateFromResponseCreate(out, requestModelForThisFrame)
meta.updateFromResponseCreate(out, model, requestModelForThisFrame)
}
return out, blocked, policyErr
}
@@ -2,14 +2,12 @@ package service
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"strings"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
"github.com/Wei-Shaw/sub2api/internal/util/responseheaders"
@@ -72,13 +70,13 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions(
}
clientStream := gjson.GetBytes(body, "stream").Bool()
// 1b. Extract reasoning effort and service tier from the raw body before any transformation.
reasoningEffort := extractOpenAIReasoningEffortFromBody(body, originalModel)
// 1b. Extract service tier from the raw body before any transformation.
serviceTier := extractOpenAIServiceTierFromBody(body)
// 2. Resolve model mapping (same as ForwardAsChatCompletions)
billingModel := resolveOpenAIForwardModel(account, originalModel, defaultMappedModel)
upstreamModel := normalizeOpenAIModelForUpstream(account, billingModel)
reasoningEffort := extractOpenAIReasoningEffortFromBody(body, upstreamModel, billingModel, originalModel)
// 国产模型默认 effort 补充:需要 mappedModel 判定,推迟到 billingModel 算出之后。
reasoningEffort = ApplyThinkingEnabledFallback(reasoningEffort, body, billingModel)
@@ -383,12 +381,9 @@ func extractCCStreamUsage(payload string) *OpenAIUsage {
if !usageResult.Exists() || !usageResult.IsObject() {
return nil
}
u := OpenAIUsage{
InputTokens: int(gjson.Get(payload, "usage.prompt_tokens").Int()),
OutputTokens: int(gjson.Get(payload, "usage.completion_tokens").Int()),
}
if cached := gjson.Get(payload, "usage.prompt_tokens_details.cached_tokens"); cached.Exists() {
u.CacheReadInputTokens = int(cached.Int())
u, ok := openAIUsageFromGJSON(usageResult)
if !ok {
return nil
}
return &u
}
@@ -414,16 +409,9 @@ func (s *OpenAIGatewayService) bufferRawChatCompletions(
return nil, fmt.Errorf("read upstream body: %w", err)
}
var ccResp apicompat.ChatCompletionsResponse
var usage OpenAIUsage
if err := json.Unmarshal(respBody, &ccResp); err == nil && ccResp.Usage != nil {
usage = OpenAIUsage{
InputTokens: ccResp.Usage.PromptTokens,
OutputTokens: ccResp.Usage.CompletionTokens,
}
if ccResp.Usage.PromptTokensDetails != nil {
usage.CacheReadInputTokens = ccResp.Usage.PromptTokensDetails.CachedTokens
}
if parsedUsage, ok := extractOpenAIUsageFromJSONBytes(respBody); ok {
usage = parsedUsage
}
if s.responseHeaderFilter != nil {
@@ -122,6 +122,90 @@ func TestForwardAsRawChatCompletions_ForcesStreamUsageUpstreamAndPassesUsageDown
require.Contains(t, rec.Body.String(), "data: [DONE]")
}
func TestForwardAsRawChatCompletions_PreservesMappedGPT56MaxEffort(t *testing.T) {
gin.SetMode(gin.TestMode)
body := []byte(`{"model":"sol","messages":[{"role":"user","content":"hello"}],"reasoning_effort":"max","stream":false}`)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(
`{"id":"chatcmpl_max","object":"chat.completion","model":"gpt-5.6-sol","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5}}`,
)),
}}
svc := &OpenAIGatewayService{
cfg: rawChatCompletionsTestConfig(),
httpUpstream: upstream,
}
account := rawChatCompletionsTestAccount()
account.Credentials["model_mapping"] = map[string]any{"sol": "gpt-5.6-sol"}
result, err := svc.forwardAsRawChatCompletions(context.Background(), c, account, body, "")
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(upstream.lastBody, "model").String())
require.Equal(t, "max", gjson.GetBytes(upstream.lastBody, "reasoning_effort").String())
require.NotNil(t, result.ReasoningEffort)
require.Equal(t, "max", *result.ReasoningEffort)
}
func TestForwardAsRawChatCompletions_NonStreamingCapturesCacheWriteUsage(t *testing.T) {
gin.SetMode(gin.TestMode)
tests := []struct {
name string
usageJSON string
wantWrite int
}{
{
name: "positive cache write",
usageJSON: `{"prompt_tokens":12,"completion_tokens":3,"total_tokens":15,"prompt_tokens_details":{"cached_tokens":4,"cache_write_tokens":6}}`,
wantWrite: 6,
},
{
name: "nested zero overrides legacy alias",
usageJSON: `{"prompt_tokens":12,"completion_tokens":3,"total_tokens":15,"cache_creation_input_tokens":19,"prompt_tokens_details":{"cached_tokens":4,"cache_write_tokens":0}}`,
wantWrite: 0,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
body := []byte(`{"model":"gpt-5.6","messages":[{"role":"user","content":"hello"}],"stream":false}`)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(
`{"id":"chatcmpl_cache","object":"chat.completion","model":"gpt-5.6","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"usage":` + tt.usageJSON + `}`,
)),
}}
svc := &OpenAIGatewayService{
cfg: rawChatCompletionsTestConfig(),
httpUpstream: upstream,
}
result, err := svc.forwardAsRawChatCompletions(context.Background(), c, rawChatCompletionsTestAccount(), body, "")
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, 12, result.Usage.InputTokens)
require.Equal(t, 4, result.Usage.CacheReadInputTokens)
require.Equal(t, tt.wantWrite, result.Usage.CacheCreationInputTokens)
})
}
}
func TestForwardAsRawChatCompletions_PreservesDeepSeekReasoningContentNonStreaming(t *testing.T) {
gin.SetMode(gin.TestMode)
@@ -35,6 +35,14 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
return nil, errors.New("codex_cli_only restriction: only codex official clients are allowed")
}
normalizedBody, normalized, err := normalizeOpenAICodexCompactReasoningEffortForAccount(c, account, body)
if err != nil {
return nil, err
}
if normalized {
body = normalizedBody
}
originalBody := body
requestView := newOpenAIRequestView(body)
reqModel, reqStream, promptCacheKey := requestView.Model, requestView.Stream, requestView.PromptCacheKey
@@ -53,6 +61,10 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
setOpenAICompatMessagesBridgeContext(c, compatMessagesBridge)
isCodexCLI := openai.IsCodexOfficialClientByHeaders(c.GetHeader("User-Agent"), c.GetHeader("originator")) || (s.cfg != nil && s.cfg.Gateway.ForceCodexCLI)
codexImageGenerationExplicitToolPolicy := codexImageGenerationExplicitToolPolicyAllow
if isCodexCLI {
codexImageGenerationExplicitToolPolicy = account.CodexImageGenerationExplicitToolPolicy()
}
wsDecision := s.getOpenAIWSProtocolResolver().Resolve(account)
clientTransport := GetOpenAIClientTransport(c)
// 仅允许 WS 入站请求走 WS 上游,避免出现 HTTP -> WS 协议混用。
@@ -87,10 +99,22 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
}
passthroughEnabled := account.IsOpenAIPassthroughEnabled()
if passthroughEnabled {
if isCodexCLI && codexImageGenerationExplicitToolPolicy == codexImageGenerationExplicitToolPolicyStrip {
strippedBody, changed, stripErr := stripOpenAIImageGenerationToolsFromRawPayload(body)
if stripErr != nil {
return nil, stripErr
}
if changed {
body = strippedBody
originalBody = strippedBody
logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Stripped /responses image_generation tool for Codex client by account policy")
}
}
// 透传分支只需要轻量提取字段,避免热路径全量 Unmarshal。
reasoningEffort := extractOpenAIReasoningEffortFromBody(body, reqModel)
mappedModel := account.GetMappedModel(reqModel)
reasoningEffort := extractOpenAIReasoningEffortFromBody(body, mappedModel)
// 国产模型默认 effort 补充:也要用 mappedModel 判定是否是 passback-required 上游。
reasoningEffort = ApplyThinkingEnabledFallback(reasoningEffort, body, account.GetMappedModel(reqModel))
reasoningEffort = ApplyThinkingEnabledFallback(reasoningEffort, body, mappedModel)
return s.forwardOpenAIPassthrough(ctx, c, account, originalBody, reqModel, reasoningEffort, reqStream, startTime)
}
@@ -150,10 +174,6 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
if apiKey != nil {
imageGenerationAllowed = GroupAllowsImageGeneration(apiKey.Group)
}
codexImageGenerationExplicitToolPolicy := codexImageGenerationExplicitToolPolicyAllow
if isCodexCLI {
codexImageGenerationExplicitToolPolicy = account.CodexImageGenerationExplicitToolPolicy()
}
codexImageGenerationBridgeEnabled := isCodexCLI && imageGenerationAllowed && codexImageGenerationExplicitToolPolicy != codexImageGenerationExplicitToolPolicyStrip && s.isCodexImageGenerationBridgeEnabled(ctx, account, apiKey)
var imageIntent bool
if isCodexCLI && codexImageGenerationExplicitToolPolicy == codexImageGenerationExplicitToolPolicyStrip {
@@ -263,7 +283,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
markDecodedModified()
logger.LegacyPrintf("service.openai_gateway", "[OpenAI] Added Codex image_generation bridge instructions")
}
} else if imageGenerationAllowed && imageIntent && openAIRequestBodyHasImageGenerationTool(body) {
} else if imageGenerationAllowed && imageIntent && openAIRequestBodyHasImageGenerationDeclaration(body) {
// 完整 image_generation tool 只做 raw 计费读取,校验/桥接/旧字段迁移命中时才展开大 input map。
logger.LegacyPrintf("service.openai_gateway", "[OpenAI] /responses image_generation request inbound_model=%s mapped_model=%s account_type=%s", requestView.Model, upstreamModel, account.Type)
}
@@ -283,7 +303,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
// gpt-5.3-codex-spark also rejects the image_generation tool (HTTP 400,
// param=tools). Strip it here so both APIKey and OAuth /responses paths are
// covered regardless of the image-generation feature gate.
if isCodexSparkModel(upstreamModel) && openAIRequestBodyHasImageGenerationTool(body) {
if isCodexSparkModel(upstreamModel) && openAIRequestBodyHasImageGenerationDeclaration(body) {
decoded, decodeErr := ensureReqBody()
if decodeErr != nil {
return nil, decodeErr
@@ -746,7 +766,7 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
}
defer func() { _ = resp.Body.Close() }()
reasoningEffort := extractOpenAIReasoningEffortFromBody(body, originalModel)
reasoningEffort := extractOpenAIReasoningEffortFromBody(body, upstreamModel, billingModel, originalModel)
// 国产模型默认 effort 补充:此处 reqModel 已被 mapping 重写为 billingModel(见
// line 2510-2515 的 GetMappedModel + reqModel 赋值),可直接作为 mappedModel。
reasoningEffort = ApplyThinkingEnabledFallback(reasoningEffort, body, reqModel)
@@ -902,6 +922,10 @@ func (s *OpenAIGatewayService) buildUpstreamRequest(ctx context.Context, c *gin.
req.Header.Set("conversation_id", isolated)
}
}
} else if isOpenAIResponsesCompactPath(c) {
// compact 上游是 unary JSON 协议:API-key 账号也显式声明 Accept,
// 避免 OpenAI 兼容网关按 SSE 返回(#3777 期望行为 4)。
req.Header.Set("accept", "application/json")
}
// Apply custom User-Agent if configured
@@ -920,6 +944,11 @@ func (s *OpenAIGatewayService) buildUpstreamRequest(ctx context.Context, c *gin.
// (Chrome/Firefox/Safari/Edge 等),替换为后台配置的 Codex UA,避免 Cloudflare 触发 JS 质询。
s.overrideBrowserUserAgent(ctx, account, req)
// 终态收口:originator 必须与最终 User-Agent 首段配套且为官方身份,否则上游 404(issue #3901)。
if account.Type == AccountTypeOAuth {
enforceCodexIdentityHeaders(req.Header)
}
// Ensure required headers exist
if req.Header.Get("content-type") == "" {
req.Header.Set("content-type", "application/json")
@@ -122,13 +122,14 @@ func (s *OpenAIGatewayService) forwardGrokResponses(
if usage == nil {
usage = &OpenAIUsage{}
}
reasoningEffort := extractOpenAIReasoningEffortFromBody(patchedBody, originalModel)
return &OpenAIForwardResult{
RequestID: firstNonEmpty(resp.Header.Get("x-request-id"), resp.Header.Get("xai-request-id")),
ResponseID: responseID,
Usage: *usage,
Model: originalModel,
UpstreamModel: upstreamModel,
ReasoningEffort: ptrStringOrNil(normalizeOpenAIReasoningEffort(gjson.GetBytes(patchedBody, "reasoning.effort").String())),
ReasoningEffort: reasoningEffort,
Stream: reqStream,
OpenAIWSMode: false,
ResponseHeaders: resp.Header.Clone(),
@@ -693,10 +694,3 @@ func (s *OpenAIGatewayService) tempUnscheduleGrok(ctx context.Context, account *
_ = s.accountRepo.SetTempUnschedulable(stateCtx, account.ID, until, reason)
}
}
func ptrStringOrNil(value string) *string {
if strings.TrimSpace(value) == "" {
return nil
}
return &value
}
@@ -41,6 +41,17 @@ func TestPatchGrokResponsesBodySetsMappedModelAndDropsUnsupportedFields(t *testi
require.Equal(t, "high", gjson.GetBytes(patched, "reasoning.effort").String())
}
func TestExtractGrokResponsesReasoningEffortSupportsOpenAICompatibleField(t *testing.T) {
t.Parallel()
effort := extractOpenAIReasoningEffortFromBody(
[]byte(`{"model":"grok-4.3","reasoning_effort":"high"}`),
"grok-4.3",
)
require.NotNil(t, effort)
require.Equal(t, "high", *effort)
}
func TestPatchGrokResponsesBodyDropsGrok45ReasoningUnsupportedFields(t *testing.T) {
t.Parallel()
@@ -656,7 +667,7 @@ func TestForwardGrokResponsesStreamingUsesXAIResponsesAndSnapshots(t *testing.T)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
body := []byte(`{"model":"grok","input":"hi","stream":true}`)
body := []byte(`{"model":"grok","input":"hi","stream":true,"reasoning_effort":"high"}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
c.Request.Header.Set("OpenAI-Beta", "responses=experimental")
@@ -708,6 +719,7 @@ func TestForwardGrokResponsesStreamingUsesXAIResponsesAndSnapshots(t *testing.T)
require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization"))
require.Equal(t, "responses=experimental", upstream.lastReq.Header.Get("OpenAI-Beta"))
require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.lastBody, "model").String())
require.Equal(t, "high", gjson.GetBytes(upstream.lastBody, "reasoning_effort").String())
require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool())
require.True(t, result.Stream)
require.Equal(t, "resp_grok", result.ResponseID)
@@ -715,6 +727,8 @@ func TestForwardGrokResponsesStreamingUsesXAIResponsesAndSnapshots(t *testing.T)
require.Equal(t, 5, result.Usage.InputTokens)
require.Equal(t, 3, result.Usage.OutputTokens)
require.Equal(t, 2, result.Usage.CacheReadInputTokens)
require.NotNil(t, result.ReasoningEffort)
require.Equal(t, "high", *result.ReasoningEffort)
require.Contains(t, recorder.Header().Get("Content-Type"), "text/event-stream")
require.Contains(t, recorder.Body.String(), "response.output_text.delta")
require.NotNil(t, repo.updates[52][grokQuotaSnapshotExtraKey])
@@ -259,6 +259,13 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic(
}
// 6. Build upstream request
if account.Type == AccountTypeOAuth && account.Platform != PlatformGrok {
// Messages 兼容桥即使 body 未带 todo-guard/prompt_cache_key 标记(如映射到非
// gpt-5/codex 模型),也必须让 buildUpstreamRequest 走 bridge 分支:不带
// originator、User-Agent 逐字透传,避免身份收口(issue #3901)误改本路径
// 刻意最小化的请求形态(下方的 Del(OpenAI-Beta/originator) 兜底保持不变)。
setOpenAICompatMessagesBridgeContext(c, true)
}
upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
var upstreamReq *http.Request
if account.Platform == PlatformGrok {
@@ -1114,8 +1121,9 @@ func copyOpenAIUsageFromResponsesUsage(usage *apicompat.ResponsesUsage) OpenAIUs
return OpenAIUsage{}
}
result := OpenAIUsage{
InputTokens: usage.InputTokens,
OutputTokens: usage.OutputTokens,
InputTokens: usage.InputTokens,
OutputTokens: usage.OutputTokens,
CacheCreationInputTokens: usage.CacheCreationInputTokens,
}
if usage.InputTokensDetails != nil {
result.CacheReadInputTokens = usage.InputTokensDetails.CachedTokens
@@ -72,7 +72,7 @@ func (s *OpenAIGatewayService) forwardAnthropicViaRawChatCompletions(
chatReq.StreamOptions = &apicompat.ChatStreamOptions{IncludeUsage: true}
}
reasoningEffort := extractOpenAIReasoningEffortFromBody(body, originalModel)
reasoningEffort := extractOpenAIReasoningEffortFromBody(body, upstreamModel, billingModel, originalModel)
reasoningEffort = ApplyThinkingEnabledFallback(reasoningEffort, body, billingModel)
serviceTier := extractOpenAIServiceTierFromBody(body)
@@ -0,0 +1,28 @@
//go:build unit
package service
import (
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
"github.com/stretchr/testify/require"
)
func TestCopyOpenAIUsageFromResponsesUsageTrustsCanonicalCacheCreationValue(t *testing.T) {
usage := &apicompat.ResponsesUsage{
InputTokens: 20,
OutputTokens: 2,
CacheCreationInputTokens: 0,
InputTokensDetails: &apicompat.ResponsesInputTokensDetails{
CachedTokens: 3,
CacheWriteTokens: 19,
},
}
got := copyOpenAIUsageFromResponsesUsage(usage)
require.Equal(t, 20, got.InputTokens)
require.Equal(t, 3, got.CacheReadInputTokens)
require.Zero(t, got.CacheCreationInputTokens)
}
@@ -17,7 +17,6 @@ import (
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
"github.com/Wei-Shaw/sub2api/internal/util/responseheaders"
"github.com/gin-gonic/gin"
"github.com/tidwall/gjson"
@@ -382,6 +381,11 @@ func (s *OpenAIGatewayService) buildUpstreamRequestOpenAIPassthrough(
if clientConversationID != "" {
req.Header.Set("conversation_id", isolateOpenAISessionID(apiKeyID, clientConversationID))
}
} else if isOpenAIResponsesCompactPath(c) {
// 透传白名单会放行客户端的 Accept: text/event-stream;compact 上游是
// unary JSON 协议,API-key 账号同样强制 Accept,避免上游按 SSE 返回
// (#3777 期望行为 4)。
req.Header.Set("accept", "application/json")
}
// 透传模式也支持账户自定义 User-Agent 与 ForceCodexCLI 兜底。
@@ -392,15 +396,17 @@ func (s *OpenAIGatewayService) buildUpstreamRequestOpenAIPassthrough(
if s.cfg != nil && s.cfg.Gateway.ForceCodexCLI {
req.Header.Set("user-agent", codexCLIUserAgent)
}
// OAuth 安全透传:对非 Codex UA 统一兜底,降低被上游风控拦截概率。
if account.Type == AccountTypeOAuth && !openai.IsCodexCLIRequest(req.Header.Get("user-agent")) {
req.Header.Set("user-agent", codexCLIUserAgent)
}
// 浏览器型 UA 兜底:仅 OAuth(ChatGPT 内部接口)账号生效,若最终 user-agent 仍为浏览器
// (Chrome/Firefox/Safari/Edge 等),替换为后台配置的 Codex UA,避免 Cloudflare 触发 JS 质询。
s.overrideBrowserUserAgent(ctx, account, req)
// 终态收口:originator 必须与最终 User-Agent 首段配套且为官方身份,非官方 UA 整体回退为
// 默认 Codex CLI 身份(承接原「非 Codex UA 安全兜底」,并修复其把 codex-tui 等官方 UA 改写为
// codex_cli_rs 造成的 originator 错配 404),详见 issue #3901。
if account.Type == AccountTypeOAuth {
enforceCodexIdentityHeaders(req.Header)
}
if req.Header.Get("content-type") == "" {
req.Header.Set("content-type", "application/json")
}
@@ -1122,6 +1128,7 @@ func (s *OpenAIGatewayService) handlePassthroughSSEToJSON(resp *http.Response, c
}
}
}
finalResponse = supplementCompactionItemFromSSE(c, finalResponse, bodyText)
body = finalResponse
if originalModel != "" && mappedModel != "" && originalModel != mappedModel {
body = s.replaceModelInResponseBody(body, mappedModel, originalModel)
@@ -247,7 +247,7 @@ func expectedOpenAICost(t *testing.T, svc *OpenAIGatewayService, model string, u
t.Helper()
cost, err := svc.billingService.CalculateCost(model, UsageTokens{
InputTokens: max(usage.InputTokens-usage.CacheReadInputTokens, 0),
InputTokens: max(usage.InputTokens-usage.CacheReadInputTokens-usage.CacheCreationInputTokens, 0),
OutputTokens: usage.OutputTokens,
CacheCreationTokens: usage.CacheCreationInputTokens,
CacheReadTokens: usage.CacheReadInputTokens,
@@ -1002,6 +1002,49 @@ func TestOpenAIGatewayServiceRecordUsage_ClampsActualInputTokensToZero(t *testin
require.Equal(t, 0, usageRepo.lastLog.InputTokens)
}
func TestOpenAIGatewayServiceRecordUsage_GPT56SeparatesCacheWriteForBillingAndStats(t *testing.T) {
usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
userRepo := &openAIRecordUsageUserRepoStub{}
subRepo := &openAIRecordUsageSubRepoStub{}
svc := newOpenAIRecordUsageServiceForTest(usageRepo, userRepo, subRepo, nil)
svc.billingService = NewBillingService(svc.cfg, &PricingService{pricingData: map[string]*LiteLLMModelPricing{
"gpt-5.6-sol": {
InputCostPerToken: 5e-6,
OutputCostPerToken: 30e-6,
CacheReadInputTokenCost: 0.5e-6,
},
}})
err := svc.RecordUsage(context.Background(), &OpenAIRecordUsageInput{
Result: &OpenAIForwardResult{
RequestID: "resp_gpt56_cache_write",
Usage: OpenAIUsage{
InputTokens: 1000,
OutputTokens: 50,
CacheCreationInputTokens: 200,
CacheReadInputTokens: 100,
},
Model: "gpt-5.6-sol",
Duration: time.Second,
},
APIKey: &APIKey{ID: 1056},
User: &User{ID: 2056},
Account: &Account{ID: 3056},
})
require.NoError(t, err)
require.NotNil(t, usageRepo.lastLog)
require.Equal(t, 700, usageRepo.lastLog.InputTokens)
require.Equal(t, 200, usageRepo.lastLog.CacheCreationTokens)
require.Equal(t, 100, usageRepo.lastLog.CacheReadTokens)
require.Equal(t, 1050, usageRepo.lastLog.TotalTokens())
require.InDelta(t, 700*5e-6, usageRepo.lastLog.InputCost, 1e-12)
require.InDelta(t, 200*6.25e-6, usageRepo.lastLog.CacheCreationCost, 1e-12)
require.InDelta(t, 100*0.5e-6, usageRepo.lastLog.CacheReadCost, 1e-12)
require.InDelta(t, 50*30e-6, usageRepo.lastLog.OutputCost, 1e-12)
require.InDelta(t, usageRepo.lastLog.TotalCost*1.1, usageRepo.lastLog.ActualCost, 1e-12)
}
func TestOpenAIGatewayServiceRecordUsage_Gpt54LongContextBillsWholeSession(t *testing.T) {
usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
userRepo := &openAIRecordUsageUserRepoStub{}
@@ -9,6 +9,7 @@ import (
"net/http"
"strings"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
"github.com/gin-gonic/gin"
"github.com/google/uuid"
@@ -199,6 +200,32 @@ func normalizeOpenAICompactRequestBody(body []byte) ([]byte, bool, error) {
return normalized, true, nil
}
func normalizeOpenAICodexCompactReasoningEffortForAccount(c *gin.Context, account *Account, body []byte) ([]byte, bool, error) {
if account == nil || !account.IsOpenAIOAuth() || !isOpenAIResponsesCompactPath(c) {
return body, false, nil
}
requestedModel := strings.TrimSpace(gjson.GetBytes(body, "model").String())
effectiveModel := account.GetMappedModel(requestedModel)
return normalizeOpenAICodexCompactReasoningEffort(body, effectiveModel)
}
func normalizeOpenAICodexCompactReasoningEffort(body []byte, effectiveModel string) ([]byte, bool, error) {
if !isOpenAIGPT56Model(effectiveModel) ||
!strings.EqualFold(strings.TrimSpace(gjson.GetBytes(body, "reasoning.effort").String()), "max") {
return body, false, nil
}
// Codex Ultra 在客户端编排层会下发 max;ChatGPT compact 端点目前只接受到
// xhigh。这里只降级 OpenAI OAuth 的 GPT-5.6 compact 子请求,普通 Responses、
// API Key 请求和其他平台的 OAuth 请求保留 max。
normalized, err := sjson.SetBytes(body, "reasoning.effort", "xhigh")
if err != nil {
return body, false, fmt.Errorf("normalize codex compact reasoning effort: %w", err)
}
return normalized, true, nil
}
func resolveOpenAICompactSessionID(c *gin.Context) string {
if c != nil {
if sessionID := strings.TrimSpace(c.GetHeader("session_id")); sessionID != "" {
@@ -259,7 +286,7 @@ func (s *OpenAIGatewayService) replaceModelInResponseBody(body []byte, fromModel
return body
}
func getOpenAIReasoningEffortFromReqBody(reqBody map[string]any) (value string, present bool) {
func getOpenAIReasoningEffortFromReqBody(reqBody map[string]any, requestedModel string) (value string, present bool) {
if reqBody == nil {
return "", false
}
@@ -267,13 +294,13 @@ func getOpenAIReasoningEffortFromReqBody(reqBody map[string]any) (value string,
// Primary: reasoning.effort
if reasoning, ok := reqBody["reasoning"].(map[string]any); ok {
if effort, ok := reasoning["effort"].(string); ok {
return normalizeOpenAIReasoningEffort(effort), true
return normalizeOpenAIReasoningEffortForModel(effort, requestedModel), true
}
}
// Fallback: some clients may use a flat field.
if effort, ok := reqBody["reasoning_effort"].(string); ok {
return normalizeOpenAIReasoningEffort(effort), true
return normalizeOpenAIReasoningEffortForModel(effort, requestedModel), true
}
return "", false
@@ -302,7 +329,18 @@ func deriveOpenAIReasoningEffortFromModel(model string) string {
return ""
}
return normalizeOpenAIReasoningEffort(parts[len(parts)-1])
return normalizeOpenAIReasoningEffortForModel(parts[len(parts)-1], modelID)
}
// deriveOpenAIReasoningEffortFromModelCandidates 依次对每个候选模型做后缀推导,
// 返回第一个非空结果。
func deriveOpenAIReasoningEffortFromModelCandidates(models []string) string {
for _, model := range models {
if value := deriveOpenAIReasoningEffortFromModel(model); value != "" {
return value
}
}
return ""
}
type openAIRequestView struct {
@@ -545,20 +583,24 @@ func detectOpenAIPassthroughInstructionsRejectReason(reqModel string, body []byt
return ""
}
func extractOpenAIReasoningEffortFromBody(body []byte, requestedModel string) *string {
// extractOpenAIReasoningEffortFromBody 按优先级传入模型候选(如 upstreamModel,
// billingModel, originalModel):显式 effort 的模型归一化(max 保留判定)用第一个
// 非空候选;body 未携带 effort 时的模型后缀推导依次尝试每个候选——OAuth 的
// normalizeCodexModel 会剥掉 upstreamModel 的 effort 后缀,只有原始模型名还留着。
func extractOpenAIReasoningEffortFromBody(body []byte, modelCandidates ...string) *string {
reasoningEffort := strings.TrimSpace(gjson.GetBytes(body, "reasoning.effort").String())
if reasoningEffort == "" {
reasoningEffort = strings.TrimSpace(gjson.GetBytes(body, "reasoning_effort").String())
}
if reasoningEffort != "" {
normalized := normalizeOpenAIReasoningEffort(reasoningEffort)
normalized := normalizeOpenAIReasoningEffortForModel(reasoningEffort, firstNonEmpty(modelCandidates...))
if normalized == "" {
return nil
}
return &normalized
}
value := deriveOpenAIReasoningEffortFromModel(requestedModel)
value := deriveOpenAIReasoningEffortFromModelCandidates(modelCandidates)
if value == "" {
return nil
}
@@ -618,9 +660,12 @@ func (e *OpenAIFastBlockedError) Error() string { return e.Message }
//
// Matching rules:
// - Scope filters by account type (all / oauth / apikey / bedrock)
// - UserIDs, when present, filters by the trusted Sub2API user that owns the API key
// - ServiceTier must be empty (= any), "all", or equal the normalized tier
// - ModelWhitelist narrows the rule to specific models; FallbackAction
// handles the non-matching case (default: pass)
// - User-specific rules take precedence over global rules; each group keeps
// the configured first-match order
//
// 与 Claude BetaPolicy 的差异(保留首条匹配 short-circuit):
// - BetaPolicy 处理的是 anthropic-beta header 中的 token 集合,不同
@@ -646,39 +691,70 @@ func (s *OpenAIGatewayService) evaluateOpenAIFastPolicy(ctx context.Context, acc
}
settings = fetched
}
return evaluateOpenAIFastPolicyWithSettings(settings, account, model, tier)
return evaluateOpenAIFastPolicyWithSettings(settings, openAIFastPolicyUserID(ctx), account, model, tier)
}
// evaluateOpenAIFastPolicyWithSettings is the pure-function core extracted so
// long-lived sessions (e.g. WS) can prefetch settings once and avoid hitting
// the settingService on every frame. See WSSession entry and
// openAIFastPolicySettingsFromContext for the caching glue.
func evaluateOpenAIFastPolicyWithSettings(settings *OpenAIFastPolicySettings, account *Account, model, tier string) (action, errMsg string) {
func evaluateOpenAIFastPolicyWithSettings(settings *OpenAIFastPolicySettings, userID int64, account *Account, model, tier string) (action, errMsg string) {
if settings == nil {
return BetaPolicyActionPass, ""
}
isOAuth := account != nil && account.IsOAuth()
isBedrock := account != nil && account.IsBedrock()
for _, rule := range settings.Rules {
if !betaPolicyScopeMatches(rule.Scope, isOAuth, isBedrock) {
continue
// 用户专属规则先于全局规则。规则组内仍按配置顺序首条命中,允许
// 管理员为某位用户配置例外,而不被先出现的全局规则覆盖。
for _, userScoped := range []bool{true, false} {
for _, rule := range settings.Rules {
if (len(rule.UserIDs) > 0) != userScoped || !openAIFastPolicyUserMatches(rule.UserIDs, userID) {
continue
}
if !betaPolicyScopeMatches(rule.Scope, isOAuth, isBedrock) {
continue
}
ruleTier := strings.ToLower(strings.TrimSpace(rule.ServiceTier))
if ruleTier != "" && ruleTier != OpenAIFastTierAny && ruleTier != tier {
continue
}
eff := BetaPolicyRule{
Action: rule.Action,
ErrorMessage: rule.ErrorMessage,
ModelWhitelist: rule.ModelWhitelist,
FallbackAction: rule.FallbackAction,
FallbackErrorMessage: rule.FallbackErrorMessage,
}
return resolveRuleAction(eff, model)
}
ruleTier := strings.ToLower(strings.TrimSpace(rule.ServiceTier))
if ruleTier != "" && ruleTier != OpenAIFastTierAny && ruleTier != tier {
continue
}
eff := BetaPolicyRule{
Action: rule.Action,
ErrorMessage: rule.ErrorMessage,
ModelWhitelist: rule.ModelWhitelist,
FallbackAction: rule.FallbackAction,
FallbackErrorMessage: rule.FallbackErrorMessage,
}
return resolveRuleAction(eff, model)
}
return BetaPolicyActionPass, ""
}
func openAIFastPolicyUserID(ctx context.Context) int64 {
if ctx == nil {
return 0
}
userID, _ := ctx.Value(ctxkey.UserID).(int64)
if userID <= 0 {
return 0
}
return userID
}
func openAIFastPolicyUserMatches(ruleUserIDs []int64, userID int64) bool {
if len(ruleUserIDs) == 0 {
return true
}
for _, ruleUserID := range ruleUserIDs {
if ruleUserID == userID {
return true
}
}
return false
}
// openAIFastPolicyCtxKey 是 context 中预取的 OpenAIFastPolicySettings 缓存
// 键,仅用于 WebSocket 长会话内多帧复用同一份策略快照,避免每帧 DB 命中。
//
@@ -772,6 +848,13 @@ func writeOpenAIFastPolicyBlockedResponse(c *gin.Context, err *OpenAIFastBlocked
return
}
MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalPolicyDenied)
// body-signal compact 心跳可能已把响应头提交为 200(长排队后才进入
// Forward),此时以 response.failed 终止事件回传;未提交时先停拍再写
// JSON,保持原状态码语义(#3887)。
if StopOpenAICompactSSEKeepaliveCommitted(c) {
writeOpenAICompactSSEFailureMessage(c, http.StatusForbidden, "permission_error", err.Message)
return
}
c.JSON(http.StatusForbidden, gin.H{
"error": gin.H{
"type": "permission_error",
@@ -1126,15 +1209,16 @@ func getOpenAIRequestBodyMap(_ *gin.Context, body []byte) (map[string]any, error
return reqBody, nil
}
func extractOpenAIReasoningEffort(reqBody map[string]any, requestedModel string) *string {
if value, present := getOpenAIReasoningEffortFromReqBody(reqBody); present {
// extractOpenAIReasoningEffort 的模型候选语义同 extractOpenAIReasoningEffortFromBody。
func extractOpenAIReasoningEffort(reqBody map[string]any, modelCandidates ...string) *string {
if value, present := getOpenAIReasoningEffortFromReqBody(reqBody, firstNonEmpty(modelCandidates...)); present {
if value == "" {
return nil
}
return &value
}
value := deriveOpenAIReasoningEffortFromModel(requestedModel)
value := deriveOpenAIReasoningEffortFromModelCandidates(modelCandidates)
if value == "" {
return nil
}
@@ -1162,3 +1246,10 @@ func normalizeOpenAIReasoningEffort(raw string) string {
return ""
}
}
func normalizeOpenAIReasoningEffortForModel(raw, model string) string {
if strings.EqualFold(strings.TrimSpace(raw), "max") && isOpenAIGPT56Model(model) {
return "max"
}
return normalizeOpenAIReasoningEffort(raw)
}
@@ -756,10 +756,8 @@ func openAIUsageFromGJSON(value gjson.Result) (OpenAIUsage, bool) {
if outputTokens == 0 {
outputTokens = value.Get("completion_tokens").Int()
}
cacheReadTokens := value.Get("input_tokens_details.cached_tokens").Int()
if cacheReadTokens == 0 {
cacheReadTokens = value.Get("prompt_tokens_details.cached_tokens").Int()
}
cacheReadTokens := openAICacheReadTokensFromUsage(value)
cacheCreationTokens := openAICacheCreationTokensFromUsage(value)
imageOutputTokens := value.Get("output_tokens_details.image_tokens").Int()
if imageOutputTokens == 0 {
imageOutputTokens = value.Get("completion_tokens_details.image_tokens").Int()
@@ -767,12 +765,49 @@ func openAIUsageFromGJSON(value gjson.Result) (OpenAIUsage, bool) {
return OpenAIUsage{
InputTokens: int(inputTokens),
OutputTokens: int(outputTokens),
CacheCreationInputTokens: int(value.Get("cache_creation_input_tokens").Int()),
CacheReadInputTokens: int(cacheReadTokens),
CacheCreationInputTokens: cacheCreationTokens,
CacheReadInputTokens: cacheReadTokens,
ImageOutputTokens: int(imageOutputTokens),
}, true
}
func openAICacheReadTokensFromUsage(value gjson.Result) int {
for _, nested := range []gjson.Result{
value.Get("input_tokens_details.cached_tokens"),
value.Get("prompt_tokens_details.cached_tokens"),
} {
if nested.Exists() {
return max(int(nested.Int()), 0)
}
}
return firstPositiveGJSONInt(
value.Get("cache_read_input_tokens"),
value.Get("cache_read_tokens"),
value.Get("cached_tokens"),
)
}
func openAICacheCreationTokensFromUsage(value gjson.Result) int {
for _, nested := range []gjson.Result{
value.Get("input_tokens_details.cache_write_tokens"),
value.Get("prompt_tokens_details.cache_write_tokens"),
value.Get("input_tokens_details.cache_creation_tokens"),
value.Get("prompt_tokens_details.cache_creation_tokens"),
} {
if nested.Exists() {
return max(int(nested.Int()), 0)
}
}
return firstPositiveGJSONInt(
value.Get("cache_write_tokens"),
value.Get("cache_creation_input_tokens"),
value.Get("cache_write_input_tokens"),
value.Get("cache_creation_tokens"),
)
}
func (s *OpenAIGatewayService) handleNonStreamingResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, originalModel, mappedModel string) (*openaiNonStreamingResult, error) {
body, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, openAITooLargeError)
if err != nil {
@@ -878,6 +913,7 @@ func (s *OpenAIGatewayService) handleSSEToJSON(resp *http.Response, c *gin.Conte
}
}
}
finalResponse = supplementCompactionItemFromSSE(c, finalResponse, bodyText)
body = finalResponse
if originalModel != mappedModel {
body = s.replaceModelInResponseBody(body, mappedModel, originalModel)
@@ -1012,6 +1048,12 @@ func (s *OpenAIGatewayService) writeOpenAINonStreamingProtocolError(resp *http.R
message = "Upstream returned an invalid non-streaming response"
}
setOpsUpstreamError(c, http.StatusBadGateway, message, "")
// body-signal compact 心跳可能已把响应头提交为 200,此时只能以
// response.failed 终止事件回传错误,不能再写 JSON+状态码。
if openAICompactClientWantsStream(c) && StopOpenAICompactSSEKeepaliveCommitted(c) {
writeOpenAICompactSSEFailureMessage(c, http.StatusBadGateway, "upstream_error", message)
return fmt.Errorf("non-streaming openai protocol error: %s", message)
}
responseheaders.WriteFilteredHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter)
c.Writer.Header().Set("Content-Type", "application/json; charset=utf-8")
c.JSON(http.StatusBadGateway, gin.H{
@@ -1081,10 +1123,152 @@ func responsesStreamEventMayContributeToOutput(eventType string) bool {
}
}
// reconstructResponseOutputFromSSE scans raw SSE body text for delta events and
// returns a JSON-encoded output array reconstructed from accumulated deltas.
// Returns (nil, false) if no content was found in deltas.
// collectRawResponsesOutputItemsFromSSE 按到达顺序收集 SSE 流中
// response.output_item.done 携带的原始 item。item 以 raw JSON 逐字节保留,
// 避免经窄结构体重建时丢弃 encrypted_content/summary/opaque 等 compact
// 专属或未来新增字段(#3777 问题 2)。若整条流没有任何 done 事件,退回
// 收集 output_item.added 中的 compaction 类 item——compaction 结果没有
// delta 事件,部分上游只在 added 事件中携带完整 item。
func collectRawResponsesOutputItemsFromSSE(bodyText string) ([]byte, bool) {
var items []json.RawMessage
seen := make(map[string]struct{})
hasCompactionItem := false
appendItem := func(item gjson.Result) {
if !item.Exists() || !item.IsObject() {
return
}
key := strings.TrimSpace(item.Get("id").String())
if key == "" {
key = item.Raw
}
if _, dup := seen[key]; dup {
return
}
seen[key] = struct{}{}
if isResponsesCompactionItemType(item.Get("type").String()) {
hasCompactionItem = true
}
items = append(items, json.RawMessage(item.Raw))
}
forEachOpenAISSEDataPayload(bodyText, func(data []byte) {
if strings.TrimSpace(gjson.GetBytes(data, "type").String()) != "response.output_item.done" {
return
}
appendItem(gjson.GetBytes(data, "item"))
})
// done 事件未携带 compaction item 时再看 added:覆盖"其他 item 有 done、
// compaction 只在 added 中"的混合形态;done 已含 compaction 时跳过,
// 避免同一 item 在无 id 可去重时被收集两份(Codex 要求恰好一个)。
if !hasCompactionItem {
forEachOpenAISSEDataPayload(bodyText, func(data []byte) {
if strings.TrimSpace(gjson.GetBytes(data, "type").String()) != "response.output_item.added" {
return
}
item := gjson.GetBytes(data, "item")
if !isResponsesCompactionItemType(item.Get("type").String()) {
return
}
appendItem(item)
})
}
if len(items) == 0 {
return nil, false
}
outputJSON, err := json.Marshal(items)
if err != nil {
return nil, false
}
return outputJSON, true
}
// isResponsesCompactionItemType reports whether the item type is the Codex
// remote-compact result item ("compaction", upstream alias "compaction_summary").
func isResponsesCompactionItemType(itemType string) bool {
switch strings.TrimSpace(itemType) {
case "compaction", "compaction_summary":
return true
default:
return false
}
}
// supplementCompactionItemFromSSE 保证 compact 请求的终态 output 携带
// compaction item:终态 output 非空但缺失 compaction、而原始事件流的
// output_item.done(或 added)中存在时(上游不一致形态),以 raw JSON 补入。
// Codex remote compact v2 只从 output_item.done 收集 item 且要求恰好一个
// compaction item——纯流式透传(v0.1.146)下客户端直接读事件流天然拿得到,
// SSE→JSON 提取链路必须给出等价结果。非 compact 请求原样返回。
func supplementCompactionItemFromSSE(c *gin.Context, finalResponse []byte, bodyText string) []byte {
if !isOpenAIResponsesCompactPath(c) {
return finalResponse
}
if len(gjson.GetBytes(finalResponse, "output").Array()) == 0 {
// 空 output 由 reconstructResponseOutputFromSSE 整体修补,不在此处理。
return finalResponse
}
if responsesOutputHasCompactionItem(finalResponse) {
return finalResponse
}
item, found := findRawCompactionItemFromSSE(bodyText)
if !found {
return finalResponse
}
patched, err := sjson.SetRawBytes(finalResponse, "output.-1", item)
if err != nil {
return finalResponse
}
return patched
}
// responsesOutputHasCompactionItem reports whether the response JSON already
// carries a compaction item in its output array.
func responsesOutputHasCompactionItem(response []byte) bool {
for _, item := range gjson.GetBytes(response, "output").Array() {
if isResponsesCompactionItemType(item.Get("type").String()) {
return true
}
}
return false
}
// findRawCompactionItemFromSSE 从原始 SSE 事件流中提取第一个 compaction 类
// item 的 raw JSON:output_item.done 优先,output_item.added 兜底。
func findRawCompactionItemFromSSE(bodyText string) (json.RawMessage, bool) {
var found json.RawMessage
pick := func(eventType string) {
forEachOpenAISSEDataPayload(bodyText, func(data []byte) {
if found != nil {
return
}
if strings.TrimSpace(gjson.GetBytes(data, "type").String()) != eventType {
return
}
item := gjson.GetBytes(data, "item")
if !item.IsObject() || !isResponsesCompactionItemType(item.Get("type").String()) {
return
}
found = json.RawMessage(item.Raw)
})
}
pick("response.output_item.done")
if found == nil {
pick("response.output_item.added")
}
return found, found != nil
}
// reconstructResponseOutputFromSSE scans raw SSE body text and returns a
// JSON-encoded output array for a terminal event whose output is empty.
// Raw output_item.done items are preferred: per the Responses protocol they
// are the authoritative final form of each item. Delta accumulation only
// covers text/function_call/reasoning content and silently drops unknown
// item types such as compaction — Codex remote compact v2 then fails with
// "expected exactly one compaction output item, got 0" (#3887).
// Returns (nil, false) if nothing could be reconstructed.
func reconstructResponseOutputFromSSE(bodyText string) ([]byte, bool) {
if outputJSON, ok := collectRawResponsesOutputItemsFromSSE(bodyText); ok {
return outputJSON, true
}
acc := apicompat.NewBufferedResponseAccumulator()
imageOutputs := make([]json.RawMessage, 0, 1)
seenImages := make(map[string]struct{})
@@ -38,7 +38,6 @@ func (s *OpenAIGatewayService) forwardResponsesViaRawChatCompletions(
}
clientStream := responsesReq.Stream
reasoningEffort := extractOpenAIReasoningEffortFromBody(body, originalModel)
serviceTier := extractOpenAIServiceTierFromBody(body)
// custom 工具(如 codex 的 exec)降级为 function 工具转发,回程需按名字还原为
// custom_tool_call 项,先记下名字集合;tool_search 工具同理,回程还原为
@@ -56,6 +55,7 @@ func (s *OpenAIGatewayService) forwardResponsesViaRawChatCompletions(
billingModel := resolveOpenAIForwardModel(account, originalModel, "")
upstreamModel := normalizeOpenAIModelForUpstream(account, billingModel)
reasoningEffort := extractOpenAIReasoningEffortFromBody(body, upstreamModel, billingModel, originalModel)
// 国产模型默认 effort 补充:需要 mappedModel 判定,推迟到 billingModel 算出之后。
reasoningEffort = ApplyThinkingEnabledFallback(reasoningEffort, body, billingModel)
chatReq.Model = upstreamModel
@@ -35,7 +35,7 @@ const (
// 与真实 Codex CLI 的 User-Agent 结构对齐:
// {originator}/{version} ({OS} {OS_version}; {arch}) {terminal}
// 旧值 "codex_cli_rs/0.125.0" 缺少 OS/架构/终端后缀,易被上游指纹识别为非官方客户端。
codexCLIUserAgent = "codex_cli_rs/0.125.0 (Ubuntu 22.4.0; x86_64) xterm-256color"
codexCLIUserAgent = "codex_cli_rs/0.144.1 (Ubuntu 22.4.0; x86_64) xterm-256color"
// codex_cli_only 拒绝时单个请求头日志长度上限(字符)
codexCLIOnlyHeaderValueMaxBytes = 256
@@ -49,7 +49,7 @@ const (
openAIWSRetryBackoffMaxDefault = 2 * time.Second
openAIWSRetryJitterRatioDefault = 0.2
openAICompactSessionSeedKey = "openai_compact_session_seed"
codexCLIVersion = "0.125.0"
codexCLIVersion = "0.144.1"
// Codex 限额快照仅用于后台展示/诊断,不需要每个成功请求都立即落库。
openAICodexSnapshotPersistMinInterval = 30 * time.Second
// 配额自动暂停时,超过该时长仍未刷新的 used% 快照视为陈旧,不再据此暂停账号。
@@ -2482,15 +2482,25 @@ func TestOpenAIBuildUpstreamRequestPreservesCompactPathForAPIKeyBaseURL(t *testi
func TestOpenAIBuildUpstreamRequestOAuthOfficialClientOriginatorCompatibility(t *testing.T) {
gin.SetMode(gin.TestMode)
// 上游要求 originator 与最终 User-Agent 首段配套(issue #3901):
// originator 一律由最终 UA 推导;推导不出官方身份时整体回退默认 Codex CLI 身份。
tests := []struct {
name string
userAgent string
originator string
wantOriginator string
wantUA string
}{
{name: "desktop originator preserved", originator: "Codex Desktop", wantOriginator: "Codex Desktop"},
{name: "vscode originator preserved", originator: "codex_vscode", wantOriginator: "codex_vscode"},
{name: "official ua fallback to codex_cli_rs", userAgent: "Codex Desktop/1.2.3", wantOriginator: "codex_cli_rs"},
{name: "official ua pairs originator", userAgent: "Codex Desktop/1.2.3", wantOriginator: "Codex Desktop", wantUA: "Codex Desktop/1.2.3"},
{
name: "mismatched originator repaired from ua",
userAgent: "codex-tui/0.140.2 (Mac OS X 14.0; arm64) iTerm (codex-tui; 0.140.2)",
originator: "codex_cli_rs",
wantOriginator: "codex-tui",
wantUA: "codex-tui/0.140.2 (Mac OS X 14.0; arm64) iTerm (codex-tui; 0.140.2)",
},
{name: "official originator without ua falls back to default identity", originator: "codex_vscode", wantOriginator: "codex_cli_rs", wantUA: codexCLIUserAgent},
{name: "third-party ua masked to default identity", userAgent: "luna/1.2.0", wantOriginator: "codex_cli_rs", wantUA: codexCLIUserAgent},
}
for _, tt := range tests {
@@ -2515,6 +2525,7 @@ func TestOpenAIBuildUpstreamRequestOAuthOfficialClientOriginatorCompatibility(t
req, err := svc.buildUpstreamRequest(c.Request.Context(), c, account, []byte(`{"model":"gpt-5"}`), "token", false, "", isCodexCLI)
require.NoError(t, err)
require.Equal(t, tt.wantOriginator, req.Header.Get("originator"))
require.Equal(t, tt.wantUA, req.Header.Get("User-Agent"))
})
}
}
@@ -2800,17 +2811,35 @@ func TestParseSSEUsage_SelectiveParsing(t *testing.T) {
}
func TestExtractOpenAIUsageFromJSONBytes_AcceptsResponseAndChatUsageShapes(t *testing.T) {
usage, ok := extractOpenAIUsageFromJSONBytes([]byte(`{"id":"resp_1","usage":{"input_tokens":3,"output_tokens":5,"input_tokens_details":{"cached_tokens":2}}}`))
usage, ok := extractOpenAIUsageFromJSONBytes([]byte(`{"id":"resp_1","usage":{"input_tokens":9,"output_tokens":5,"input_tokens_details":{"cached_tokens":2,"cache_write_tokens":4}}}`))
require.True(t, ok)
require.Equal(t, 3, usage.InputTokens)
require.Equal(t, 9, usage.InputTokens)
require.Equal(t, 5, usage.OutputTokens)
require.Equal(t, 2, usage.CacheReadInputTokens)
require.Equal(t, 4, usage.CacheCreationInputTokens)
usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"type":"response.completed","response":{"usage":{"prompt_tokens":13,"completion_tokens":7,"prompt_tokens_details":{"cached_tokens":4}}}}`))
usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"type":"response.completed","response":{"usage":{"prompt_tokens":13,"completion_tokens":7,"prompt_tokens_details":{"cached_tokens":4,"cache_creation_tokens":3}}}}`))
require.True(t, ok)
require.Equal(t, 13, usage.InputTokens)
require.Equal(t, 7, usage.OutputTokens)
require.Equal(t, 4, usage.CacheReadInputTokens)
require.Equal(t, 3, usage.CacheCreationInputTokens)
usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":11,"output_tokens":2,"cache_write_input_tokens":6}}`))
require.True(t, ok)
require.Equal(t, 6, usage.CacheCreationInputTokens)
usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":20,"output_tokens":2,"cache_creation_input_tokens":19,"input_tokens_details":{"cache_write_tokens":7}}}`))
require.True(t, ok)
require.Equal(t, 7, usage.CacheCreationInputTokens, "官方嵌套字段应优先于兼容顶层别名")
usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":20,"output_tokens":2,"cache_creation_input_tokens":19,"input_tokens_details":{"cache_write_tokens":0}}}`))
require.True(t, ok)
require.Zero(t, usage.CacheCreationInputTokens, "官方嵌套字段显式为零时仍应优先于兼容顶层别名")
usage, ok = extractOpenAIUsageFromJSONBytes([]byte(`{"usage":{"input_tokens":20,"output_tokens":2,"cache_read_input_tokens":19,"input_tokens_details":{"cached_tokens":0}}}`))
require.True(t, ok)
require.Zero(t, usage.CacheReadInputTokens, "官方嵌套缓存读取字段显式为零时仍应优先于兼容顶层别名")
}
func TestExtractCodexFinalResponse_SampleReplay(t *testing.T) {
@@ -119,9 +119,9 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec
ApplyOpenAIImageBillingResolution(result)
}
// 计算实际的新输入token(减去缓存读取的token)
// 因为 input_tokens 包含了 cache_read_tokens,而缓存读取的token不应按输入价格计费
actualInputTokens := result.Usage.InputTokens - result.Usage.CacheReadInputTokens
// OpenAI input_tokens 是总输入,包含缓存读取和缓存写入明细。
// 将三类 token 拆成互斥桶,避免缓存写入同时按普通输入和 cache_write 重复计费。
actualInputTokens := result.Usage.InputTokens - result.Usage.CacheReadInputTokens - result.Usage.CacheCreationInputTokens
if actualInputTokens < 0 {
actualInputTokens = 0
}
@@ -0,0 +1,266 @@
package service
import (
"context"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
func TestNormalizeOpenAIReasoningEffortForGPT56(t *testing.T) {
tests := []struct {
name string
raw string
model string
want string
}{
{name: "Sol 保留 max", raw: "max", model: "gpt-5.6-sol", want: "max"},
{name: "Terra 保留 max", raw: "max", model: "openai/gpt-5.6-terra", want: "max"},
{name: "Luna 后缀保留 max", raw: "max", model: "gpt-5.6-luna-2026-07-09", want: "max"},
{name: "其他模型沿用 xhigh", raw: "max", model: "deepseek-v4-pro", want: "xhigh"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
require.Equal(t, tt.want, normalizeOpenAIReasoningEffortForModel(tt.raw, tt.model))
})
}
}
func TestNormalizeOpenAICodexCompactReasoningEffortDowngradesMax(t *testing.T) {
body := []byte(`{"model":"gpt-5.6-sol","input":"compact me","reasoning":{"effort":"max","summary":"auto"}}`)
normalized, changed, err := normalizeOpenAICodexCompactReasoningEffort(body, "gpt-5.6-sol")
require.NoError(t, err)
require.True(t, changed)
require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(normalized, "model").String())
require.Equal(t, "xhigh", gjson.GetBytes(normalized, "reasoning.effort").String())
require.Equal(t, "auto", gjson.GetBytes(normalized, "reasoning.summary").String())
}
func TestNormalizeOpenAICodexCompactReasoningEffortForAccountScopesCompatibility(t *testing.T) {
gin.SetMode(gin.TestMode)
body := []byte(`{"model":"gpt-5.6-sol","input":"compact me","reasoning":{"effort":"max"}}`)
tests := []struct {
name string
path string
account *Account
changed bool
want string
}{
{
name: "OpenAI OAuth compact 降级",
path: "/openai/v1/responses/compact",
account: &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth},
changed: true,
want: "xhigh",
},
{
name: "OpenAI OAuth 普通请求保留",
path: "/openai/v1/responses",
account: &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth},
want: "max",
},
{
name: "OpenAI API Key compact 保留",
path: "/openai/v1/responses/compact",
account: &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey},
want: "max",
},
{
name: "Grok OAuth compact 保留",
path: "/openai/v1/responses/compact",
account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth},
want: "max",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, tt.path, nil)
normalized, changed, err := normalizeOpenAICodexCompactReasoningEffortForAccount(c, tt.account, body)
require.NoError(t, err)
require.Equal(t, tt.changed, changed)
require.Equal(t, tt.want, gjson.GetBytes(normalized, "reasoning.effort").String())
})
}
}
func TestOpenAIGatewayServiceForwardPreservesGPT56MaxEffort(t *testing.T) {
gin.SetMode(gin.TestMode)
upstream := &httpUpstreamRecorder{
resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)),
},
}
cfg := &config.Config{}
cfg.Security.URLAllowlist.Enabled = false
svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream}
account := &Account{
ID: 7,
Name: "openai-apikey",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://example.com",
},
Extra: map[string]any{"use_responses_api": true},
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
SetOpenAIClientTransport(c, OpenAIClientTransportHTTP)
body := []byte(`{"model":"gpt-5.6-sol","stream":false,"reasoning":{"effort":"max"},"input":"hello"}`)
result, err := svc.Forward(context.Background(), c, account, body)
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, "max", gjson.GetBytes(upstream.lastBody, "reasoning.effort").String())
require.NotNil(t, result.ReasoningEffort)
require.Equal(t, "max", *result.ReasoningEffort)
}
func TestOpenAIGatewayServiceForwardPreservesMappedGPT56MaxEffort(t *testing.T) {
gin.SetMode(gin.TestMode)
upstream := &httpUpstreamRecorder{
resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)),
},
}
cfg := &config.Config{}
cfg.Security.URLAllowlist.Enabled = false
svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream}
account := &Account{
ID: 9,
Name: "openai-apikey-mapped",
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://example.com",
"model_mapping": map[string]any{
"sol": "gpt-5.6-sol",
},
},
Extra: map[string]any{"use_responses_api": true},
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
SetOpenAIClientTransport(c, OpenAIClientTransportHTTP)
body := []byte(`{"model":"sol","stream":false,"reasoning":{"effort":"max"},"input":"hello"}`)
result, err := svc.Forward(context.Background(), c, account, body)
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(upstream.lastBody, "model").String())
require.Equal(t, "max", gjson.GetBytes(upstream.lastBody, "reasoning.effort").String())
require.NotNil(t, result.ReasoningEffort)
require.Equal(t, "max", *result.ReasoningEffort)
}
func TestOpenAIGatewayServiceForwardOAuthCompactDowngradesMaxEffort(t *testing.T) {
gin.SetMode(gin.TestMode)
upstream := &httpUpstreamRecorder{
resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)),
},
}
cfg := &config.Config{}
cfg.Security.URLAllowlist.Enabled = false
svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream}
account := &Account{
ID: 8,
Name: "openai-oauth",
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Concurrency: 1,
Credentials: map[string]any{
"access_token": "oauth-token",
"chatgpt_account_id": "chatgpt-acc",
},
Status: StatusActive,
Schedulable: true,
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses/compact", nil)
SetOpenAIClientTransport(c, OpenAIClientTransportHTTP)
body := []byte(`{"model":"gpt-5.6-sol","instructions":"compact-test","input":"hello","reasoning":{"effort":"max"}}`)
result, err := svc.Forward(context.Background(), c, account, body)
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, upstream.lastReq)
require.Equal(t, chatgptCodexURL+"/compact", upstream.lastReq.URL.String())
require.Equal(t, "xhigh", gjson.GetBytes(upstream.lastBody, "reasoning.effort").String())
require.NotNil(t, result.ReasoningEffort)
require.Equal(t, "xhigh", *result.ReasoningEffort)
}
func TestOpenAIGatewayServiceForwardOAuthResponsesPreservesMaxEffort(t *testing.T) {
gin.SetMode(gin.TestMode)
upstream := &httpUpstreamRecorder{
resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)),
},
}
cfg := &config.Config{}
cfg.Security.URLAllowlist.Enabled = false
svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream}
account := &Account{
ID: 10,
Name: "openai-oauth-responses",
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Concurrency: 1,
Credentials: map[string]any{
"access_token": "oauth-token",
"chatgpt_account_id": "chatgpt-acc",
},
Status: StatusActive,
Schedulable: true,
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
SetOpenAIClientTransport(c, OpenAIClientTransportHTTP)
body := []byte(`{"model":"gpt-5.6-sol","instructions":"response-test","input":"hello","reasoning":{"effort":"max"}}`)
result, err := svc.Forward(context.Background(), c, account, body)
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, upstream.lastReq)
require.Equal(t, chatgptCodexURL, upstream.lastReq.URL.String())
require.Equal(t, "max", gjson.GetBytes(upstream.lastBody, "reasoning.effort").String())
require.NotNil(t, result.ReasoningEffort)
require.Equal(t, "max", *result.ReasoningEffort)
}
@@ -2,6 +2,7 @@ package service
import (
"context"
"encoding/json"
"io"
"net/http"
"net/http/httptest"
@@ -191,6 +192,64 @@ func TestOpenAIGatewayServiceForward_AccountPolicyStripsExplicitImageTool(t *tes
require.NotContains(t, instructions, "image_generation")
}
func TestOpenAIGatewayServiceForward_AccountPolicyStripsImageNamespaceTools(t *testing.T) {
gin.SetMode(gin.TestMode)
tests := []struct {
name string
passthrough bool
}{
{name: "managed forwarding"},
{name: "passthrough forwarding", passthrough: true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
upstream := &httpUpstreamRecorder{
resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"id":"resp_stripped_namespace","model":"gpt-5.5","usage":{"input_tokens":2,"output_tokens":1}}`)),
},
}
svc := newOpenAIImageGenerationControlTestService(upstream)
c, _ := newOpenAIImageGenerationControlTestContext(false, "codex_cli_rs/0.144.1")
account := newOpenAIImageGenerationControlTestAccount()
account.Extra = map[string]any{
featureKeyCodexImageGenerationExplicitToolPolicy: codexImageGenerationExplicitToolPolicyStrip,
"openai_passthrough": tt.passthrough,
}
body := []byte(`{
"model":"gpt-5.5",
"stream":false,
"tools":[
{"type":"function","name":"shell","parameters":{"type":"object"}},
{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]},
{"type":"namespace","name":"code_tools","tools":[{"type":"function","name":"run"}]}
],
"input":[
{"type":"message","role":"user","content":[{"type":"input_text","text":"write code"}]},
{"type":"additional_tools","tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}]}
],
"tool_choice":"auto"
}`)
result, err := svc.Forward(context.Background(), c, account, body)
require.NoError(t, err)
require.NotNil(t, result)
require.NotNil(t, upstream.lastReq)
var forwarded map[string]any
require.NoError(t, json.Unmarshal(upstream.lastBody, &forwarded))
require.False(t, hasOpenAIImageGenerationTool(forwarded))
require.Equal(t, "auto", forwarded["tool_choice"])
require.True(t, gjson.GetBytes(upstream.lastBody, `tools.#(name=="shell")`).Exists())
require.True(t, gjson.GetBytes(upstream.lastBody, `tools.#(name=="code_tools")`).Exists())
require.Equal(t, "write code", gjson.GetBytes(upstream.lastBody, "input.0.content.0.text").String())
})
}
}
func TestOpenAIGatewayServiceForward_ChannelBridgeOverrideEnablesCodexInjection(t *testing.T) {
gin.SetMode(gin.TestMode)
@@ -71,6 +71,14 @@ func normalizeKnownOpenAICodexModel(model string) string {
return "gpt-5.6-terra"
case strings.Contains(normalized, "gpt-5.6-luna"):
return "gpt-5.6-luna"
case normalized == "gpt-5.6":
return "gpt-5.6-sol"
case strings.HasPrefix(normalized, "gpt-5.6-"):
suffix := strings.TrimPrefix(normalized, "gpt-5.6-")
if suffix == "max" || isKnownCodexModelSuffix(suffix) {
return "gpt-5.6-sol"
}
return ""
case strings.Contains(normalized, "gpt-5.5-pro"):
return "gpt-5.5-pro"
case strings.Contains(normalized, "gpt-5.5"):
@@ -98,6 +106,24 @@ func normalizeKnownOpenAICodexModel(model string) string {
}
}
// isOpenAIGPT56Model 判断是否 GPT-5.6 系列模型;入参可为原始模型名
// (含大小写/路径/后缀变体)或已归一化的基名,两者均能正确识别。
func isOpenAIGPT56Model(model string) bool {
normalized := canonicalizeOpenAIModelAliasSpelling(model)
if normalized == "gpt-5.6" {
return true
}
if suffix, ok := strings.CutPrefix(normalized, "gpt-5.6-"); ok && (suffix == "max" || isKnownCodexModelSuffix(suffix)) {
return true
}
for _, prefix := range []string{"gpt-5.6-sol", "gpt-5.6-terra", "gpt-5.6-luna"} {
if normalized == prefix || strings.HasPrefix(normalized, prefix+"-") {
return true
}
}
return false
}
func appendUsageBillingModelCandidate(candidates []string, seen map[string]struct{}, model string) []string {
trimmed := strings.TrimSpace(model)
if trimmed == "" {
@@ -0,0 +1,36 @@
package service
import (
"testing"
"github.com/stretchr/testify/require"
)
func TestNormalizeKnownOpenAICodexModel_BareGPT56RoutesToSol(t *testing.T) {
tests := map[string]string{
"gpt-5.6": "gpt-5.6-sol",
"openai/gpt-5.6": "gpt-5.6-sol",
"gpt5.6": "gpt-5.6-sol",
"gpt-5.6-high": "gpt-5.6-sol",
"gpt-5.6-max": "gpt-5.6-sol",
"gpt-5.6-2026-07-09": "gpt-5.6-sol",
"openai/gpt-5.6-max": "gpt-5.6-sol",
}
for input, expected := range tests {
t.Run(input, func(t *testing.T) {
require.Equal(t, expected, normalizeKnownOpenAICodexModel(input))
})
}
}
func TestUsageBillingModelCandidates_BareGPT56IncludesSol(t *testing.T) {
require.Equal(t,
[]string{"gpt-5.6", "gpt-5.6-sol"},
usageBillingModelCandidates("gpt-5.6"),
)
require.Equal(t,
[]string{"openai/gpt-5.6", "gpt-5.6", "gpt-5.6-sol"},
usageBillingModelCandidates("openai/gpt-5.6"),
)
}
@@ -252,6 +252,18 @@ func TestNormalizeOpenAIModelForUpstream(t *testing.T) {
model string
want string
}{
{
name: "oauth routes bare GPT-5.6 alias to Sol",
account: &Account{Type: AccountTypeOAuth},
model: "gpt-5.6",
want: "gpt-5.6-sol",
},
{
name: "oauth routes provider-prefixed GPT-5.6 alias to Sol",
account: &Account{Type: AccountTypeOAuth},
model: "openai/gpt-5.6",
want: "gpt-5.6-sol",
},
{
name: "oauth preserves unknown non codex model",
account: &Account{Type: AccountTypeOAuth},
@@ -282,6 +294,12 @@ func TestNormalizeOpenAIModelForUpstream(t *testing.T) {
model: "codex-auto-review",
want: "codex-auto-review",
},
{
name: "apikey preserves official bare GPT-5.6 alias",
account: &Account{Type: AccountTypeAPIKey},
model: "gpt-5.6",
want: "gpt-5.6",
},
{
name: "apikey preserves custom compatible model",
account: &Account{Type: AccountTypeAPIKey},
@@ -703,7 +703,9 @@ func TestOpenAIGatewayService_OAuthLegacy_CompositeCodexUAUsesCodexOriginator(t
_, err := svc.Forward(context.Background(), c, account, inputBody)
require.NoError(t, err)
require.NotNil(t, upstream.lastReq)
require.Equal(t, "codex_cli_rs", upstream.lastReq.Header.Get("originator"))
// 浏览器型复合 UA 被替换为默认 Codex UA(codex-tui 形态),originator 随最终 UA 配套(issue #3901)。
require.Equal(t, DefaultOpenAICodexUserAgent, upstream.lastReq.Header.Get("User-Agent"))
require.Equal(t, "codex-tui", upstream.lastReq.Header.Get("originator"))
require.NotEqual(t, "opencode", upstream.lastReq.Header.Get("originator"))
}
@@ -1111,6 +1113,55 @@ func TestOpenAIGatewayService_OAuthPassthrough_NonCodexUAFallbackToCodexUA(t *te
require.Equal(t, codexCLIUserAgent, upstream.lastReq.Header.Get("User-Agent"))
}
// 回归(issue #3901):codex-tui 等官方 UA 在透传模式下必须逐字保留,且 originator
// 由最终 UA 推导配套——历史实现会把 codex-tui UA 强改为 codex_cli_rs,而 originator
// 保留客户端原值,造成 originator/UA 首段错配被上游 404。
func TestOpenAIGatewayService_OAuthPassthrough_CodexTuiIdentityPreservedAndPaired(t *testing.T) {
gin.SetMode(gin.TestMode)
const tuiUA = "codex-tui/0.140.2 (Mac OS X 14.0; arm64) iTerm (codex-tui; 0.140.2)"
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(nil))
c.Request.Header.Set("User-Agent", tuiUA)
// 客户端携带错配的 originator,也必须按最终 UA 重配。
c.Request.Header.Set("originator", "codex_cli_rs")
inputBody := []byte(`{"model":"gpt-5.2","stream":false,"store":true,"input":[{"type":"text","text":"hi"}]}`)
resp := &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"rid"}},
Body: io.NopCloser(strings.NewReader("data: [DONE]\n\n")),
}
upstream := &httpUpstreamRecorder{resp: resp}
svc := &OpenAIGatewayService{
cfg: &config.Config{Gateway: config.GatewayConfig{ForceCodexCLI: false}},
httpUpstream: upstream,
}
account := &Account{
ID: 123,
Name: "acc",
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Concurrency: 1,
Credentials: map[string]any{"access_token": "oauth-token", "chatgpt_account_id": "chatgpt-acc"},
Extra: map[string]any{"openai_passthrough": true},
Status: StatusActive,
Schedulable: true,
RateMultiplier: f64p(1),
}
_, err := svc.Forward(context.Background(), c, account, inputBody)
require.NoError(t, err)
require.NotNil(t, upstream.lastReq)
require.Equal(t, tuiUA, upstream.lastReq.Header.Get("User-Agent"))
require.Equal(t, "codex-tui", upstream.lastReq.Header.Get("originator"))
}
func TestOpenAIGatewayService_CodexCLIOnly_RejectsNonCodexClient(t *testing.T) {
gin.SetMode(gin.TestMode)
@@ -0,0 +1,122 @@
package service
import (
"context"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
func TestExtractOpenAIReasoningEffortFromBodyModelCandidates(t *testing.T) {
bodyWithoutEffort := []byte(`{"model":"whatever","input":"hello"}`)
bodyWithMax := []byte(`{"model":"sol","reasoning":{"effort":"max"},"input":"hello"}`)
tests := []struct {
name string
body []byte
candidates []string
want string // "" 表示期望 nil
}{
{
name: "后缀推导回退到原始模型(OAuth 上游模型已剥后缀)",
body: bodyWithoutEffort,
candidates: []string{"gpt-5.4", "gpt-5.4", "gpt-5.4-xhigh"},
want: "xhigh",
},
{
name: "GPT-5.6 后缀 max 经原始模型推导保留",
body: bodyWithoutEffort,
candidates: []string{"gpt-5.6-sol", "gpt-5.6-sol", "gpt-5.6-sol-max"},
want: "max",
},
{
name: "显式 max 用第一个非空候选(映射后模型)判定",
body: bodyWithMax,
candidates: []string{"gpt-5.6-sol", "sol"},
want: "max",
},
{
name: "显式 max 非 5.6 首候选仍折叠为 xhigh",
body: bodyWithMax,
candidates: []string{"gpt-5.4", "sol"},
want: "xhigh",
},
{
name: "所有候选均无后缀时返回 nil",
body: bodyWithoutEffort,
candidates: []string{"gpt-5.4", "gpt-5.4", "gpt-5.4"},
want: "",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := extractOpenAIReasoningEffortFromBody(tt.body, tt.candidates...)
if tt.want == "" {
require.Nil(t, got)
return
}
require.NotNil(t, got)
require.Equal(t, tt.want, *got)
})
}
}
func TestExtractOpenAIReasoningEffortModelCandidates(t *testing.T) {
reqBody := map[string]any{"model": "gpt-5.3-codex-high", "input": "hello"}
got := extractOpenAIReasoningEffort(reqBody, "gpt-5.3-codex", "gpt-5.3-codex-high")
require.NotNil(t, got)
require.Equal(t, "high", *got)
}
// 回归:OAuth 账号请求后缀式模型(无显式 reasoning 字段)时,上游模型被
// normalizeCodexModel 剥掉 effort 后缀,用量元数据的 effort 必须仍能从
// 原始模型名后缀推导出来。
func TestOpenAIGatewayServiceForwardOAuthDerivesEffortFromSuffixModel(t *testing.T) {
gin.SetMode(gin.TestMode)
upstream := &httpUpstreamRecorder{
resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)),
},
}
cfg := &config.Config{}
cfg.Security.URLAllowlist.Enabled = false
svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream}
account := &Account{
ID: 11,
Name: "openai-oauth-suffix",
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Concurrency: 1,
Credentials: map[string]any{
"access_token": "oauth-token",
"chatgpt_account_id": "chatgpt-acc",
},
Status: StatusActive,
Schedulable: true,
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
SetOpenAIClientTransport(c, OpenAIClientTransportHTTP)
body := []byte(`{"model":"gpt-5.3-codex-xhigh","instructions":"suffix-test","input":"hello","stream":false}`)
result, err := svc.Forward(context.Background(), c, account, body)
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, "gpt-5.3-codex", gjson.GetBytes(upstream.lastBody, "model").String())
require.NotNil(t, result.ReasoningEffort)
require.Equal(t, "xhigh", *result.ReasoningEffort)
}
@@ -270,7 +270,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
normalized = next
}
if isCodexCLI && codexImageGenerationExplicitToolPolicy == codexImageGenerationExplicitToolPolicyStrip {
if stripped, changed, stripErr := stripOpenAIImageGenerationToolFromRawPayload(normalized); stripErr != nil {
if stripped, changed, stripErr := stripOpenAIImageGenerationToolsFromRawPayload(normalized); stripErr != nil {
return openAIWSClientPayload{}, NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", stripErr)
} else if changed {
normalized = stripped
@@ -920,7 +920,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
Model: originalModel,
UpstreamModel: mappedModel,
ServiceTier: extractOpenAIServiceTierFromBody(payload),
ReasoningEffort: ApplyThinkingEnabledFallback(extractOpenAIReasoningEffortFromBody(payload, originalModel), payload, mappedModel),
ReasoningEffort: ApplyThinkingEnabledFallback(extractOpenAIReasoningEffortFromBody(payload, mappedModel, originalModel), payload, mappedModel),
Stream: reqStream,
OpenAIWSMode: true,
ResponseHeaders: lease.HandshakeHeaders(),
@@ -33,6 +33,7 @@ func TestIsOpenAIWSClientDisconnectError(t *testing.T) {
{name: "ws_policy_violation", err: coderws.CloseError{Code: coderws.StatusPolicyViolation}, want: false},
{name: "wrapped_eof_message", err: errors.New("failed to get reader: failed to read frame header: EOF"), want: true},
{name: "connection_reset_by_peer", err: errors.New("failed to read frame header: read tcp 127.0.0.1:1234->127.0.0.1:5678: read: connection reset by peer"), want: true},
{name: "windows_connection_reset", err: errors.New("failed to get reader: failed to read frame header: read tcp 127.0.0.1:1234->127.0.0.1:5678: wsarecv: An existing connection was forcibly closed by the remote host."), want: true},
{name: "broken_pipe", err: errors.New("write tcp 127.0.0.1:1234->127.0.0.1:5678: write: broken pipe"), want: true},
}
@@ -152,6 +153,24 @@ func TestStripCodexSparkImageGenerationToolFromRawPayload(t *testing.T) {
require.True(t, gjson.GetBytes(updated, `tools.#(type=="function")`).Exists())
})
t.Run("strips_namespace_tools_for_spark", func(t *testing.T) {
payload := []byte(`{
"type":"response.create",
"model":"gpt-5.3-codex-spark",
"input":[
{"type":"message","role":"user","content":"hello"},
{"type":"additional_tools","tools":[{"type":"namespace","name":"image_gen"}]}
],
"tool_choice":{"type":"namespace","name":"image_gen"}
}`)
updated, changed, err := stripCodexSparkImageGenerationToolFromRawPayload(payload, "gpt-5.3-codex-spark")
require.NoError(t, err)
require.True(t, changed)
require.False(t, IsImageGenerationIntent(openAIResponsesEndpoint, "gpt-5.3-codex-spark", updated))
require.Equal(t, "hello", gjson.GetBytes(updated, "input.0.content").String())
require.False(t, gjson.GetBytes(updated, "tool_choice").Exists())
})
t.Run("keeps_image_generation_for_non_spark", func(t *testing.T) {
payload := []byte(`{"type":"response.create","model":"gpt-5.3-codex","tools":[{"type":"image_generation","output_format":"png"}]}`)
updated, changed, err := stripCodexSparkImageGenerationToolFromRawPayload(payload, "gpt-5.3-codex")
@@ -169,24 +188,61 @@ func TestStripCodexSparkImageGenerationToolFromRawPayload(t *testing.T) {
})
}
func TestStripOpenAIImageGenerationToolFromRawPayload(t *testing.T) {
payload := []byte(`{
"type":"response.create",
"model":"gpt-5.4",
"tools":[
{"type":"function","name":"shell"},
{"type":"image_generation","output_format":"png"}
],
"tool_choice":{"type":"image_generation"}
}`)
func TestStripOpenAIImageGenerationToolsFromRawPayload(t *testing.T) {
t.Run("flat image tool", func(t *testing.T) {
payload := []byte(`{
"type":"response.create",
"model":"gpt-5.4",
"tools":[
{"type":"function","name":"shell"},
{"type":"image_generation","output_format":"png"}
],
"tool_choice":{"type":"image_generation"}
}`)
updated, changed, err := stripOpenAIImageGenerationToolFromRawPayload(payload)
updated, changed, err := stripOpenAIImageGenerationToolsFromRawPayload(payload)
require.NoError(t, err)
require.True(t, changed)
require.False(t, gjson.GetBytes(updated, `tools.#(type=="image_generation")`).Exists())
require.True(t, gjson.GetBytes(updated, `tools.#(type=="function")`).Exists())
require.False(t, gjson.GetBytes(updated, "tool_choice").Exists())
require.NoError(t, err)
require.True(t, changed)
require.False(t, gjson.GetBytes(updated, `tools.#(type=="image_generation")`).Exists())
require.True(t, gjson.GetBytes(updated, `tools.#(type=="function")`).Exists())
require.False(t, gjson.GetBytes(updated, "tool_choice").Exists())
})
t.Run("namespace and Responses Lite tools", func(t *testing.T) {
payload := []byte(`{
"type":"response.create",
"model":"gpt-5.5",
"tools":[
{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]},
{"type":"namespace","name":"code_tools","tools":[{"type":"function","name":"run"}]}
],
"input":[
{"type":"message","role":"user","content":"hello"},
{"type":"additional_tools","tools":[{"type":"namespace","name":"image_gen"}]}
],
"tool_choice":{"type":"namespace","name":"image_gen"}
}`)
updated, changed, err := stripOpenAIImageGenerationToolsFromRawPayload(payload)
require.NoError(t, err)
require.True(t, changed)
require.False(t, IsImageGenerationIntent(openAIResponsesEndpoint, "gpt-5.5", updated))
require.True(t, gjson.GetBytes(updated, `tools.#(name=="code_tools")`).Exists())
require.Equal(t, "hello", gjson.GetBytes(updated, "input.0.content").String())
require.False(t, gjson.GetBytes(updated, "tool_choice").Exists())
})
t.Run("non-image namespace is unchanged", func(t *testing.T) {
payload := []byte(`{"type":"response.create","model":"gpt-5.5","tools":[{"type":"namespace","name":"code_tools"}]}`)
updated, changed, err := stripOpenAIImageGenerationToolsFromRawPayload(payload)
require.NoError(t, err)
require.False(t, changed)
require.Equal(t, payload, updated)
})
}
func TestAlignStoreDisabledPreviousResponseID(t *testing.T) {
@@ -660,6 +660,7 @@ func isOpenAIWSClientDisconnectError(err error) bool {
strings.Contains(message, "use of closed network connection") ||
strings.Contains(message, "connection reset by peer") ||
strings.Contains(message, "broken pipe") ||
strings.Contains(message, "an existing connection was forcibly closed by the remote host") ||
strings.Contains(message, "an established connection was aborted")
}
@@ -10,7 +10,6 @@ import (
"net/url"
"strings"
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
"github.com/gin-gonic/gin"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
@@ -127,8 +126,11 @@ func (s *OpenAIGatewayService) buildOpenAIWSHeaders(
if s != nil && s.cfg != nil && s.cfg.Gateway.ForceCodexCLI {
headers.Set("user-agent", codexCLIUserAgent)
}
if account != nil && account.Type == AccountTypeOAuth && !openai.IsCodexCLIRequest(headers.Get("user-agent")) {
headers.Set("user-agent", codexCLIUserAgent)
// 终态收口:originator 必须与最终 user-agent 首段配套且为官方身份,非官方 UA 整体回退为
// 默认 Codex CLI 身份(承接原「非 Codex UA 兜底」,并修复其把 codex-tui 等官方 UA 改写为
// codex_cli_rs 造成的 originator 错配 404),详见 issue #3901。
if account != nil && account.Type == AccountTypeOAuth {
enforceCodexIdentityHeaders(headers)
}
// 账号级请求头覆写(仅 openai api_key 账号启用时生效;OAuth 路径 no-op)。
@@ -670,15 +670,24 @@ func TestOpenAIGatewayService_Forward_WSv2_OAuthStoreFalseByDefault(t *testing.T
func TestOpenAIGatewayService_Forward_WSv2_OAuthOriginatorCompatibility(t *testing.T) {
gin.SetMode(gin.TestMode)
// 上游要求 originator 与最终 user-agent 首段配套(issue #3901):
// originator 一律由最终 UA 推导;推导不出官方身份时整体回退默认 Codex CLI 身份。
tests := []struct {
name string
userAgent string
originator string
wantOriginator string
wantUA string
}{
{name: "desktop originator preserved", originator: "Codex Desktop", wantOriginator: "Codex Desktop"},
{name: "vscode originator preserved", originator: "codex_vscode", wantOriginator: "codex_vscode"},
{name: "official ua fallback to codex_cli_rs", userAgent: "Codex Desktop/1.2.3", wantOriginator: "codex_cli_rs"},
{name: "official ua pairs originator", userAgent: "Codex Desktop/1.2.3", wantOriginator: "Codex Desktop", wantUA: "Codex Desktop/1.2.3"},
{
name: "mismatched originator repaired from ua",
userAgent: "codex-tui/0.140.2 (Mac OS X 14.0; arm64) iTerm (codex-tui; 0.140.2)",
originator: "codex_cli_rs",
wantOriginator: "codex-tui",
wantUA: "codex-tui/0.140.2 (Mac OS X 14.0; arm64) iTerm (codex-tui; 0.140.2)",
},
{name: "official originator without ua falls back to default identity", originator: "codex_vscode", wantOriginator: "codex_cli_rs", wantUA: codexCLIUserAgent},
}
for _, tt := range tests {
@@ -743,6 +752,7 @@ func TestOpenAIGatewayService_Forward_WSv2_OAuthOriginatorCompatibility(t *testi
require.NoError(t, err)
require.NotNil(t, result)
require.Equal(t, tt.wantOriginator, captureDialer.lastHeaders.Get("originator"))
require.Equal(t, tt.wantUA, captureDialer.lastHeaders.Get("user-agent"))
})
}
}
@@ -262,15 +262,9 @@ func populateOpenAIUsageFromResponseJSON(body []byte, usage *OpenAIUsage) {
if usage == nil || len(body) == 0 {
return
}
values := gjson.GetManyBytes(
body,
"usage.input_tokens",
"usage.output_tokens",
"usage.input_tokens_details.cached_tokens",
)
usage.InputTokens = int(values[0].Int())
usage.OutputTokens = int(values[1].Int())
usage.CacheReadInputTokens = int(values[2].Int())
if parsed, ok := extractOpenAIUsageFromJSONBytes(body); ok {
*usage = parsed
}
}
func getOpenAIGroupIDFromContext(c *gin.Context) int64 {
@@ -3,7 +3,6 @@ package service
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
@@ -693,7 +692,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2(
ImageCount: imageCounter.Count(),
ImageOutputSizes: imageCounter.Sizes(),
ServiceTier: extractOpenAIServiceTier(reqBody),
ReasoningEffort: extractOpenAIReasoningEffort(reqBody, originalModel),
ReasoningEffort: extractOpenAIReasoningEffort(reqBody, mappedModel, originalModel),
Stream: reqStream,
OpenAIWSMode: true,
ResponseHeaders: lease.HandshakeHeaders(),
@@ -710,23 +709,8 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2(
// Codex clients advertise it by default. Returns the (possibly unchanged) payload,
// whether it changed, and any JSON decode error.
func stripCodexSparkImageGenerationToolFromRawPayload(payload []byte, model string) ([]byte, bool, error) {
if !isCodexSparkModel(model) || !openAIRequestBodyHasImageGenerationTool(payload) {
if !isCodexSparkModel(model) {
return payload, false, nil
}
return stripOpenAIImageGenerationToolFromRawPayload(payload)
}
func stripOpenAIImageGenerationToolFromRawPayload(payload []byte) ([]byte, bool, error) {
payloadMap := make(map[string]any)
if err := json.Unmarshal(payload, &payloadMap); err != nil {
return payload, false, err
}
if !stripOpenAIImageGenerationTools(payloadMap) {
return payload, false, nil
}
rebuilt, err := json.Marshal(payloadMap)
if err != nil {
return payload, false, err
}
return rebuilt, true, nil
return stripOpenAIImageGenerationToolsFromRawPayload(payload)
}
@@ -263,7 +263,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
Model: originalModel,
UpstreamModel: mappedModel,
ServiceTier: extractOpenAIServiceTierFromBody(body),
ReasoningEffort: ApplyThinkingEnabledFallback(extractOpenAIReasoningEffortFromBody(body, originalModel), body, mappedModel),
ReasoningEffort: ApplyThinkingEnabledFallback(extractOpenAIReasoningEffortFromBody(body, mappedModel, originalModel), body, mappedModel),
Stream: reqStream,
OpenAIWSMode: true,
ResponseHeaders: cloneHeader(resp.Header),
@@ -787,7 +787,7 @@ func parseUsageAndAccumulate(
parsedUsage := Usage{
InputTokens: inputTokens,
OutputTokens: outputTokens,
CacheCreationInputTokens: int(usageResult.Get("cache_creation_input_tokens").Int()),
CacheCreationInputTokens: openAICacheCreationTokensFromUsage(usageResult),
CacheReadInputTokens: cachedTokens,
ImageOutputTokens: int(imageTokens),
}
@@ -810,6 +810,31 @@ func parseUsageIntField(value gjson.Result, required bool) (int, bool) {
return int(value.Int()), true
}
func openAICacheCreationTokensFromUsage(value gjson.Result) int {
for _, field := range []string{
"input_tokens_details.cache_write_tokens",
"prompt_tokens_details.cache_write_tokens",
"input_tokens_details.cache_creation_tokens",
"prompt_tokens_details.cache_creation_tokens",
} {
result := value.Get(field)
if result.Exists() {
return max(int(result.Int()), 0)
}
}
for _, field := range []string{
"cache_write_tokens",
"cache_creation_input_tokens",
"cache_write_input_tokens",
"cache_creation_tokens",
} {
if tokens := int(value.Get(field).Int()); tokens > 0 {
return tokens
}
}
return 0
}
func enrichResult(result *RelayResult, state *relayState, duration time.Duration) {
if result == nil {
return
@@ -300,7 +300,7 @@ func TestParseUsageAndEnrichCoverage(t *testing.T) {
require.Equal(t, 0, state.usage.OutputTokens)
require.Equal(t, 0, state.usage.CacheReadInputTokens)
parseUsageAndAccumulate(state, []byte(`{"type":"response.completed","response":{"usage":{"input_tokens":2,"output_tokens":1,"input_tokens_details":{"cached_tokens":1},"cache_creation_input_tokens":4,"output_tokens_details":{"image_tokens":3}}}}`), "response.completed", nil)
parseUsageAndAccumulate(state, []byte(`{"type":"response.completed","response":{"usage":{"input_tokens":2,"output_tokens":1,"input_tokens_details":{"cached_tokens":1,"cache_write_tokens":4},"output_tokens_details":{"image_tokens":3}}}}`), "response.completed", nil)
require.Equal(t, 2, state.usage.InputTokens)
require.Equal(t, 1, state.usage.OutputTokens)
require.Equal(t, 1, state.usage.CacheReadInputTokens)
@@ -335,6 +335,13 @@ func TestParseUsageAndAccumulateAcceptsChatUsageAliases(t *testing.T) {
require.Equal(t, got, state.usage)
}
func TestOpenAICacheCreationTokensFromUsageNestedZeroWins(t *testing.T) {
t.Parallel()
usage := gjson.Parse(`{"input_tokens_details":{"cache_write_tokens":0},"cache_creation_input_tokens":19}`)
require.Zero(t, openAICacheCreationTokensFromUsage(usage))
}
func TestEmitTurnCompleteCoverage(t *testing.T) {
t.Parallel()
@@ -142,12 +142,12 @@ func newOpenAIWSPassthroughUsageMeta(initialRequestModel string, firstFrame []by
return meta
}
func (m *openAIWSPassthroughUsageMeta) initFromFirstFrame(policyOutput []byte) {
func (m *openAIWSPassthroughUsageMeta) initFromFirstFrame(policyOutput []byte, mappedModel string) {
if m == nil {
return
}
m.serviceTier.Store(extractOpenAIServiceTierFromBody(policyOutput))
m.reasoningEffort.Store(extractOpenAIReasoningEffortFromBody(policyOutput, m.sessionRequestModel))
m.reasoningEffort.Store(extractOpenAIReasoningEffortFromBody(policyOutput, mappedModel, m.sessionRequestModel))
}
func (m *openAIWSPassthroughUsageMeta) updateSessionRequestModel(payload []byte) {
@@ -169,12 +169,12 @@ func (m *openAIWSPassthroughUsageMeta) requestModelForFrame(payload []byte) stri
return m.sessionRequestModel
}
func (m *openAIWSPassthroughUsageMeta) updateFromResponseCreate(policyOutput []byte, requestModelForFrame string) {
func (m *openAIWSPassthroughUsageMeta) updateFromResponseCreate(policyOutput []byte, mappedModel string, requestModelForFrame string) {
if m == nil {
return
}
m.serviceTier.Store(extractOpenAIServiceTierFromBody(policyOutput))
m.reasoningEffort.Store(extractOpenAIReasoningEffortFromBody(policyOutput, requestModelForFrame))
m.reasoningEffort.Store(extractOpenAIReasoningEffortFromBody(policyOutput, mappedModel, requestModelForFrame))
}
func openAIWSPassthroughRequestModelForFrame(payload []byte) string {
@@ -311,7 +311,7 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough(
// 因此使用 atomic.Pointer[string] 在 filter(runClientToUpstream
// goroutine)和 OnTurnComplete / final result(runUpstreamToClient
// goroutine)之间同步当前 turn 的 usage metadata。
usageMeta.initFromFirstFrame(firstClientMessage)
usageMeta.initFromFirstFrame(firstClientMessage, capturedSessionModel)
promptCacheKey := strings.TrimSpace(gjson.GetBytes(firstClientMessage, "prompt_cache_key").String())
wsURL, err := s.buildOpenAIResponsesWSURL(account)
@@ -455,7 +455,7 @@ func (s *OpenAIGatewayService) proxyResponsesWebSocketV2Passthrough(
// service_tier 时按 default 处理,billing 应如实反映。
if policyErr == nil && blocked == nil &&
strings.TrimSpace(gjson.GetBytes(payload, "type").String()) == "response.create" {
usageMeta.updateFromResponseCreate(out, requestModelForThisFrame)
usageMeta.updateFromResponseCreate(out, model, requestModelForThisFrame)
}
return out, blocked, policyErr
},
@@ -0,0 +1,42 @@
//go:build unit
package service
import (
"testing"
"github.com/stretchr/testify/require"
)
func TestWSPassthroughUsageMeta_InitFromFirstFrame_MappedModelCandidate(t *testing.T) {
body := []byte(`{"type":"response.create","model":"sol","reasoning":{"effort":"max"}}`)
meta := newOpenAIWSPassthroughUsageMeta("sol", body)
meta.initFromFirstFrame(body, "gpt-5.6-sol")
got := meta.reasoningEffort.Load()
require.NotNil(t, got, "reasoning effort should be set")
require.Equal(t, "max", *got, "mapped model gpt-5.6-sol should preserve max")
}
func TestWSPassthroughUsageMeta_InitFromFirstFrame_NonGPT56FallsBackToXHigh(t *testing.T) {
body := []byte(`{"type":"response.create","model":"gpt-5.4","reasoning":{"effort":"max"}}`)
meta := newOpenAIWSPassthroughUsageMeta("gpt-5.4", body)
meta.initFromFirstFrame(body, "gpt-5.4")
got := meta.reasoningEffort.Load()
require.NotNil(t, got)
require.Equal(t, "xhigh", *got, "non-5.6 model should normalize max to xhigh")
}
func TestWSPassthroughUsageMeta_UpdateFromResponseCreate_MappedModelCandidate(t *testing.T) {
body := []byte(`{"type":"response.create","model":"sol","reasoning":{"effort":"max"}}`)
meta := newOpenAIWSPassthroughUsageMeta("sol", body)
meta.updateFromResponseCreate(body, "gpt-5.6-sol", "sol")
got := meta.reasoningEffort.Load()
require.NotNil(t, got)
require.Equal(t, "max", *got, "mapped model should preserve max on multi-turn update")
}
+212 -54
View File
@@ -28,6 +28,12 @@ import (
// misconfigured to point at us, or when our orders table has been wiped).
var ErrOrderNotFound = errors.New("payment order not found")
const paymentFulfillmentLeaseDuration = 5 * time.Minute
type paymentFulfillmentLease struct {
version time.Time
}
// --- Payment Notification & Fulfillment ---
func (s *PaymentService) HandlePaymentNotification(ctx context.Context, n *payment.PaymentNotification, pk string) error {
@@ -188,10 +194,8 @@ func (s *PaymentService) alreadyProcessed(ctx context.Context, o *dbent.PaymentO
switch cur.Status {
case OrderStatusCompleted, OrderStatusRefunded:
return nil
case OrderStatusFailed:
case OrderStatusFailed, OrderStatusPaid, OrderStatusRecharging:
return s.executeFulfillment(ctx, o.ID)
case OrderStatusPaid, OrderStatusRecharging:
return fmt.Errorf("order %d is being processed", o.ID)
case OrderStatusExpired:
slog.Warn("webhook payment success for expired order beyond grace period",
"orderID", o.ID,
@@ -231,23 +235,74 @@ func (s *PaymentService) ExecuteBalanceFulfillment(ctx context.Context, oid int6
if psIsRefundStatus(o.Status) {
return infraerrors.BadRequest("INVALID_STATUS", "refund-related order cannot fulfill")
}
if o.Status != OrderStatusPaid && o.Status != OrderStatusFailed {
if o.Status != OrderStatusPaid && o.Status != OrderStatusFailed && o.Status != OrderStatusRecharging {
return infraerrors.BadRequest("INVALID_STATUS", "order cannot fulfill in status "+o.Status)
}
c, err := s.entClient.PaymentOrder.Update().Where(paymentorder.IDEQ(oid), paymentorder.StatusIn(OrderStatusPaid, OrderStatusFailed)).SetStatus(OrderStatusRecharging).Save(ctx)
lease, err := s.acquirePaymentFulfillmentLease(ctx, o)
if err != nil {
return fmt.Errorf("lock: %w", err)
return err
}
if c == 0 {
if lease == nil {
return nil
}
if err := s.doBalance(ctx, o); err != nil {
s.markFailed(ctx, oid, err)
if err := s.doBalance(ctx, o, lease); err != nil {
s.markFailed(ctx, oid, lease, err)
return err
}
return nil
}
func (s *PaymentService) acquirePaymentFulfillmentLease(ctx context.Context, o *dbent.PaymentOrder) (*paymentFulfillmentLease, error) {
if o == nil {
return nil, infraerrors.BadRequest("INVALID_STATUS", "nil payment order")
}
now := time.Now().UTC().Truncate(time.Microsecond)
staleBefore := now.Add(-paymentFulfillmentLeaseDuration)
updated, err := s.entClient.PaymentOrder.Update().
Where(
paymentorder.IDEQ(o.ID),
paymentorder.Or(
paymentorder.StatusIn(OrderStatusPaid, OrderStatusFailed),
paymentorder.And(
paymentorder.StatusEQ(OrderStatusRecharging),
paymentorder.UpdatedAtLTE(staleBefore),
),
),
).
SetStatus(OrderStatusRecharging).
SetUpdatedAt(now).
ClearFailedAt().
ClearFailedReason().
Save(ctx)
if err != nil {
return nil, fmt.Errorf("acquire fulfillment lease: %w", err)
}
if updated == 0 {
current, getErr := s.entClient.PaymentOrder.Get(ctx, o.ID)
if getErr != nil {
return nil, fmt.Errorf("reload fulfillment lease: %w", getErr)
}
if current.Status == OrderStatusCompleted {
return nil, nil
}
if current.Status == OrderStatusRecharging {
return nil, infraerrors.Conflict("CONFLICT", "order is being processed")
}
return nil, infraerrors.Conflict("CONFLICT", "order status changed while acquiring fulfillment lease")
}
// Reload the persisted timestamp instead of trusting application clock precision.
claimed, err := s.entClient.PaymentOrder.Get(ctx, o.ID)
if err != nil {
return nil, fmt.Errorf("reload acquired fulfillment lease: %w", err)
}
if claimed.Status != OrderStatusRecharging {
return nil, infraerrors.Conflict("CONFLICT", "fulfillment lease was lost")
}
return &paymentFulfillmentLease{version: claimed.UpdatedAt}, nil
}
// redeemAction represents the idempotency decision for balance fulfillment.
type redeemAction int
@@ -272,7 +327,7 @@ func resolveRedeemAction(existing *RedeemCode, lookupErr error) redeemAction {
return redeemActionRedeem
}
func (s *PaymentService) doBalance(ctx context.Context, o *dbent.PaymentOrder) error {
func (s *PaymentService) doBalance(ctx context.Context, o *dbent.PaymentOrder, lease *paymentFulfillmentLease) error {
// Idempotency: check if redeem code already exists (from a previous partial run)
existing, lookupErr := s.redeemService.GetByCode(ctx, o.RechargeCode)
action := resolveRedeemAction(existing, lookupErr)
@@ -283,7 +338,7 @@ func (s *PaymentService) doBalance(ctx context.Context, o *dbent.PaymentOrder) e
return err
}
// Code already created and redeemed — just mark completed
return s.markCompleted(ctx, o, "RECHARGE_SUCCESS")
return s.markCompleted(ctx, o, lease, "RECHARGE_SUCCESS")
case redeemActionCreate:
rc := &RedeemCode{Code: o.RechargeCode, Type: RedeemTypeBalance, Value: o.Amount, Status: StatusUnused}
if err := s.redeemService.CreateCode(ctx, rc); err != nil {
@@ -298,21 +353,37 @@ func (s *PaymentService) doBalance(ctx context.Context, o *dbent.PaymentOrder) e
if err := s.applyAffiliateRebateForOrder(ctx, o); err != nil {
return err
}
return s.markCompleted(ctx, o, "RECHARGE_SUCCESS")
return s.markCompleted(ctx, o, lease, "RECHARGE_SUCCESS")
}
func (s *PaymentService) markCompleted(ctx context.Context, o *dbent.PaymentOrder, auditAction string) error {
func (s *PaymentService) markCompleted(ctx context.Context, o *dbent.PaymentOrder, lease *paymentFulfillmentLease, auditAction string) error {
if lease == nil {
return errors.New("missing payment fulfillment lease")
}
now := time.Now()
_, err := s.entClient.PaymentOrder.Update().Where(paymentorder.IDEQ(o.ID), paymentorder.StatusEQ(OrderStatusRecharging)).SetStatus(OrderStatusCompleted).SetCompletedAt(now).Save(ctx)
updated, err := s.entClient.PaymentOrder.Update().Where(
paymentorder.IDEQ(o.ID),
paymentorder.StatusEQ(OrderStatusRecharging),
paymentorder.UpdatedAtEQ(lease.version),
).SetStatus(OrderStatusCompleted).SetCompletedAt(now).Save(ctx)
if err != nil {
return fmt.Errorf("mark completed: %w", err)
}
s.writeAuditLog(ctx, o.ID, auditAction, "system", map[string]any{
"rechargeCode": o.RechargeCode,
"creditedAmount": o.Amount,
"payAmount": o.PayAmount,
})
s.dispatchPaymentFulfillmentNotification(o, auditAction)
if updated == 0 {
current, getErr := s.entClient.PaymentOrder.Get(ctx, o.ID)
if getErr == nil && current.Status == OrderStatusCompleted {
return nil
}
return infraerrors.Conflict("CONFLICT", "fulfillment lease was lost before completion")
}
if !s.hasAuditLog(ctx, o.ID, auditAction) {
s.writeAuditLog(ctx, o.ID, auditAction, "system", map[string]any{
"rechargeCode": o.RechargeCode,
"creditedAmount": o.Amount,
"payAmount": o.PayAmount,
})
s.dispatchPaymentFulfillmentNotification(o, auditAction)
}
return nil
}
@@ -404,51 +475,138 @@ func (s *PaymentService) ExecuteSubscriptionFulfillment(ctx context.Context, oid
if psIsRefundStatus(o.Status) {
return infraerrors.BadRequest("INVALID_STATUS", "refund-related order cannot fulfill")
}
if o.Status != OrderStatusPaid && o.Status != OrderStatusFailed {
if o.Status != OrderStatusPaid && o.Status != OrderStatusFailed && o.Status != OrderStatusRecharging {
return infraerrors.BadRequest("INVALID_STATUS", "order cannot fulfill in status "+o.Status)
}
if o.SubscriptionGroupID == nil || o.SubscriptionDays == nil {
return infraerrors.BadRequest("INVALID_STATUS", "missing subscription info")
}
c, err := s.entClient.PaymentOrder.Update().Where(paymentorder.IDEQ(oid), paymentorder.StatusIn(OrderStatusPaid, OrderStatusFailed)).SetStatus(OrderStatusRecharging).Save(ctx)
lease, err := s.acquirePaymentFulfillmentLease(ctx, o)
if err != nil {
return fmt.Errorf("lock: %w", err)
return err
}
if c == 0 {
if lease == nil {
return nil
}
if err := s.doSub(ctx, o); err != nil {
s.markFailed(ctx, oid, err)
if err := s.doSub(ctx, o, lease); err != nil {
s.markFailed(ctx, oid, lease, err)
return err
}
return nil
}
func (s *PaymentService) doSub(ctx context.Context, o *dbent.PaymentOrder) error {
func (s *PaymentService) doSub(ctx context.Context, o *dbent.PaymentOrder, lease *paymentFulfillmentLease) error {
gid := *o.SubscriptionGroupID
days := *o.SubscriptionDays
g, err := s.groupRepo.GetByID(ctx, gid)
if err != nil || g.Status != payment.EntityStatusActive {
return fmt.Errorf("group %d no longer exists or inactive", gid)
}
assigned := s.hasAuditLog(ctx, o.ID, "SUBSCRIPTION_ASSIGNED") || s.hasAuditLog(ctx, o.ID, "SUBSCRIPTION_SUCCESS")
if !assigned {
orderNote := fmt.Sprintf("payment order %d", o.ID)
_, _, err = s.subscriptionSvc.AssignOrExtendSubscription(ctx, &AssignSubscriptionInput{UserID: o.UserID, GroupID: gid, ValidityDays: days, AssignedBy: 0, Notes: orderNote})
if err != nil {
return fmt.Errorf("assign subscription: %w", err)
}
s.writeAuditLog(ctx, o.ID, "SUBSCRIPTION_ASSIGNED", "system", map[string]any{
"groupID": gid,
"validityDays": days,
})
} else {
slog.Info("subscription already assigned for order, skipping", "orderID", o.ID, "groupID", gid)
if err := s.ensurePaymentSubscriptionAssigned(ctx, o, gid, days); err != nil {
return err
}
if err := s.applyAffiliateRebateForOrder(ctx, o); err != nil {
return err
}
return s.markCompleted(ctx, o, "SUBSCRIPTION_SUCCESS")
return s.markCompleted(ctx, o, lease, "SUBSCRIPTION_SUCCESS")
}
func (s *PaymentService) ensurePaymentSubscriptionAssigned(ctx context.Context, o *dbent.PaymentOrder, groupID int64, days int) error {
if s.subscriptionSvc == nil {
return errors.New("subscription service is unavailable")
}
tx, err := s.entClient.Tx(ctx)
if err != nil {
return fmt.Errorf("begin subscription fulfillment tx: %w", err)
}
defer func() { _ = tx.Rollback() }()
txCtx := dbent.NewTxContext(ctx, tx)
txClient := tx.Client()
alreadyAssigned, err := hasPaymentSubscriptionAssignmentAudit(txCtx, txClient, o.ID)
if err != nil {
return fmt.Errorf("check subscription assignment audit: %w", err)
}
recoveredFromNote := false
if !alreadyAssigned {
orderNote := paymentSubscriptionOrderNote(o.ID)
existing, lookupErr := s.subscriptionSvc.userSubRepo.GetByUserIDAndGroupID(txCtx, o.UserID, groupID)
switch {
case lookupErr == nil && existing != nil && hasPaymentSubscriptionOrderNote(existing.Notes, orderNote):
recoveredFromNote = true
case lookupErr != nil && !errors.Is(lookupErr, ErrSubscriptionNotFound):
return fmt.Errorf("check existing subscription assignment: %w", lookupErr)
default:
if _, _, err := s.subscriptionSvc.assignOrExtendSubscription(txCtx, &AssignSubscriptionInput{
UserID: o.UserID,
GroupID: groupID,
ValidityDays: days,
AssignedBy: 0,
Notes: orderNote,
}, true); err != nil {
return fmt.Errorf("assign subscription: %w", err)
}
}
detail, _ := json.Marshal(map[string]any{
"groupID": groupID,
"validityDays": days,
"recoveredFromNote": recoveredFromNote,
})
if _, err := txClient.PaymentAuditLog.Create().
SetOrderID(strconv.FormatInt(o.ID, 10)).
SetAction("SUBSCRIPTION_ASSIGNED").
SetDetail(string(detail)).
SetOperator("system").
Save(txCtx); err != nil {
if dbent.IsConstraintError(err) {
_ = tx.Rollback()
claimed, checkErr := hasPaymentSubscriptionAssignmentAudit(ctx, s.entClient, o.ID)
if checkErr == nil && claimed {
return s.subscriptionSvc.invalidateSubscriptionCaches(o.UserID, groupID)
}
}
return fmt.Errorf("record subscription assignment audit: %w", err)
}
} else {
slog.Info("subscription already assigned for order, skipping", "orderID", o.ID, "groupID", groupID)
}
if err := tx.Commit(); err != nil {
return fmt.Errorf("commit subscription fulfillment tx: %w", err)
}
// Assignment cache invalidation is deferred while this transaction is open,
// then performed synchronously against the committed subscription.
if err := s.subscriptionSvc.invalidateSubscriptionCaches(o.UserID, groupID); err != nil {
return fmt.Errorf("invalidate subscription cache after fulfillment: %w", err)
}
return nil
}
func hasPaymentSubscriptionAssignmentAudit(ctx context.Context, client *dbent.Client, orderID int64) (bool, error) {
count, err := client.PaymentAuditLog.Query().
Where(
paymentauditlog.OrderIDEQ(strconv.FormatInt(orderID, 10)),
paymentauditlog.ActionIn("SUBSCRIPTION_ASSIGNED", "SUBSCRIPTION_SUCCESS"),
).
Limit(1).
Count(ctx)
return count > 0, err
}
func paymentSubscriptionOrderNote(orderID int64) string {
return fmt.Sprintf("payment order %d", orderID)
}
func hasPaymentSubscriptionOrderNote(notes string, orderNote string) bool {
for _, line := range strings.Split(strings.ReplaceAll(notes, "\r\n", "\n"), "\n") {
if strings.TrimSpace(line) == orderNote {
return true
}
}
return false
}
func (s *PaymentService) hasAuditLog(ctx context.Context, orderID int64, action string) bool {
@@ -642,13 +800,20 @@ func (s *PaymentService) updateClaimedAffiliateRebateAudit(ctx context.Context,
return nil
}
func (s *PaymentService) markFailed(ctx context.Context, oid int64, cause error) {
func (s *PaymentService) markFailed(ctx context.Context, oid int64, lease *paymentFulfillmentLease, cause error) {
if lease == nil {
slog.Error("mark FAILED without fulfillment lease", "orderID", oid)
return
}
now := time.Now()
r := psErrMsg(cause)
// Only mark FAILED if still in RECHARGING state — prevents overwriting
// a COMPLETED order when markCompleted failed but fulfillment succeeded.
// The lease version prevents a stale worker from overwriting a newer owner.
c, e := s.entClient.PaymentOrder.Update().
Where(paymentorder.IDEQ(oid), paymentorder.StatusEQ(OrderStatusRecharging)).
Where(
paymentorder.IDEQ(oid),
paymentorder.StatusEQ(OrderStatusRecharging),
paymentorder.UpdatedAtEQ(lease.version),
).
SetStatus(OrderStatusFailed).SetFailedAt(now).SetFailedReason(r).Save(ctx)
if e != nil {
slog.Error("mark FAILED", "orderID", oid, "error", e)
@@ -669,18 +834,11 @@ func (s *PaymentService) RetryFulfillment(ctx context.Context, oid int64) error
if psIsRefundStatus(o.Status) {
return infraerrors.BadRequest("INVALID_STATUS", "refund-related order cannot retry")
}
if o.Status == OrderStatusRecharging {
return infraerrors.Conflict("CONFLICT", "order is being processed")
}
if o.Status == OrderStatusCompleted {
return infraerrors.BadRequest("INVALID_STATUS", "order already completed")
}
if o.Status != OrderStatusFailed && o.Status != OrderStatusPaid {
return infraerrors.BadRequest("INVALID_STATUS", "only paid and failed orders can retry")
}
_, err = s.entClient.PaymentOrder.Update().Where(paymentorder.IDEQ(oid), paymentorder.StatusIn(OrderStatusFailed, OrderStatusPaid)).SetStatus(OrderStatusPaid).ClearFailedAt().ClearFailedReason().Save(ctx)
if err != nil {
return fmt.Errorf("reset for retry: %w", err)
if o.Status != OrderStatusFailed && o.Status != OrderStatusPaid && o.Status != OrderStatusRecharging {
return infraerrors.BadRequest("INVALID_STATUS", "only paid, failed, and recoverable recharging orders can retry")
}
s.writeAuditLog(ctx, oid, "RECHARGE_RETRY", "admin", map[string]any{"detail": "admin manual retry"})
return s.executeFulfillment(ctx, oid)
@@ -13,6 +13,7 @@ import (
dbent "github.com/Wei-Shaw/sub2api/ent"
"github.com/Wei-Shaw/sub2api/ent/paymentauditlog"
"github.com/Wei-Shaw/sub2api/internal/payment"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
@@ -586,6 +587,238 @@ func TestPaymentAmountToleranceForThreeDecimalCurrency(t *testing.T) {
assert.InDelta(t, 0.0005, paymentAmountToleranceForCurrency("KWD"), 1e-12)
}
func TestRetryFulfillmentRejectsFreshRechargingLease(t *testing.T) {
ctx := context.Background()
client := newPaymentConfigServiceTestClient(t)
order := createPaymentFulfillmentSubscriptionOrder(t, ctx, client, OrderStatusRecharging, time.Now())
svc := &PaymentService{entClient: client}
err := svc.RetryFulfillment(ctx, order.ID)
require.Error(t, err)
require.Equal(t, "CONFLICT", infraerrors.Reason(err))
reloaded, getErr := client.PaymentOrder.Get(ctx, order.ID)
require.NoError(t, getErr)
require.Equal(t, OrderStatusRecharging, reloaded.Status)
}
func TestAlreadyProcessedRecoversStaleRechargingLease(t *testing.T) {
ctx := context.Background()
client := newPaymentConfigServiceTestClient(t)
ensurePaymentAuditOrderActionUniqueIndex(t, ctx, client)
order := createPaymentFulfillmentSubscriptionOrder(
t,
ctx,
client,
OrderStatusRecharging,
time.Now().Add(-paymentFulfillmentLeaseDuration-time.Minute),
)
_, err := client.PaymentAuditLog.Create().
SetOrderID(strconv.FormatInt(order.ID, 10)).
SetAction("SUBSCRIPTION_ASSIGNED").
SetDetail(`{"groupID":7,"validityDays":30}`).
SetOperator("system").
Save(ctx)
require.NoError(t, err)
groupRepo := &subscriptionGroupRepoStub{
group: &Group{ID: 7, Status: payment.EntityStatusActive, SubscriptionType: SubscriptionTypeSubscription},
}
svc := &PaymentService{
entClient: client,
groupRepo: groupRepo,
subscriptionSvc: NewSubscriptionService(groupRepo, userSubRepoNoop{}, nil, nil, nil),
}
require.NoError(t, svc.alreadyProcessed(ctx, order))
reloaded, err := client.PaymentOrder.Get(ctx, order.ID)
require.NoError(t, err)
require.Equal(t, OrderStatusCompleted, reloaded.Status)
}
func TestFulfillmentLeaseVersionRejectsStaleWorker(t *testing.T) {
ctx := context.Background()
client := newPaymentConfigServiceTestClient(t)
staleAt := time.Now().Add(-paymentFulfillmentLeaseDuration - time.Minute)
order := createPaymentFulfillmentSubscriptionOrder(t, ctx, client, OrderStatusRecharging, staleAt)
svc := &PaymentService{entClient: client}
firstLease, err := svc.acquirePaymentFulfillmentLease(ctx, order)
require.NoError(t, err)
require.NotNil(t, firstLease)
_, err = client.PaymentOrder.UpdateOneID(order.ID).SetUpdatedAt(staleAt).Save(ctx)
require.NoError(t, err)
time.Sleep(time.Millisecond)
staleOrder, err := client.PaymentOrder.Get(ctx, order.ID)
require.NoError(t, err)
secondLease, err := svc.acquirePaymentFulfillmentLease(ctx, staleOrder)
require.NoError(t, err)
require.NotNil(t, secondLease)
require.False(t, firstLease.version.Equal(secondLease.version))
err = svc.markCompleted(ctx, order, firstLease, "SUBSCRIPTION_SUCCESS")
require.Error(t, err)
require.Equal(t, "CONFLICT", infraerrors.Reason(err))
svc.markFailed(ctx, order.ID, firstLease, errors.New("stale worker failure"))
reloaded, err := client.PaymentOrder.Get(ctx, order.ID)
require.NoError(t, err)
require.Equal(t, OrderStatusRecharging, reloaded.Status)
require.NoError(t, svc.markCompleted(ctx, order, secondLease, "SUBSCRIPTION_SUCCESS"))
}
func TestExecuteBalanceFulfillmentRecoversAfterRedeemWithoutCreditingAgain(t *testing.T) {
ctx := context.Background()
client := newPaymentConfigServiceTestClient(t)
ensurePaymentAuditOrderActionUniqueIndex(t, ctx, client)
staleAt := time.Now().Add(-paymentFulfillmentLeaseDuration - time.Minute)
order := createPaymentFulfillmentSubscriptionOrder(t, ctx, client, OrderStatusRecharging, staleAt)
order, err := client.PaymentOrder.UpdateOneID(order.ID).
SetOrderType(payment.OrderTypeBalance).
ClearPlanID().
ClearSubscriptionGroupID().
ClearSubscriptionDays().
SetUpdatedAt(staleAt).
Save(ctx)
require.NoError(t, err)
redeemRepo := &redeemCodeRepoStub{codesByCode: map[string]*RedeemCode{
order.RechargeCode: {
ID: 101,
Code: order.RechargeCode,
Type: RedeemTypeBalance,
Value: order.Amount,
Status: StatusUsed,
},
}}
svc := &PaymentService{
entClient: client,
redeemService: &RedeemService{redeemRepo: redeemRepo},
}
require.NoError(t, svc.ExecuteBalanceFulfillment(ctx, order.ID))
require.Empty(t, redeemRepo.useCalls, "an already-used order code must not be redeemed again")
reloaded, err := client.PaymentOrder.Get(ctx, order.ID)
require.NoError(t, err)
require.Equal(t, OrderStatusCompleted, reloaded.Status)
}
func TestExecuteSubscriptionFulfillmentRecoversCommittedAssignmentWithoutExtendingAgain(t *testing.T) {
ctx := context.Background()
client := newPaymentConfigServiceTestClient(t)
ensurePaymentAuditOrderActionUniqueIndex(t, ctx, client)
staleAt := time.Now().Add(-paymentFulfillmentLeaseDuration - time.Minute)
order := createPaymentFulfillmentSubscriptionOrder(t, ctx, client, OrderStatusRecharging, staleAt)
expiresAt := time.Now().Add(30 * 24 * time.Hour).Truncate(time.Second)
subRepo := newSubscriptionUserSubRepoStub()
subRepo.seed(&UserSubscription{
ID: 99,
UserID: order.UserID,
GroupID: *order.SubscriptionGroupID,
StartsAt: time.Now().Add(-time.Hour),
ExpiresAt: expiresAt,
Status: SubscriptionStatusActive,
Notes: "manual note\n" + paymentSubscriptionOrderNote(order.ID) + "\nretained note",
})
groupRepo := &subscriptionGroupRepoStub{
group: &Group{ID: 7, Status: payment.EntityStatusActive, SubscriptionType: SubscriptionTypeSubscription},
}
svc := &PaymentService{
entClient: client,
groupRepo: groupRepo,
subscriptionSvc: NewSubscriptionService(groupRepo, subRepo, nil, nil, nil),
}
require.NoError(t, svc.ExecuteSubscriptionFulfillment(ctx, order.ID))
assertPaymentSubscriptionExpiry(t, subRepo, order, expiresAt)
assignmentAuditCount, err := client.PaymentAuditLog.Query().
Where(
paymentauditlog.OrderIDEQ(strconv.FormatInt(order.ID, 10)),
paymentauditlog.ActionEQ("SUBSCRIPTION_ASSIGNED"),
).
Count(ctx)
require.NoError(t, err)
require.Equal(t, 1, assignmentAuditCount)
// Simulate another stale recovery attempt after completion. The durable audit
// must make replay a no-op for the subscription entitlement.
_, err = client.PaymentOrder.UpdateOneID(order.ID).
SetStatus(OrderStatusRecharging).
SetUpdatedAt(staleAt).
ClearCompletedAt().
Save(ctx)
require.NoError(t, err)
require.NoError(t, svc.ExecuteSubscriptionFulfillment(ctx, order.ID))
assertPaymentSubscriptionExpiry(t, subRepo, order, expiresAt)
assignmentAuditCount, err = client.PaymentAuditLog.Query().
Where(
paymentauditlog.OrderIDEQ(strconv.FormatInt(order.ID, 10)),
paymentauditlog.ActionEQ("SUBSCRIPTION_ASSIGNED"),
).
Count(ctx)
require.NoError(t, err)
require.Equal(t, 1, assignmentAuditCount)
}
func TestHasPaymentSubscriptionOrderNoteRequiresIndependentExactLine(t *testing.T) {
t.Parallel()
require.True(t, hasPaymentSubscriptionOrderNote("before\r\npayment order 42\r\nafter", "payment order 42"))
require.False(t, hasPaymentSubscriptionOrderNote("payment order 420", "payment order 42"))
require.False(t, hasPaymentSubscriptionOrderNote("prefix payment order 42 suffix", "payment order 42"))
}
func createPaymentFulfillmentSubscriptionOrder(
t *testing.T,
ctx context.Context,
client *dbent.Client,
status string,
updatedAt time.Time,
) *dbent.PaymentOrder {
t.Helper()
user, err := client.User.Create().
SetEmail("fulfillment-" + strconv.FormatInt(time.Now().UnixNano(), 10) + "@example.com").
SetPasswordHash("hash").
SetUsername("payment-fulfillment-user").
Save(ctx)
require.NoError(t, err)
order, err := client.PaymentOrder.Create().
SetUserID(user.ID).
SetUserEmail(user.Email).
SetUserName(user.Username).
SetAmount(80).
SetPayAmount(80).
SetFeeRate(0).
SetRechargeCode("PAY-SUB-" + strconv.FormatInt(time.Now().UnixNano(), 10)).
SetOutTradeNo("sub2_fulfillment_" + strconv.FormatInt(time.Now().UnixNano(), 10)).
SetPaymentType(payment.TypeAlipay).
SetPaymentTradeNo("trade-fulfillment").
SetOrderType(payment.OrderTypeSubscription).
SetPlanID(100).
SetSubscriptionGroupID(7).
SetSubscriptionDays(30).
SetStatus(status).
SetPaidAt(time.Now().Add(-time.Hour)).
SetExpiresAt(time.Now().Add(time.Hour)).
SetClientIP("127.0.0.1").
SetSrcHost("api.example.com").
SetUpdatedAt(updatedAt).
Save(ctx)
require.NoError(t, err)
return order
}
func assertPaymentSubscriptionExpiry(t *testing.T, repo *subscriptionUserSubRepoStub, order *dbent.PaymentOrder, expected time.Time) {
t.Helper()
sub, err := repo.GetByUserIDAndGroupID(context.Background(), order.UserID, *order.SubscriptionGroupID)
require.NoError(t, err)
require.True(t, sub.ExpiresAt.Equal(expected), "subscription expiry changed from %s to %s", expected, sub.ExpiresAt)
}
func TestExecuteSubscriptionFulfillmentAppliesAffiliateRebate(t *testing.T) {
ctx := context.Background()
client := newPaymentConfigServiceTestClient(t)
+87 -4
View File
@@ -35,6 +35,57 @@ var (
Mode: "chat",
SupportsPromptCaching: true,
}
openAIGPT56SolFallbackPricing = &LiteLLMModelPricing{
InputCostPerToken: 5e-06,
InputCostPerTokenPriority: 1e-05,
OutputCostPerToken: 3e-05,
OutputCostPerTokenPriority: 6e-05,
CacheCreationInputTokenCost: 6.25e-06,
CacheCreationInputTokenCostPriority: 1.25e-05,
CacheReadInputTokenCost: 5e-07,
CacheReadInputTokenCostPriority: 1e-06,
LongContextInputTokenThreshold: openAIGPT54LongContextInputThreshold,
LongContextInputCostMultiplier: openAIGPT54LongContextInputMultiplier,
LongContextOutputCostMultiplier: openAIGPT54LongContextOutputMultiplier,
SupportsServiceTier: true,
LiteLLMProvider: "openai",
Mode: "chat",
SupportsPromptCaching: true,
}
openAIGPT56TerraFallbackPricing = &LiteLLMModelPricing{
InputCostPerToken: 2.5e-06,
InputCostPerTokenPriority: 5e-06,
OutputCostPerToken: 1.5e-05,
OutputCostPerTokenPriority: 3e-05,
CacheCreationInputTokenCost: 3.125e-06,
CacheCreationInputTokenCostPriority: 6.25e-06,
CacheReadInputTokenCost: 2.5e-07,
CacheReadInputTokenCostPriority: 5e-07,
LongContextInputTokenThreshold: openAIGPT54LongContextInputThreshold,
LongContextInputCostMultiplier: openAIGPT54LongContextInputMultiplier,
LongContextOutputCostMultiplier: openAIGPT54LongContextOutputMultiplier,
SupportsServiceTier: true,
LiteLLMProvider: "openai",
Mode: "chat",
SupportsPromptCaching: true,
}
openAIGPT56LunaFallbackPricing = &LiteLLMModelPricing{
InputCostPerToken: 1e-06,
InputCostPerTokenPriority: 2e-06,
OutputCostPerToken: 6e-06,
OutputCostPerTokenPriority: 1.2e-05,
CacheCreationInputTokenCost: 1.25e-06,
CacheCreationInputTokenCostPriority: 2.5e-06,
CacheReadInputTokenCost: 1e-07,
CacheReadInputTokenCostPriority: 2e-07,
LongContextInputTokenThreshold: openAIGPT54LongContextInputThreshold,
LongContextInputCostMultiplier: openAIGPT54LongContextInputMultiplier,
LongContextOutputCostMultiplier: openAIGPT54LongContextOutputMultiplier,
SupportsServiceTier: true,
LiteLLMProvider: "openai",
Mode: "chat",
SupportsPromptCaching: true,
}
openAIGPT54MiniFallbackPricing = &LiteLLMModelPricing{
InputCostPerToken: 7.5e-07,
OutputCostPerToken: 4.5e-06,
@@ -61,6 +112,7 @@ type LiteLLMModelPricing struct {
OutputCostPerToken float64 `json:"output_cost_per_token"`
OutputCostPerTokenPriority float64 `json:"output_cost_per_token_priority"`
CacheCreationInputTokenCost float64 `json:"cache_creation_input_token_cost"`
CacheCreationInputTokenCostPriority float64 `json:"cache_creation_input_token_cost_priority"`
CacheCreationInputTokenCostAbove1hr float64 `json:"cache_creation_input_token_cost_above_1hr"`
CacheReadInputTokenCost float64 `json:"cache_read_input_token_cost"`
CacheReadInputTokenCostPriority float64 `json:"cache_read_input_token_cost_priority"`
@@ -93,9 +145,13 @@ type LiteLLMRawEntry struct {
OutputCostPerToken *float64 `json:"output_cost_per_token"`
OutputCostPerTokenPriority *float64 `json:"output_cost_per_token_priority"`
CacheCreationInputTokenCost *float64 `json:"cache_creation_input_token_cost"`
CacheCreationInputTokenCostPriority *float64 `json:"cache_creation_input_token_cost_priority"`
CacheCreationInputTokenCostAbove1hr *float64 `json:"cache_creation_input_token_cost_above_1hr"`
CacheReadInputTokenCost *float64 `json:"cache_read_input_token_cost"`
CacheReadInputTokenCostPriority *float64 `json:"cache_read_input_token_cost_priority"`
LongContextInputTokenThreshold *int `json:"long_context_input_token_threshold"`
LongContextInputCostMultiplier *float64 `json:"long_context_input_cost_multiplier"`
LongContextOutputCostMultiplier *float64 `json:"long_context_output_cost_multiplier"`
SupportsServiceTier bool `json:"supports_service_tier"`
LiteLLMProvider string `json:"litellm_provider"`
Mode string `json:"mode"`
@@ -406,6 +462,9 @@ func (s *PricingService) parsePricingData(body []byte) (map[string]*LiteLLMModel
if entry.CacheCreationInputTokenCost != nil {
pricing.CacheCreationInputTokenCost = *entry.CacheCreationInputTokenCost
}
if entry.CacheCreationInputTokenCostPriority != nil {
pricing.CacheCreationInputTokenCostPriority = *entry.CacheCreationInputTokenCostPriority
}
if entry.CacheCreationInputTokenCostAbove1hr != nil {
pricing.CacheCreationInputTokenCostAbove1hr = *entry.CacheCreationInputTokenCostAbove1hr
}
@@ -415,6 +474,15 @@ func (s *PricingService) parsePricingData(body []byte) (map[string]*LiteLLMModel
if entry.CacheReadInputTokenCostPriority != nil {
pricing.CacheReadInputTokenCostPriority = *entry.CacheReadInputTokenCostPriority
}
if entry.LongContextInputTokenThreshold != nil {
pricing.LongContextInputTokenThreshold = *entry.LongContextInputTokenThreshold
}
if entry.LongContextInputCostMultiplier != nil {
pricing.LongContextInputCostMultiplier = *entry.LongContextInputCostMultiplier
}
if entry.LongContextOutputCostMultiplier != nil {
pricing.LongContextOutputCostMultiplier = *entry.LongContextOutputCostMultiplier
}
if entry.OutputCostPerImage != nil {
pricing.OutputCostPerImage = *entry.OutputCostPerImage
}
@@ -666,6 +734,12 @@ func normalizeModelNameForPricing(model string) string {
model = strings.TrimLeft(model, "/")
if canonical := canonicalizeOpenAIModelAliasSpelling(model); canonical != "" {
if canonical == "gpt-5.6" {
return "gpt-5.6-sol"
}
if suffix, ok := strings.CutPrefix(canonical, "gpt-5.6-"); ok && (suffix == "max" || isKnownCodexModelSuffix(suffix)) {
return "gpt-5.6-sol"
}
return canonical
}
return model
@@ -837,11 +911,20 @@ func (s *PricingService) matchOpenAIModel(model string) *LiteLLMModelPricing {
}
}
// GPT-5.6(sol / terra / luna)回退到 GPT-5.4 定价
if strings.HasPrefix(model, "gpt-5.6") {
if strings.HasPrefix(model, "gpt-5.6-sol") {
logger.With(zap.String("component", "service.pricing")).
Info(fmt.Sprintf("[Pricing] OpenAI fallback matched %s -> %s", model, "gpt-5.4(static)"))
return openAIGPT54FallbackPricing
Info(fmt.Sprintf("[Pricing] OpenAI fallback matched %s -> %s", model, "gpt-5.6-sol(static)"))
return openAIGPT56SolFallbackPricing
}
if strings.HasPrefix(model, "gpt-5.6-terra") {
logger.With(zap.String("component", "service.pricing")).
Info(fmt.Sprintf("[Pricing] OpenAI fallback matched %s -> %s", model, "gpt-5.6-terra(static)"))
return openAIGPT56TerraFallbackPricing
}
if strings.HasPrefix(model, "gpt-5.6-luna") {
logger.With(zap.String("component", "service.pricing")).
Info(fmt.Sprintf("[Pricing] OpenAI fallback matched %s -> %s", model, "gpt-5.6-luna(static)"))
return openAIGPT56LunaFallbackPricing
}
// GPT-5.5 回退到 GPT-5.4 定价
@@ -19,8 +19,12 @@ func TestParsePricingData_ParsesPriorityAndServiceTierFields(t *testing.T) {
"output_cost_per_token": 0.000015,
"output_cost_per_token_priority": 0.00003,
"cache_creation_input_token_cost": 0.0000025,
"cache_creation_input_token_cost_priority": 0.000005,
"cache_read_input_token_cost": 0.00000025,
"cache_read_input_token_cost_priority": 0.0000005,
"long_context_input_token_threshold": 272000,
"long_context_input_cost_multiplier": 2,
"long_context_output_cost_multiplier": 1.5,
"supports_service_tier": true,
"supports_prompt_caching": true,
"litellm_provider": "openai",
@@ -34,10 +38,226 @@ func TestParsePricingData_ParsesPriorityAndServiceTierFields(t *testing.T) {
require.NotNil(t, pricing)
require.InDelta(t, 5e-6, pricing.InputCostPerTokenPriority, 1e-12)
require.InDelta(t, 3e-5, pricing.OutputCostPerTokenPriority, 1e-12)
require.InDelta(t, 5e-6, pricing.CacheCreationInputTokenCostPriority, 1e-12)
require.InDelta(t, 5e-7, pricing.CacheReadInputTokenCostPriority, 1e-12)
require.Equal(t, 272000, pricing.LongContextInputTokenThreshold)
require.InDelta(t, 2.0, pricing.LongContextInputCostMultiplier, 1e-12)
require.InDelta(t, 1.5, pricing.LongContextOutputCostMultiplier, 1e-12)
require.True(t, pricing.SupportsServiceTier)
}
func TestBillingService_GPT56CacheWritePricingUsesOfficialMultiplier(t *testing.T) {
tests := []struct {
model string
input float64
inputPriority float64
output float64
outputPriority float64
cacheRead float64
cacheReadPriority float64
}{
{model: "gpt-5.6-sol", input: 5e-6, inputPriority: 10e-6, output: 30e-6, outputPriority: 60e-6, cacheRead: 0.5e-6, cacheReadPriority: 1e-6},
{model: "gpt-5.6-terra", input: 2.5e-6, inputPriority: 5e-6, output: 15e-6, outputPriority: 30e-6, cacheRead: 0.25e-6, cacheReadPriority: 0.5e-6},
{model: "gpt-5.6-luna", input: 1e-6, inputPriority: 2e-6, output: 6e-6, outputPriority: 12e-6, cacheRead: 0.1e-6, cacheReadPriority: 0.2e-6},
}
for _, tt := range tests {
t.Run(tt.model, func(t *testing.T) {
pricingSvc := &PricingService{pricingData: map[string]*LiteLLMModelPricing{
tt.model: {
InputCostPerToken: tt.input,
InputCostPerTokenPriority: tt.inputPriority,
OutputCostPerToken: tt.output,
OutputCostPerTokenPriority: tt.outputPriority,
CacheReadInputTokenCost: tt.cacheRead,
CacheReadInputTokenCostPriority: tt.cacheReadPriority,
},
}}
svc := NewBillingService(&config.Config{}, pricingSvc)
pricing, err := svc.GetModelPricing(tt.model)
require.NoError(t, err)
require.InDelta(t, tt.input*1.25, pricing.CacheCreationPricePerToken, 1e-12)
require.InDelta(t, tt.inputPriority*1.25, pricing.CacheCreationPricePerTokenPriority, 1e-12)
require.Equal(t, 272000, pricing.LongContextInputThreshold)
require.InDelta(t, 2.0, pricing.LongContextInputMultiplier, 1e-12)
require.InDelta(t, 1.5, pricing.LongContextOutputMultiplier, 1e-12)
tokens := UsageTokens{InputTokens: 700, OutputTokens: 50, CacheCreationTokens: 200, CacheReadTokens: 100}
standard, err := svc.CalculateCostWithServiceTier(tt.model, tokens, 1, "")
require.NoError(t, err)
require.InDelta(t, 200*tt.input*1.25, standard.CacheCreationCost, 1e-12)
priority, err := svc.CalculateCostWithServiceTier(tt.model, tokens, 1, "priority")
require.NoError(t, err)
require.InDelta(t, 200*tt.inputPriority*1.25, priority.CacheCreationCost, 1e-12)
flex, err := svc.CalculateCostWithServiceTier(tt.model, tokens, 1, "flex")
require.NoError(t, err)
require.InDelta(t, 200*tt.input*1.25*0.5, flex.CacheCreationCost, 1e-12)
})
}
}
func TestBillingService_GPT56UsesLongContextPricingAcrossModelsAndTiers(t *testing.T) {
models := []struct {
name string
input, cached float64
cacheWrite, output float64
}{
{name: "gpt-5.6-sol", input: 5e-6, cached: 0.5e-6, cacheWrite: 6.25e-6, output: 30e-6},
{name: "gpt-5.6-terra", input: 2.5e-6, cached: 0.25e-6, cacheWrite: 3.125e-6, output: 15e-6},
{name: "gpt-5.6-luna", input: 1e-6, cached: 0.1e-6, cacheWrite: 1.25e-6, output: 6e-6},
}
tiers := []struct {
name string
priceScale float64
}{
{name: "standard", priceScale: 1},
{name: "priority", priceScale: 2},
{name: "flex", priceScale: 0.5},
}
tokens := UsageTokens{
InputTokens: 100000,
CacheCreationTokens: 100000,
CacheReadTokens: 73000,
OutputTokens: 10,
}
for _, model := range models {
for _, tier := range tiers {
t.Run(model.name+"/"+tier.name, func(t *testing.T) {
svc := NewBillingService(&config.Config{}, nil)
serviceTier := ""
if tier.name != "standard" {
serviceTier = tier.name
}
cost, err := svc.CalculateCostWithServiceTier(model.name, tokens, 1, serviceTier)
require.NoError(t, err)
require.InDelta(t, float64(tokens.InputTokens)*model.input*tier.priceScale*2, cost.InputCost, 1e-12)
require.InDelta(t, float64(tokens.CacheCreationTokens)*model.cacheWrite*tier.priceScale*2, cost.CacheCreationCost, 1e-12)
require.InDelta(t, float64(tokens.CacheReadTokens)*model.cached*tier.priceScale*2, cost.CacheReadCost, 1e-12)
require.InDelta(t, float64(tokens.OutputTokens)*model.output*tier.priceScale*1.5, cost.OutputCost, 1e-12)
})
}
}
}
func TestBillingService_GPT56LongContextBoundaryIsExclusive(t *testing.T) {
svc := NewBillingService(&config.Config{}, nil)
tokens := UsageTokens{InputTokens: 100000, CacheCreationTokens: 100000, CacheReadTokens: 72000, OutputTokens: 10}
cost, err := svc.CalculateCost("gpt-5.6-sol", tokens, 1)
require.NoError(t, err)
require.InDelta(t, 100000*5e-6, cost.InputCost, 1e-12)
require.InDelta(t, 100000*6.25e-6, cost.CacheCreationCost, 1e-12)
require.InDelta(t, 72000*0.5e-6, cost.CacheReadCost, 1e-12)
require.InDelta(t, 10*30e-6, cost.OutputCost, 1e-12)
}
func TestPricingService_BareGPT56AliasDeterministicallyUsesSol(t *testing.T) {
pricingSvc := &PricingService{pricingData: map[string]*LiteLLMModelPricing{
"gpt-5.6-sol": {InputCostPerToken: 5e-6},
"gpt-5.6-terra": {InputCostPerToken: 2.5e-6},
"gpt-5.6-luna": {InputCostPerToken: 1e-6},
"gpt-5.4": {InputCostPerToken: 2.5e-6},
}}
for i := 0; i < 100; i++ {
for _, alias := range []string{"gpt-5.6", "openai/gpt-5.6"} {
pricing := pricingSvc.GetModelPricing(alias)
require.NotNil(t, pricing)
require.InDelta(t, 5e-6, pricing.InputCostPerToken, 1e-12, "iteration=%d alias=%s", i, alias)
}
}
billingSvc := NewBillingService(&config.Config{}, pricingSvc)
for _, alias := range []string{"gpt-5.6", "openai/gpt-5.6"} {
pricing, err := billingSvc.GetModelPricing(alias)
require.NoError(t, err)
require.InDelta(t, 5e-6, pricing.InputPricePerToken, 1e-12)
require.InDelta(t, 6.25e-6, pricing.CacheCreationPricePerToken, 1e-12)
}
}
func TestDefaultPricingIncludesOfficialGPT56Rates(t *testing.T) {
data, err := os.ReadFile(filepath.Join("..", "..", "resources", "model-pricing", "model_prices_and_context_window.json"))
require.NoError(t, err)
pricingSvc := &PricingService{}
pricingData, err := pricingSvc.parsePricingData(data)
require.NoError(t, err)
pricingSvc.pricingData = pricingData
billingSvc := NewBillingService(&config.Config{}, pricingSvc)
tests := []struct {
model string
input, cached, cacheWrite, output float64
inputPriority, cachedPriority, cacheWritePriority, outputPriority float64
}{
{model: "gpt-5.6-sol", input: 5e-6, cached: 0.5e-6, cacheWrite: 6.25e-6, output: 30e-6, inputPriority: 10e-6, cachedPriority: 1e-6, cacheWritePriority: 12.5e-6, outputPriority: 60e-6},
{model: "gpt-5.6-terra", input: 2.5e-6, cached: 0.25e-6, cacheWrite: 3.125e-6, output: 15e-6, inputPriority: 5e-6, cachedPriority: 0.5e-6, cacheWritePriority: 6.25e-6, outputPriority: 30e-6},
{model: "gpt-5.6-luna", input: 1e-6, cached: 0.1e-6, cacheWrite: 1.25e-6, output: 6e-6, inputPriority: 2e-6, cachedPriority: 0.2e-6, cacheWritePriority: 2.5e-6, outputPriority: 12e-6},
}
for _, tt := range tests {
t.Run(tt.model, func(t *testing.T) {
pricing, err := billingSvc.GetModelPricing(tt.model)
require.NoError(t, err)
require.InDelta(t, tt.input, pricing.InputPricePerToken, 1e-12)
require.InDelta(t, tt.cached, pricing.CacheReadPricePerToken, 1e-12)
require.InDelta(t, tt.cacheWrite, pricing.CacheCreationPricePerToken, 1e-12)
require.InDelta(t, tt.output, pricing.OutputPricePerToken, 1e-12)
require.InDelta(t, tt.inputPriority, pricing.InputPricePerTokenPriority, 1e-12)
require.InDelta(t, tt.cachedPriority, pricing.CacheReadPricePerTokenPriority, 1e-12)
require.InDelta(t, tt.cacheWritePriority, pricing.CacheCreationPricePerTokenPriority, 1e-12)
require.InDelta(t, tt.outputPriority, pricing.OutputPricePerTokenPriority, 1e-12)
require.Equal(t, 272000, pricing.LongContextInputThreshold)
require.InDelta(t, 2.0, pricing.LongContextInputMultiplier, 1e-12)
require.InDelta(t, 1.5, pricing.LongContextOutputMultiplier, 1e-12)
})
}
}
func TestGPT56DedicatedFallbacksUseOfficialRates(t *testing.T) {
tests := []struct {
model string
input, cached, cacheWrite, output float64
}{
{model: "gpt-5.6-sol", input: 5e-6, cached: 0.5e-6, cacheWrite: 6.25e-6, output: 30e-6},
{model: "gpt-5.6-terra", input: 2.5e-6, cached: 0.25e-6, cacheWrite: 3.125e-6, output: 15e-6},
{model: "gpt-5.6-luna", input: 1e-6, cached: 0.1e-6, cacheWrite: 1.25e-6, output: 6e-6},
}
for _, tt := range tests {
t.Run(tt.model+"/pricing_service", func(t *testing.T) {
pricingSvc := &PricingService{pricingData: map[string]*LiteLLMModelPricing{
"gpt-5.1-codex": {InputCostPerToken: 1.25e-6},
}}
svc := NewBillingService(&config.Config{}, pricingSvc)
pricing, err := svc.GetModelPricing(tt.model + "-preview")
require.NoError(t, err)
assertGPT56FallbackPricing(t, pricing, tt.input, tt.cached, tt.cacheWrite, tt.output)
})
t.Run(tt.model+"/billing_service", func(t *testing.T) {
svc := NewBillingService(&config.Config{}, nil)
pricing, err := svc.GetModelPricing(tt.model)
require.NoError(t, err)
assertGPT56FallbackPricing(t, pricing, tt.input, tt.cached, tt.cacheWrite, tt.output)
})
}
}
func assertGPT56FallbackPricing(t *testing.T, pricing *ModelPricing, input, cached, cacheWrite, output float64) {
t.Helper()
require.InDelta(t, input, pricing.InputPricePerToken, 1e-12)
require.InDelta(t, cached, pricing.CacheReadPricePerToken, 1e-12)
require.InDelta(t, cacheWrite, pricing.CacheCreationPricePerToken, 1e-12)
require.InDelta(t, output, pricing.OutputPricePerToken, 1e-12)
require.Equal(t, 272000, pricing.LongContextInputThreshold)
require.InDelta(t, 2.0, pricing.LongContextInputMultiplier, 1e-12)
require.InDelta(t, 1.5, pricing.LongContextOutputMultiplier, 1e-12)
}
func TestParsePricingData_KeepsImageOnlyPricing(t *testing.T) {
svc := &PricingService{}
body := []byte(`{
+20 -11
View File
@@ -135,6 +135,7 @@ type RedeemCodeBatchUpdateResult struct {
type RedeemService struct {
redeemRepo RedeemCodeRepository
userRepo UserRepository
redeemUserRepo RedeemUserAdjustmentRepository
subscriptionService *SubscriptionService
cache RedeemCache
billingCacheService *BillingCacheService
@@ -154,9 +155,11 @@ func NewRedeemService(
authCacheInvalidator APIKeyAuthCacheInvalidator,
affiliateService *AffiliateService,
) *RedeemService {
redeemUserRepo, _ := userRepo.(RedeemUserAdjustmentRepository)
return &RedeemService{
redeemRepo: redeemRepo,
userRepo: userRepo,
redeemUserRepo: redeemUserRepo,
subscriptionService: subscriptionService,
cache: cache,
billingCacheService: billingCacheService,
@@ -426,7 +429,7 @@ func (s *RedeemService) Redeem(ctx context.Context, userID int64, code string) (
}
// 获取用户信息
user, err := s.userRepo.GetByID(ctx, userID)
_, err = s.userRepo.GetByID(ctx, userID)
if err != nil {
return nil, fmt.Errorf("get user: %w", err)
}
@@ -454,21 +457,27 @@ func (s *RedeemService) Redeem(ctx context.Context, userID int64, code string) (
switch redeemCode.Type {
case RedeemTypeBalance:
amount := redeemCode.Value
// 负数为退款扣减,余额最低为 0
if amount < 0 && user.Balance+amount < 0 {
amount = -user.Balance
}
if err := s.userRepo.UpdateBalance(txCtx, userID, amount); err != nil {
if amount < 0 {
if s.redeemUserRepo == nil {
return nil, errors.New("user repository does not support atomic redeem balance adjustments")
}
if err := s.redeemUserRepo.ApplyRedeemBalanceAdjustment(txCtx, userID, amount); err != nil {
return nil, fmt.Errorf("update user balance: %w", err)
}
} else if err := s.userRepo.UpdateBalance(txCtx, userID, amount); err != nil {
return nil, fmt.Errorf("update user balance: %w", err)
}
case RedeemTypeConcurrency:
delta := int(redeemCode.Value)
// 负数为退款扣减,并发数最低为 0
if delta < 0 && user.Concurrency+delta < 0 {
delta = -user.Concurrency
}
if err := s.userRepo.UpdateConcurrency(txCtx, userID, delta); err != nil {
if delta < 0 {
if s.redeemUserRepo == nil {
return nil, errors.New("user repository does not support atomic redeem concurrency adjustments")
}
if err := s.redeemUserRepo.ApplyRedeemConcurrencyAdjustment(txCtx, userID, delta); err != nil {
return nil, fmt.Errorf("update user concurrency: %w", err)
}
} else if err := s.userRepo.UpdateConcurrency(txCtx, userID, delta); err != nil {
return nil, fmt.Errorf("update user concurrency: %w", err)
}
@@ -801,6 +801,16 @@ func (s *SettingService) SetOpenAIFastPolicySettings(ctx context.Context, settin
if !validScopes[rule.Scope] {
return fmt.Errorf("rule[%d]: invalid scope %q", i, rule.Scope)
}
seenUserIDs := make(map[int64]struct{}, len(rule.UserIDs))
for j, userID := range rule.UserIDs {
if userID <= 0 {
return fmt.Errorf("rule[%d]: user_ids[%d] must be positive", i, j)
}
if _, exists := seenUserIDs[userID]; exists {
return fmt.Errorf("rule[%d]: user_ids[%d] duplicates user_id %d", i, j, userID)
}
seenUserIDs[userID] = struct{}{}
}
for j, pattern := range rule.ModelWhitelist {
trimmed := strings.TrimSpace(pattern)
if trimmed == "" {
@@ -83,7 +83,7 @@ const antigravityUserAgentVersionErrorTTL = 5 * time.Second
const antigravityUserAgentVersionDBTimeout = 5 * time.Second
// DefaultOpenAICodexUserAgent OpenAI Codex 默认 User-Agent(用于规避 Cloudflare 对浏览器 UA 的质询)
const DefaultOpenAICodexUserAgent = "codex-tui/0.125.0 (Ubuntu 22.4.0; x86_64) xterm-256color (codex-tui; 0.125.0)"
const DefaultOpenAICodexUserAgent = "codex-tui/0.144.1 (Ubuntu 22.4.0; x86_64) xterm-256color (codex-tui; 0.144.1)"
// cachedOpenAICodexUserAgent 缓存 OpenAI Codex UA(进程内缓存,60s TTL)
type cachedOpenAICodexUserAgent struct {
@@ -586,6 +586,7 @@ type OpenAIFastPolicyRule struct {
ServiceTier string `json:"service_tier"` // "priority" | "flex" | "auto" | "default" | "scale" | "all"
Action string `json:"action"` // "pass" | "filter" | "block" | "force_priority"
Scope string `json:"scope"` // "all" | "oauth" | "apikey" | "bedrock"
UserIDs []int64 `json:"user_ids,omitempty"` // 空=所有 Sub2API 用户;非空=仅指定 API Key 所属用户
ErrorMessage string `json:"error_message,omitempty"` // 自定义错误消息 (action=block 时生效)
ModelWhitelist []string `json:"model_whitelist,omitempty"` // 模型匹配模式列表(为空=对所有模型生效)
FallbackAction string `json:"fallback_action,omitempty"` // 未匹配白名单的模型的处理方式
@@ -6,11 +6,49 @@ import (
"testing"
"time"
dbent "github.com/Wei-Shaw/sub2api/ent"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/dgraph-io/ristretto"
"github.com/stretchr/testify/require"
)
func TestWithSubscriptionUpdateTx_ReusesExistingTransaction(t *testing.T) {
existingTx := &dbent.Tx{}
ctx := dbent.NewTxContext(context.Background(), existingTx)
svc := &SubscriptionService{entClient: &dbent.Client{}}
called := false
err := svc.withSubscriptionUpdateTx(ctx, func(txCtx context.Context) error {
called = true
require.Same(t, existingTx, dbent.TxFromContext(txCtx))
return nil
})
require.NoError(t, err)
require.True(t, called)
}
func TestMaybeInvalidateAssignmentCaches_DefersForOuterTransactionOwner(t *testing.T) {
cache, err := ristretto.NewCache(&ristretto.Config{NumCounters: 1_000, MaxCost: 100, BufferItems: 64})
require.NoError(t, err)
t.Cleanup(cache.Close)
svc := &SubscriptionService{subCacheL1: cache}
key := subCacheKey(7, 9)
require.True(t, cache.Set(key, &UserSubscription{ID: 42}, 1))
cache.Wait()
svc.maybeInvalidateAssignmentCaches(7, 9, true)
_, cachedBeforeCommit := cache.Get(key)
require.True(t, cachedBeforeCommit, "outer transaction must retain caches until its owner commits")
svc.maybeInvalidateAssignmentCaches(7, 9, false)
cache.Wait()
_, cachedAfterCommit := cache.Get(key)
require.False(t, cachedAfterCommit, "post-commit invalidation must remove the cached subscription")
}
type groupRepoNoop struct{}
func (groupRepoNoop) Create(context.Context, *Group) error { panic("unexpected Create call") }
@@ -119,13 +157,16 @@ func (userSubRepoNoop) UpdateNotes(context.Context, int64, string) error {
func (userSubRepoNoop) ActivateWindows(context.Context, int64, time.Time) error {
panic("unexpected ActivateWindows call")
}
func (userSubRepoNoop) ResetDailyUsage(context.Context, int64, time.Time) error {
func (userSubRepoNoop) ResetUsageWindows(context.Context, int64, bool, bool, bool, time.Time) error {
panic("unexpected ResetUsageWindows call")
}
func (userSubRepoNoop) ResetDailyUsage(context.Context, int64, *time.Time, time.Time) error {
panic("unexpected ResetDailyUsage call")
}
func (userSubRepoNoop) ResetWeeklyUsage(context.Context, int64, time.Time) error {
func (userSubRepoNoop) ResetWeeklyUsage(context.Context, int64, *time.Time, time.Time) error {
panic("unexpected ResetWeeklyUsage call")
}
func (userSubRepoNoop) ResetMonthlyUsage(context.Context, int64, time.Time) error {
func (userSubRepoNoop) ResetMonthlyUsage(context.Context, int64, *time.Time, time.Time) error {
panic("unexpected ResetMonthlyUsage call")
}
func (userSubRepoNoop) IncrementUsage(context.Context, int64, float64) error {
@@ -87,15 +87,19 @@ func (r *subscriptionExpiryRepoStub) ActivateWindows(context.Context, int64, tim
return nil
}
func (r *subscriptionExpiryRepoStub) ResetDailyUsage(context.Context, int64, time.Time) error {
func (r *subscriptionExpiryRepoStub) ResetUsageWindows(context.Context, int64, bool, bool, bool, time.Time) error {
return nil
}
func (r *subscriptionExpiryRepoStub) ResetWeeklyUsage(context.Context, int64, time.Time) error {
func (r *subscriptionExpiryRepoStub) ResetDailyUsage(context.Context, int64, *time.Time, time.Time) error {
return nil
}
func (r *subscriptionExpiryRepoStub) ResetMonthlyUsage(context.Context, int64, time.Time) error {
func (r *subscriptionExpiryRepoStub) ResetWeeklyUsage(context.Context, int64, *time.Time, time.Time) error {
return nil
}
func (r *subscriptionExpiryRepoStub) ResetMonthlyUsage(context.Context, int64, *time.Time, time.Time) error {
return nil
}
@@ -11,7 +11,7 @@ import (
"github.com/stretchr/testify/require"
)
// resetQuotaUserSubRepoStub 支持 GetByID、ResetDailyUsage、ResetWeeklyUsage、ResetMonthlyUsage,
// resetQuotaUserSubRepoStub 支持 GetByID、ResetUsageWindows,
// 其余方法继承 userSubRepoNoop(panic)。
type resetQuotaUserSubRepoStub struct {
userSubRepoNoop
@@ -34,7 +34,38 @@ func (r *resetQuotaUserSubRepoStub) GetByID(_ context.Context, id int64) (*UserS
return &cp, nil
}
func (r *resetQuotaUserSubRepoStub) ResetDailyUsage(_ context.Context, _ int64, windowStart time.Time) error {
func (r *resetQuotaUserSubRepoStub) ResetUsageWindows(_ context.Context, _ int64, resetDaily, resetWeekly, resetMonthly bool, windowStart time.Time) error {
r.resetDailyCalled = resetDaily
r.resetWeeklyCalled = resetWeekly
r.resetMonthlyCalled = resetMonthly
if resetDaily && r.resetDailyErr != nil {
return r.resetDailyErr
}
if resetWeekly && r.resetWeeklyErr != nil {
return r.resetWeeklyErr
}
if resetMonthly && r.resetMonthlyErr != nil {
return r.resetMonthlyErr
}
if r.sub == nil {
return nil
}
if resetDaily {
r.sub.DailyUsageUSD = 0
r.sub.DailyWindowStart = &windowStart
}
if resetWeekly {
r.sub.WeeklyUsageUSD = 0
r.sub.WeeklyWindowStart = &windowStart
}
if resetMonthly {
r.sub.MonthlyUsageUSD = 0
r.sub.MonthlyWindowStart = &windowStart
}
return nil
}
func (r *resetQuotaUserSubRepoStub) ResetDailyUsage(_ context.Context, _ int64, _ *time.Time, windowStart time.Time) error {
r.resetDailyCalled = true
if r.resetDailyErr == nil && r.sub != nil {
r.sub.DailyUsageUSD = 0
@@ -43,12 +74,12 @@ func (r *resetQuotaUserSubRepoStub) ResetDailyUsage(_ context.Context, _ int64,
return r.resetDailyErr
}
func (r *resetQuotaUserSubRepoStub) ResetWeeklyUsage(_ context.Context, _ int64, _ time.Time) error {
func (r *resetQuotaUserSubRepoStub) ResetWeeklyUsage(_ context.Context, _ int64, _ *time.Time, _ time.Time) error {
r.resetWeeklyCalled = true
return r.resetWeeklyErr
}
func (r *resetQuotaUserSubRepoStub) ResetMonthlyUsage(_ context.Context, _ int64, _ time.Time) error {
func (r *resetQuotaUserSubRepoStub) ResetMonthlyUsage(_ context.Context, _ int64, _ *time.Time, _ time.Time) error {
r.resetMonthlyCalled = true
return r.resetMonthlyErr
}
@@ -140,7 +171,7 @@ func TestAdminResetQuota_ResetDailyUsageError(t *testing.T) {
require.ErrorIs(t, err, dbErr)
require.True(t, stub.resetDailyCalled)
require.False(t, stub.resetWeeklyCalled, "daily 失败后不应继续调用 weekly")
require.True(t, stub.resetWeeklyCalled, "原子重置应在一次调用中提交所选窗口")
}
func TestAdminResetQuota_ResetWeeklyUsageError(t *testing.T) {
@@ -200,7 +231,7 @@ func TestAdminResetQuota_ReturnsRefreshedSub(t *testing.T) {
result, err := svc.AdminResetQuota(context.Background(), 6, true, false, false)
require.NoError(t, err)
// ResetDailyUsage stub 会将 sub.DailyUsageUSD 归零,
// ResetUsageWindows stub 会将 sub.DailyUsageUSD 归零,
// 服务应返回第二次 GetByID 的刷新值而非初始的 99.9
require.Equal(t, float64(0), result.DailyUsageUSD, "返回的订阅应反映已归零的用量")
require.True(t, stub.resetDailyCalled)

Some files were not shown because too many files have changed in this diff Show More