Merge remote-tracking branch 'upstream/main'

This commit is contained in:
xlplbo committed 2026-07-10 20:46:38 +08:00
commit eb4d005031
124 files changed
+5597 -577

No files matched your search

+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)
Loaded 100 of 124 files, more files were not shown because too many files have changed in this diff. Show more