mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-24 16:05:44 +08:00
Merge remote-tracking branch 'upstream/main'
This commit is contained in:
@@ -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 @@
|
||||
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)
|
||||
}
|
||||
|
||||
@@ -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: ¶llel,
|
||||
}
|
||||
|
||||
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,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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"`
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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') <> ''")
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 = ¤t
|
||||
fresh.WeeklyWindowStart = ¤t
|
||||
fresh.MonthlyWindowStart = ¤t
|
||||
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 := ""
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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(`{
|
||||
|
||||
@@ -135,6 +135,7 @@ type RedeemCodeBatchUpdateResult struct {
|
||||
type RedeemService struct {
|
||||
redeemRepo RedeemCodeRepository
|
||||
userRepo UserRepository
|
||||
redeemUserRepo RedeemUserAdjustmentRepository
|
||||
subscriptionService *SubscriptionService
|
||||
cache RedeemCache
|
||||
billingCacheService *BillingCacheService
|
||||
@@ -154,9 +155,11 @@ func NewRedeemService(
|
||||
authCacheInvalidator APIKeyAuthCacheInvalidator,
|
||||
affiliateService *AffiliateService,
|
||||
) *RedeemService {
|
||||
redeemUserRepo, _ := userRepo.(RedeemUserAdjustmentRepository)
|
||||
return &RedeemService{
|
||||
redeemRepo: redeemRepo,
|
||||
userRepo: userRepo,
|
||||
redeemUserRepo: redeemUserRepo,
|
||||
subscriptionService: subscriptionService,
|
||||
cache: cache,
|
||||
billingCacheService: billingCacheService,
|
||||
@@ -426,7 +429,7 @@ func (s *RedeemService) Redeem(ctx context.Context, userID int64, code string) (
|
||||
}
|
||||
|
||||
// 获取用户信息
|
||||
user, err := s.userRepo.GetByID(ctx, userID)
|
||||
_, err = s.userRepo.GetByID(ctx, userID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get user: %w", err)
|
||||
}
|
||||
@@ -454,21 +457,27 @@ func (s *RedeemService) Redeem(ctx context.Context, userID int64, code string) (
|
||||
switch redeemCode.Type {
|
||||
case RedeemTypeBalance:
|
||||
amount := redeemCode.Value
|
||||
// 负数为退款扣减,余额最低为 0
|
||||
if amount < 0 && user.Balance+amount < 0 {
|
||||
amount = -user.Balance
|
||||
}
|
||||
if err := s.userRepo.UpdateBalance(txCtx, userID, amount); err != nil {
|
||||
if amount < 0 {
|
||||
if s.redeemUserRepo == nil {
|
||||
return nil, errors.New("user repository does not support atomic redeem balance adjustments")
|
||||
}
|
||||
if err := s.redeemUserRepo.ApplyRedeemBalanceAdjustment(txCtx, userID, amount); err != nil {
|
||||
return nil, fmt.Errorf("update user balance: %w", err)
|
||||
}
|
||||
} else if err := s.userRepo.UpdateBalance(txCtx, userID, amount); err != nil {
|
||||
return nil, fmt.Errorf("update user balance: %w", err)
|
||||
}
|
||||
|
||||
case RedeemTypeConcurrency:
|
||||
delta := int(redeemCode.Value)
|
||||
// 负数为退款扣减,并发数最低为 0
|
||||
if delta < 0 && user.Concurrency+delta < 0 {
|
||||
delta = -user.Concurrency
|
||||
}
|
||||
if err := s.userRepo.UpdateConcurrency(txCtx, userID, delta); err != nil {
|
||||
if delta < 0 {
|
||||
if s.redeemUserRepo == nil {
|
||||
return nil, errors.New("user repository does not support atomic redeem concurrency adjustments")
|
||||
}
|
||||
if err := s.redeemUserRepo.ApplyRedeemConcurrencyAdjustment(txCtx, userID, delta); err != nil {
|
||||
return nil, fmt.Errorf("update user concurrency: %w", err)
|
||||
}
|
||||
} else if err := s.userRepo.UpdateConcurrency(txCtx, userID, delta); err != nil {
|
||||
return nil, fmt.Errorf("update user concurrency: %w", err)
|
||||
}
|
||||
|
||||
|
||||
@@ -801,6 +801,16 @@ func (s *SettingService) SetOpenAIFastPolicySettings(ctx context.Context, settin
|
||||
if !validScopes[rule.Scope] {
|
||||
return fmt.Errorf("rule[%d]: invalid scope %q", i, rule.Scope)
|
||||
}
|
||||
seenUserIDs := make(map[int64]struct{}, len(rule.UserIDs))
|
||||
for j, userID := range rule.UserIDs {
|
||||
if userID <= 0 {
|
||||
return fmt.Errorf("rule[%d]: user_ids[%d] must be positive", i, j)
|
||||
}
|
||||
if _, exists := seenUserIDs[userID]; exists {
|
||||
return fmt.Errorf("rule[%d]: user_ids[%d] duplicates user_id %d", i, j, userID)
|
||||
}
|
||||
seenUserIDs[userID] = struct{}{}
|
||||
}
|
||||
for j, pattern := range rule.ModelWhitelist {
|
||||
trimmed := strings.TrimSpace(pattern)
|
||||
if trimmed == "" {
|
||||
|
||||
@@ -83,7 +83,7 @@ const antigravityUserAgentVersionErrorTTL = 5 * time.Second
|
||||
const antigravityUserAgentVersionDBTimeout = 5 * time.Second
|
||||
|
||||
// DefaultOpenAICodexUserAgent OpenAI Codex 默认 User-Agent(用于规避 Cloudflare 对浏览器 UA 的质询)
|
||||
const DefaultOpenAICodexUserAgent = "codex-tui/0.125.0 (Ubuntu 22.4.0; x86_64) xterm-256color (codex-tui; 0.125.0)"
|
||||
const DefaultOpenAICodexUserAgent = "codex-tui/0.144.1 (Ubuntu 22.4.0; x86_64) xterm-256color (codex-tui; 0.144.1)"
|
||||
|
||||
// cachedOpenAICodexUserAgent 缓存 OpenAI Codex UA(进程内缓存,60s TTL)
|
||||
type cachedOpenAICodexUserAgent struct {
|
||||
|
||||
@@ -586,6 +586,7 @@ type OpenAIFastPolicyRule struct {
|
||||
ServiceTier string `json:"service_tier"` // "priority" | "flex" | "auto" | "default" | "scale" | "all"
|
||||
Action string `json:"action"` // "pass" | "filter" | "block" | "force_priority"
|
||||
Scope string `json:"scope"` // "all" | "oauth" | "apikey" | "bedrock"
|
||||
UserIDs []int64 `json:"user_ids,omitempty"` // 空=所有 Sub2API 用户;非空=仅指定 API Key 所属用户
|
||||
ErrorMessage string `json:"error_message,omitempty"` // 自定义错误消息 (action=block 时生效)
|
||||
ModelWhitelist []string `json:"model_whitelist,omitempty"` // 模型匹配模式列表(为空=对所有模型生效)
|
||||
FallbackAction string `json:"fallback_action,omitempty"` // 未匹配白名单的模型的处理方式
|
||||
|
||||
@@ -6,11 +6,49 @@ import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
"github.com/dgraph-io/ristretto"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestWithSubscriptionUpdateTx_ReusesExistingTransaction(t *testing.T) {
|
||||
existingTx := &dbent.Tx{}
|
||||
ctx := dbent.NewTxContext(context.Background(), existingTx)
|
||||
svc := &SubscriptionService{entClient: &dbent.Client{}}
|
||||
|
||||
called := false
|
||||
err := svc.withSubscriptionUpdateTx(ctx, func(txCtx context.Context) error {
|
||||
called = true
|
||||
require.Same(t, existingTx, dbent.TxFromContext(txCtx))
|
||||
return nil
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.True(t, called)
|
||||
}
|
||||
|
||||
func TestMaybeInvalidateAssignmentCaches_DefersForOuterTransactionOwner(t *testing.T) {
|
||||
cache, err := ristretto.NewCache(&ristretto.Config{NumCounters: 1_000, MaxCost: 100, BufferItems: 64})
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(cache.Close)
|
||||
|
||||
svc := &SubscriptionService{subCacheL1: cache}
|
||||
key := subCacheKey(7, 9)
|
||||
require.True(t, cache.Set(key, &UserSubscription{ID: 42}, 1))
|
||||
cache.Wait()
|
||||
|
||||
svc.maybeInvalidateAssignmentCaches(7, 9, true)
|
||||
_, cachedBeforeCommit := cache.Get(key)
|
||||
require.True(t, cachedBeforeCommit, "outer transaction must retain caches until its owner commits")
|
||||
|
||||
svc.maybeInvalidateAssignmentCaches(7, 9, false)
|
||||
cache.Wait()
|
||||
_, cachedAfterCommit := cache.Get(key)
|
||||
require.False(t, cachedAfterCommit, "post-commit invalidation must remove the cached subscription")
|
||||
}
|
||||
|
||||
type groupRepoNoop struct{}
|
||||
|
||||
func (groupRepoNoop) Create(context.Context, *Group) error { panic("unexpected Create call") }
|
||||
@@ -119,13 +157,16 @@ func (userSubRepoNoop) UpdateNotes(context.Context, int64, string) error {
|
||||
func (userSubRepoNoop) ActivateWindows(context.Context, int64, time.Time) error {
|
||||
panic("unexpected ActivateWindows call")
|
||||
}
|
||||
func (userSubRepoNoop) ResetDailyUsage(context.Context, int64, time.Time) error {
|
||||
func (userSubRepoNoop) ResetUsageWindows(context.Context, int64, bool, bool, bool, time.Time) error {
|
||||
panic("unexpected ResetUsageWindows call")
|
||||
}
|
||||
func (userSubRepoNoop) ResetDailyUsage(context.Context, int64, *time.Time, time.Time) error {
|
||||
panic("unexpected ResetDailyUsage call")
|
||||
}
|
||||
func (userSubRepoNoop) ResetWeeklyUsage(context.Context, int64, time.Time) error {
|
||||
func (userSubRepoNoop) ResetWeeklyUsage(context.Context, int64, *time.Time, time.Time) error {
|
||||
panic("unexpected ResetWeeklyUsage call")
|
||||
}
|
||||
func (userSubRepoNoop) ResetMonthlyUsage(context.Context, int64, time.Time) error {
|
||||
func (userSubRepoNoop) ResetMonthlyUsage(context.Context, int64, *time.Time, time.Time) error {
|
||||
panic("unexpected ResetMonthlyUsage call")
|
||||
}
|
||||
func (userSubRepoNoop) IncrementUsage(context.Context, int64, float64) error {
|
||||
|
||||
@@ -87,15 +87,19 @@ func (r *subscriptionExpiryRepoStub) ActivateWindows(context.Context, int64, tim
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *subscriptionExpiryRepoStub) ResetDailyUsage(context.Context, int64, time.Time) error {
|
||||
func (r *subscriptionExpiryRepoStub) ResetUsageWindows(context.Context, int64, bool, bool, bool, time.Time) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *subscriptionExpiryRepoStub) ResetWeeklyUsage(context.Context, int64, time.Time) error {
|
||||
func (r *subscriptionExpiryRepoStub) ResetDailyUsage(context.Context, int64, *time.Time, time.Time) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *subscriptionExpiryRepoStub) ResetMonthlyUsage(context.Context, int64, time.Time) error {
|
||||
func (r *subscriptionExpiryRepoStub) ResetWeeklyUsage(context.Context, int64, *time.Time, time.Time) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *subscriptionExpiryRepoStub) ResetMonthlyUsage(context.Context, int64, *time.Time, time.Time) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
@@ -11,7 +11,7 @@ import (
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// resetQuotaUserSubRepoStub 支持 GetByID、ResetDailyUsage、ResetWeeklyUsage、ResetMonthlyUsage,
|
||||
// resetQuotaUserSubRepoStub 支持 GetByID、ResetUsageWindows,
|
||||
// 其余方法继承 userSubRepoNoop(panic)。
|
||||
type resetQuotaUserSubRepoStub struct {
|
||||
userSubRepoNoop
|
||||
@@ -34,7 +34,38 @@ func (r *resetQuotaUserSubRepoStub) GetByID(_ context.Context, id int64) (*UserS
|
||||
return &cp, nil
|
||||
}
|
||||
|
||||
func (r *resetQuotaUserSubRepoStub) ResetDailyUsage(_ context.Context, _ int64, windowStart time.Time) error {
|
||||
func (r *resetQuotaUserSubRepoStub) ResetUsageWindows(_ context.Context, _ int64, resetDaily, resetWeekly, resetMonthly bool, windowStart time.Time) error {
|
||||
r.resetDailyCalled = resetDaily
|
||||
r.resetWeeklyCalled = resetWeekly
|
||||
r.resetMonthlyCalled = resetMonthly
|
||||
if resetDaily && r.resetDailyErr != nil {
|
||||
return r.resetDailyErr
|
||||
}
|
||||
if resetWeekly && r.resetWeeklyErr != nil {
|
||||
return r.resetWeeklyErr
|
||||
}
|
||||
if resetMonthly && r.resetMonthlyErr != nil {
|
||||
return r.resetMonthlyErr
|
||||
}
|
||||
if r.sub == nil {
|
||||
return nil
|
||||
}
|
||||
if resetDaily {
|
||||
r.sub.DailyUsageUSD = 0
|
||||
r.sub.DailyWindowStart = &windowStart
|
||||
}
|
||||
if resetWeekly {
|
||||
r.sub.WeeklyUsageUSD = 0
|
||||
r.sub.WeeklyWindowStart = &windowStart
|
||||
}
|
||||
if resetMonthly {
|
||||
r.sub.MonthlyUsageUSD = 0
|
||||
r.sub.MonthlyWindowStart = &windowStart
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *resetQuotaUserSubRepoStub) ResetDailyUsage(_ context.Context, _ int64, _ *time.Time, windowStart time.Time) error {
|
||||
r.resetDailyCalled = true
|
||||
if r.resetDailyErr == nil && r.sub != nil {
|
||||
r.sub.DailyUsageUSD = 0
|
||||
@@ -43,12 +74,12 @@ func (r *resetQuotaUserSubRepoStub) ResetDailyUsage(_ context.Context, _ int64,
|
||||
return r.resetDailyErr
|
||||
}
|
||||
|
||||
func (r *resetQuotaUserSubRepoStub) ResetWeeklyUsage(_ context.Context, _ int64, _ time.Time) error {
|
||||
func (r *resetQuotaUserSubRepoStub) ResetWeeklyUsage(_ context.Context, _ int64, _ *time.Time, _ time.Time) error {
|
||||
r.resetWeeklyCalled = true
|
||||
return r.resetWeeklyErr
|
||||
}
|
||||
|
||||
func (r *resetQuotaUserSubRepoStub) ResetMonthlyUsage(_ context.Context, _ int64, _ time.Time) error {
|
||||
func (r *resetQuotaUserSubRepoStub) ResetMonthlyUsage(_ context.Context, _ int64, _ *time.Time, _ time.Time) error {
|
||||
r.resetMonthlyCalled = true
|
||||
return r.resetMonthlyErr
|
||||
}
|
||||
@@ -140,7 +171,7 @@ func TestAdminResetQuota_ResetDailyUsageError(t *testing.T) {
|
||||
|
||||
require.ErrorIs(t, err, dbErr)
|
||||
require.True(t, stub.resetDailyCalled)
|
||||
require.False(t, stub.resetWeeklyCalled, "daily 失败后不应继续调用 weekly")
|
||||
require.True(t, stub.resetWeeklyCalled, "原子重置应在一次调用中提交所选窗口")
|
||||
}
|
||||
|
||||
func TestAdminResetQuota_ResetWeeklyUsageError(t *testing.T) {
|
||||
@@ -200,7 +231,7 @@ func TestAdminResetQuota_ReturnsRefreshedSub(t *testing.T) {
|
||||
result, err := svc.AdminResetQuota(context.Background(), 6, true, false, false)
|
||||
|
||||
require.NoError(t, err)
|
||||
// ResetDailyUsage stub 会将 sub.DailyUsageUSD 归零,
|
||||
// ResetUsageWindows stub 会将 sub.DailyUsageUSD 归零,
|
||||
// 服务应返回第二次 GetByID 的刷新值而非初始的 99.9
|
||||
require.Equal(t, float64(0), result.DailyUsageUSD, "返回的订阅应反映已归零的用量")
|
||||
require.True(t, stub.resetDailyCalled)
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user