From 8f28b01c23b109a1f3cec07a67ddb58cd67619b0 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Sat, 10 Oct 2026 10:35:57 +0800 Subject: [PATCH 01/24] update --- controller/relay_finalize_replay_test.go | 222 +++++++++ controller/relay_metadata_funding_test.go | 175 ++++++++ .../relay_responses_stream_error_test.go | 281 ++++++++++++ controller/relay_stream_error.go | 10 +- controller/relay_subscription_stream_test.go | 251 +++++++++++ dto/billing_usage.go | 26 +- dto/openai_response.go | 30 +- model/subscription.go | 82 +++- model/subscription_reset_billing_test.go | 259 +++++++++++ relay/channel/baidu/relay-baidu.go | 2 +- relay/channel/claude/adaptor.go | 8 +- relay/channel/claude/relay-claude.go | 219 ++++++++- relay/channel/claude/relay_claude_test.go | 147 ++++++ relay/channel/claude/responses_repair_test.go | 104 +++++ .../claude/responses_revalidation_test.go | 60 +++ relay/channel/dify/relay-dify.go | 2 +- relay/channel/gemini/adaptor_responses.go | 37 +- .../channel/gemini/adaptor_responses_test.go | 41 +- relay/channel/gemini/relay-gemini.go | 43 +- .../channel/gemini/relay_gemini_usage_test.go | 140 +++++- relay/channel/gemini/relay_responses.go | 153 ++++++- .../gemini/responses_revalidation_test.go | 213 +++++++++ relay/channel/gemini/stream_function_calls.go | 108 +++++ relay/channel/openai/chat_stream.go | 74 ++- relay/channel/openai/chat_via_responses.go | 101 ++++- relay/channel/openai/empty_completion.go | 8 +- relay/channel/openai/empty_completion_test.go | 12 +- .../file_search_stream_delivery_test.go | 87 ++++ relay/channel/openai/first_result_test.go | 143 ++++++ relay/channel/openai/helper.go | 6 +- relay/channel/openai/relay-openai.go | 2 +- relay/channel/openai/relay_responses.go | 133 ++++-- .../openai/responses_terminal_repair_test.go | 195 ++++++++ relay/channel/openai/responses_usage_test.go | 23 + relay/channel/openai/responses_via_chat.go | 110 ++++- relay/channel/xai/text.go | 2 +- relay/chat_completions_via_responses.go | 2 +- relay/chat_stream_billing_test.go | 56 +++ relay/common/relay_info.go | 23 +- relay/common/stream_first_result.go | 66 +++ relay/common/stream_first_result_test.go | 153 +++++++ relay/common/stream_status.go | 62 ++- relay/common/stream_status_test.go | 32 ++ relay/compatible_handler.go | 3 + relay/file_search_stream_settlement_test.go | 282 ++++++++++++ relay/first_result_adapter_test.go | 93 ++++ relay/helper/common.go | 98 +++- relay/helper/first_result_timeout.go | 18 +- relay/helper/first_result_timeout_test.go | 34 +- relay/helper/responses_stream_state.go | 156 +++++++ relay/helper/responses_stream_state_test.go | 103 +++++ relay/helper/responses_write_test.go | 37 ++ relay/helper/stream_result.go | 1 + relay/helper/stream_scanner.go | 14 +- relay/metadata_funding_settlement_test.go | 246 ++++++++++ relay/request_settlement_replay_test.go | 424 ++++++++++++++++++ ...sponses_custom_tool_eof_settlement_test.go | 308 +++++++++++++ relay/responses_failure_boundary_test.go | 168 +++++++ .../responses_first_result_settlement_test.go | 232 ++++++++++ relay/responses_handler.go | 9 + relay/responses_partial_settlement_test.go | 245 ++++++++++ .../responses_to_chat_tool_settlement_test.go | 194 ++++++++ relay/responses_tools_roundtrip_test.go | 97 ++++ relay/subscription_period_settlement_test.go | 240 ++++++++++ relay/subscription_stream_settlement_test.go | 394 ++++++++++++++++ service/billing_playground_replay_test.go | 84 ++++ service/billing_session.go | 12 + service/log_info_generate.go | 3 + service/log_info_generate_test.go | 13 + service/openai_chat_responses_compat.go | 27 ++ .../chat_to_responses_response.go | 182 ++++++-- .../openaicompat/custom_tool_finalize_test.go | 45 ++ .../openaicompat/custom_tool_response_test.go | 62 +++ service/openaicompat/responses_compat_test.go | 108 ++++- .../openaicompat/responses_request_to_chat.go | 115 ++++- .../responses_revalidation_test.go | 39 ++ service/openaicompat/responses_to_chat.go | 1 + service/quota.go | 3 +- service/responses_partial_quota_test.go | 83 ++++ service/text_quota.go | 19 +- .../smart-ops/model-performance.test.tsx | 49 +- 81 files changed, 7896 insertions(+), 248 deletions(-) create mode 100644 controller/relay_finalize_replay_test.go create mode 100644 controller/relay_metadata_funding_test.go create mode 100644 controller/relay_responses_stream_error_test.go create mode 100644 controller/relay_subscription_stream_test.go create mode 100644 model/subscription_reset_billing_test.go create mode 100644 relay/channel/claude/responses_repair_test.go create mode 100644 relay/channel/claude/responses_revalidation_test.go create mode 100644 relay/channel/gemini/responses_revalidation_test.go create mode 100644 relay/channel/gemini/stream_function_calls.go create mode 100644 relay/channel/openai/file_search_stream_delivery_test.go create mode 100644 relay/channel/openai/first_result_test.go create mode 100644 relay/channel/openai/responses_terminal_repair_test.go create mode 100644 relay/common/stream_first_result.go create mode 100644 relay/common/stream_first_result_test.go create mode 100644 relay/file_search_stream_settlement_test.go create mode 100644 relay/first_result_adapter_test.go create mode 100644 relay/helper/responses_stream_state.go create mode 100644 relay/helper/responses_stream_state_test.go create mode 100644 relay/helper/responses_write_test.go create mode 100644 relay/metadata_funding_settlement_test.go create mode 100644 relay/request_settlement_replay_test.go create mode 100644 relay/responses_custom_tool_eof_settlement_test.go create mode 100644 relay/responses_failure_boundary_test.go create mode 100644 relay/responses_first_result_settlement_test.go create mode 100644 relay/responses_partial_settlement_test.go create mode 100644 relay/responses_to_chat_tool_settlement_test.go create mode 100644 relay/responses_tools_roundtrip_test.go create mode 100644 relay/subscription_period_settlement_test.go create mode 100644 relay/subscription_stream_settlement_test.go create mode 100644 service/billing_playground_replay_test.go create mode 100644 service/openaicompat/custom_tool_finalize_test.go create mode 100644 service/openaicompat/custom_tool_response_test.go create mode 100644 service/openaicompat/responses_revalidation_test.go create mode 100644 service/responses_partial_quota_test.go diff --git a/controller/relay_finalize_replay_test.go b/controller/relay_finalize_replay_test.go new file mode 100644 index 00000000..72e97b2b --- /dev/null +++ b/controller/relay_finalize_replay_test.go @@ -0,0 +1,222 @@ +package controller + +import ( + "crypto/sha256" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/constant" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/model" + "github.com/MAX-API-Next/MAX-API/service" + "github.com/MAX-API-Next/MAX-API/setting/model_setting" + "github.com/MAX-API-Next/MAX-API/setting/ratio_setting" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "gorm.io/gorm" + "gorm.io/gorm/logger" +) + +// Real paid controller requests include admission, protocol errors and the +// deferred refund before workers replay or attempt to replace their finalize. +func TestRelayFinalizeIdentityAndConcurrentReplay(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitHttpClient() + service.InitTokenEncoders() + oldRatios, oldCompletion := ratio_setting.ModelRatio2JSONString(), ratio_setting.CompletionRatio2JSONString() + oldFloor, oldCount, oldRetry, oldEmpty := common.PreConsumedQuota, constant.CountToken, common.RetryTimes, common.EmptyCompletionRetryEnabled + common.PreConsumedQuota, constant.CountToken, common.RetryTimes, common.EmptyCompletionRetryEnabled = 100, false, 2, false + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"gpt-finalize-controller":1}`)) + require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"gpt-finalize-controller":1}`)) + t.Cleanup(func() { + common.PreConsumedQuota, constant.CountToken, common.RetryTimes, common.EmptyCompletionRetryEnabled = oldFloor, oldCount, oldRetry, oldEmpty + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(oldRatios)) + require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(oldCompletion)) + }) + for _, chat := range []bool{false, true} { + for _, funding := range []struct { + name, preference, source string + wallet, subUsed int64 + allow bool + }{ + {"wallet", "wallet_only", "wallet", 1000, 0, false}, + {"subscription", "subscription_only", "subscription", 1000, 0, false}, + {"wallet_fallback", "subscription_first", "wallet", 1000, 1000, true}, + {"subscription_fallback", "wallet_first", "subscription", 0, 0, false}, + } { + for _, usage := range []struct { + name string + prompt, comp int + }{{"negative", 4, 2}, {"zero_delta", 80, 20}, {"positive", 80, 50}, {"zero_quota", 0, 0}} { + t.Run(fmt.Sprintf("chat_%t/%s/%s", chat, funding.name, usage.name), func(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + t.Cleanup(func() { _ = sqlDB.Close() }) + require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}, &model.Channel{}, &model.Log{}, &model.BillingLogReceipt{}, + &model.BillingSettlement{}, &model.BillingPreConsumeSelection{}, &model.CacheInvalidationTask{}, + &model.UserSubscription{}, &model.SubscriptionPreConsumeRecord{}, &model.SubscriptionPlan{})) + oldDB, oldLogDB, oldSQLite := model.DB, model.LOG_DB, common.UsingSQLite + oldRedis, oldBatch, oldConsume := common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled + oldGlobal := *model_setting.GetGlobalSettings() + model.DB, model.LOG_DB, common.UsingSQLite = db, db, true + common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled = false, false, true + *model_setting.GetGlobalSettings() = model_setting.GlobalSettings{} + t.Cleanup(func() { + model.DB, model.LOG_DB, common.UsingSQLite = oldDB, oldLogDB, oldSQLite + common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled = oldRedis, oldBatch, oldConsume + *model_setting.GetGlobalSettings() = oldGlobal + }) + require.NoError(t, db.Create(&model.User{Id: 971, Username: "controller-finalize", Quota: funding.wallet}).Error) + require.NoError(t, db.Create(&model.Token{Id: 972, UserId: 971, Key: "synthetic-controller-finalize", RemainQuota: 1000, Status: common.TokenStatusEnabled}).Error) + require.NoError(t, db.Create(&model.Channel{Id: 973}).Error) + require.NoError(t, db.Create(&model.SubscriptionPlan{Id: 974, Title: "controller-finalize", Enabled: true, TotalAmount: 1000, + AllowWalletOverflow: &funding.allow, QuotaResetPeriod: model.SubscriptionResetNever}).Error) + now := time.Now().Unix() + require.NoError(t, db.Create(&model.UserSubscription{Id: 975, UserId: 971, PlanId: 974, Status: "active", AmountTotal: 1000, + AmountUsed: funding.subUsed, AllowWalletOverflow: funding.allow, StartTime: now - 3600, EndTime: now + 3600, LastResetTime: now - 3600}).Error) + frame := map[string]any{"type": "response.created", "response": map[string]any{"id": "resp_controller_finalize", + "usage": map[string]any{"input_tokens": usage.prompt, "output_tokens": usage.comp, "total_tokens": usage.prompt + usage.comp}}} + output := `{"type":"response.output_text.delta","delta":"partial output"}` + if chat { + frame = map[string]any{"choices": []any{map[string]any{"delta": map[string]any{"content": "partial output"}}}, + "usage": map[string]any{"prompt_tokens": usage.prompt, "completion_tokens": usage.comp, "total_tokens": usage.prompt + usage.comp}} + output = "" + } + encoded, err := common.Marshal(frame) + require.NoError(t, err) + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprintf(w, "data: %s\n\n", encoded) + if output != "" { + _, _ = fmt.Fprintf(w, "data: %s\n\n", output) + } + _, _ = fmt.Fprint(w, "data: {broken\n\n") + })) + t.Cleanup(upstream.Close) + fingerprint := sha256.Sum256([]byte(t.Name())) + requestID := fmt.Sprintf("controller-replay-%x", fingerprint[:16]) + path, format := "/v1/responses", types.RelayFormat(types.RelayFormatOpenAIResponses) + body := `{"model":"gpt-finalize-controller","input":"test","stream":true,"max_output_tokens":16}` + if chat { + path, format = "/v1/chat/completions", types.RelayFormatOpenAI + body = `{"model":"gpt-finalize-controller","messages":[{"role":"user","content":"test"}],"stream":true,"max_tokens":16}` + } + router := gin.New() + router.POST(path, func(c *gin.Context) { + c.Set(common.RequestIdKey, requestID) + c.Set(string(constant.ContextKeyUserId), 971) + c.Set(string(constant.ContextKeyUserQuota), funding.wallet) + c.Set(string(constant.ContextKeyTokenId), 972) + c.Set(string(constant.ContextKeyTokenKey), "synthetic-controller-finalize") + c.Set(string(constant.ContextKeyUserSetting), dto.UserSetting{BillingPreference: funding.preference, QuotaWarningThreshold: 1}) + c.Set(string(constant.ContextKeyOriginalModel), "gpt-finalize-controller") + c.Set(string(constant.ContextKeyUsingGroup), "default") + c.Set(string(constant.ContextKeyUserGroup), "default") + c.Set(string(constant.ContextKeyChannelId), 973) + c.Set(string(constant.ContextKeyChannelType), constant.ChannelTypeOpenAI) + c.Set(string(constant.ContextKeyChannelBaseUrl), upstream.URL) + c.Set(string(constant.ContextKeyChannelKey), "synthetic-controller-finalize-key") + Relay(c, format) + }) + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, path, strings.NewReader(body)) + request.Header.Set("Content-Type", "application/json") + router.ServeHTTP(recorder, request) + require.EqualValues(t, 1, calls.Load()) + require.Contains(t, recorder.Body.String(), "partial output") + var record model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).First(&record).Error) + quota := usage.prompt + usage.comp + require.Equal(t, funding.source, record.Source) + require.Equal(t, model.BillingSettlementStatusApplied, record.Status) + require.EqualValues(t, quota-100, record.AppliedFundingDelta) + canonical := model.BillingSettlementInput{OperationKey: record.OperationKey, Source: record.Source, UserID: record.UserID, + TokenID: record.TokenID, SubscriptionID: record.SubscriptionID, FundingDelta: record.FundingDelta, TokenDelta: record.TokenDelta, + SubscriptionPreConsumeRequestID: record.SubscriptionPreConsumeRequestID, FinalizeSubscriptionPreConsume: record.FinalizeSubscriptionPreConsume} + var workers sync.WaitGroup + results := make(chan error, 8) + start := make(chan struct{}) + for worker := range 8 { + workers.Add(1) + go func() { + defer workers.Done() + <-start + input := canonical + if worker >= 4 { + input.FundingDelta, input.TokenDelta = -100, -100 + input.AllowMissingToken = true + input.FinalizeSubscriptionPreConsume = input.Source == model.BillingSettlementSourceSubscription + } + applied, already, applyErr := model.ApplyBillingSettlementOnce(input) + if worker >= 4 { + if !errors.Is(applyErr, model.ErrBillingSettlementOperationConflict) || applied != 0 || already { + results <- fmt.Errorf("refund replaced original finalize: applied=%d already=%t err=%v", applied, already, applyErr) + return + } + } else if applyErr != nil || applied != int64(quota-100) || !already { + results <- fmt.Errorf("canonical replay differs: applied=%d already=%t err=%v", applied, already, applyErr) + return + } + model.ProcessPendingBillingSettlementsOnce() + results <- nil + }() + } + close(start) + workers.Wait() + close(results) + for workerErr := range results { + require.NoError(t, workerErr) + } + var after model.BillingSettlement + require.NoError(t, db.Where("id = ?", record.ID).First(&after).Error) + require.Equal(t, record, after) + var user model.User + var token model.Token + var sub model.UserSubscription + var channel model.Channel + require.NoError(t, db.First(&user, 971).Error) + require.NoError(t, db.First(&token, 972).Error) + require.NoError(t, db.First(&sub, 975).Error) + require.NoError(t, db.First(&channel, 973).Error) + wallet, subUsed := funding.wallet, funding.subUsed + if funding.source == "wallet" { + wallet -= int64(quota) + } else { + subUsed += int64(quota) + } + require.EqualValues(t, wallet, user.Quota) + require.EqualValues(t, subUsed, sub.AmountUsed) + require.EqualValues(t, 1000-quota, token.RemainQuota) + require.EqualValues(t, quota, token.UsedQuota) + require.EqualValues(t, quota, user.UsedQuota) + require.EqualValues(t, quota, channel.UsedQuota) + requests := 1 + if quota == 0 { + requests = 0 + } + require.EqualValues(t, requests, user.RequestCount) + for _, table := range []any{&model.BillingLogReceipt{}, &model.Log{}} { + var count int64 + require.NoError(t, db.Model(table).Count(&count).Error) + require.EqualValues(t, 1, count) + } + }) + } + } + } +} diff --git a/controller/relay_metadata_funding_test.go b/controller/relay_metadata_funding_test.go new file mode 100644 index 00000000..220f2802 --- /dev/null +++ b/controller/relay_metadata_funding_test.go @@ -0,0 +1,175 @@ +package controller + +import ( + "crypto/sha256" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/constant" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/model" + "github.com/MAX-API-Next/MAX-API/service" + "github.com/MAX-API-Next/MAX-API/setting/model_setting" + "github.com/MAX-API-Next/MAX-API/setting/ratio_setting" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "gorm.io/gorm" + "gorm.io/gorm/logger" +) + +func TestRelayMetadataOnlyFailureRefundsOriginalFunding(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitHttpClient() + service.InitTokenEncoders() + oldRatios, oldCompletion := ratio_setting.ModelRatio2JSONString(), ratio_setting.CompletionRatio2JSONString() + oldFloor, oldRetry, oldCount, oldEmpty := common.PreConsumedQuota, common.RetryTimes, constant.CountToken, common.EmptyCompletionRetryEnabled + common.PreConsumedQuota, common.RetryTimes, constant.CountToken, common.EmptyCompletionRetryEnabled = 100, 2, false, false + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"gpt-metadata-controller":1}`)) + require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"gpt-metadata-controller":1}`)) + t.Cleanup(func() { + common.PreConsumedQuota, common.RetryTimes, constant.CountToken, common.EmptyCompletionRetryEnabled = oldFloor, oldRetry, oldCount, oldEmpty + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(oldRatios)) + require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(oldCompletion)) + }) + for _, funding := range []struct { + name, preference, source string + wallet, subUsed int64 + allow bool + }{ + {"subscription_only", "subscription_only", "subscription", 1000, 0, false}, + {"subscription_first", "subscription_first", "subscription", 1000, 0, false}, + {"wallet_first_fallback", "wallet_first", "subscription", 0, 0, false}, + {"wallet_only", "wallet_only", "wallet", 1000, 0, false}, + {"wallet_first", "wallet_first", "wallet", 1000, 0, false}, + {"subscription_first_fallback", "subscription_first", "wallet", 1000, 1000, true}, + } { + for _, usage := range []string{"known", "zero", "missing"} { + for _, ending := range []string{"eof", "transport", "malformed"} { + t.Run(strings.Join([]string{funding.name, usage, ending}, "/"), func(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + t.Cleanup(func() { _ = sqlDB.Close() }) + require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}, &model.Channel{}, &model.Log{}, &model.BillingLogReceipt{}, + &model.BillingSettlement{}, &model.BillingPreConsumeSelection{}, &model.CacheInvalidationTask{}, + &model.UserSubscription{}, &model.SubscriptionPreConsumeRecord{}, &model.SubscriptionPlan{})) + oldDB, oldLogDB := model.DB, model.LOG_DB + oldSQLite, oldRedis, oldBatch, oldConsume := common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled + oldGlobal := *model_setting.GetGlobalSettings() + model.DB, model.LOG_DB = db, db + common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled = true, false, false, true + *model_setting.GetGlobalSettings() = model_setting.GlobalSettings{} + t.Cleanup(func() { + model.DB, model.LOG_DB = oldDB, oldLogDB + common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled = oldSQLite, oldRedis, oldBatch, oldConsume + *model_setting.GetGlobalSettings() = oldGlobal + }) + require.NoError(t, db.Create(&model.User{Id: 981, Username: "metadata-controller", Quota: funding.wallet}).Error) + require.NoError(t, db.Create(&model.Token{Id: 982, UserId: 981, Key: "synthetic-metadata-controller", RemainQuota: 1000, Status: common.TokenStatusEnabled}).Error) + require.NoError(t, db.Create(&model.Channel{Id: 983}).Error) + require.NoError(t, db.Create(&model.SubscriptionPlan{Id: 984, Title: "metadata-controller", Enabled: true, TotalAmount: 1000, + AllowWalletOverflow: &funding.allow, QuotaResetPeriod: model.SubscriptionResetNever}).Error) + now := time.Now().Unix() + require.NoError(t, db.Create(&model.UserSubscription{Id: 985, UserId: 981, PlanId: 984, Status: "active", AmountTotal: 1000, + AmountUsed: funding.subUsed, AllowWalletOverflow: funding.allow, StartTime: now - 3600, EndTime: now + 3600, LastResetTime: now - 3600}).Error) + response := map[string]any{"id": "resp_metadata_controller", "status": "in_progress", "model": "gpt-metadata-controller", "output": []any{}} + if usage != "missing" { + prompt, completion := 4, 2 + if usage == "zero" { + prompt, completion = 0, 0 + } + response["usage"] = map[string]any{"input_tokens": prompt, "output_tokens": completion, "total_tokens": prompt + completion} + } + frame, err := common.Marshal(map[string]any{"type": "response.created", "response": response}) + require.NoError(t, err) + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + if ending == "transport" { + w.Header().Set("Content-Length", "100000") + } + _, _ = fmt.Fprintf(w, "data: %s\n\n", frame) + if ending == "malformed" { + _, _ = fmt.Fprint(w, "data: {broken\n\n") + } + })) + t.Cleanup(upstream.Close) + fingerprint := sha256.Sum256([]byte(t.Name())) + requestID := fmt.Sprintf("metadata-controller-%x", fingerprint[:16]) + router := gin.New() + router.POST("/v1/responses", func(c *gin.Context) { + c.Set(common.RequestIdKey, requestID) + c.Set(string(constant.ContextKeyUserId), 981) + c.Set(string(constant.ContextKeyUserQuota), funding.wallet) + c.Set(string(constant.ContextKeyTokenId), 982) + c.Set(string(constant.ContextKeyTokenKey), "synthetic-metadata-controller") + c.Set(string(constant.ContextKeyUserSetting), dto.UserSetting{BillingPreference: funding.preference, QuotaWarningThreshold: 1}) + c.Set(string(constant.ContextKeyOriginalModel), "gpt-metadata-controller") + c.Set(string(constant.ContextKeyUsingGroup), "default") + c.Set(string(constant.ContextKeyUserGroup), "default") + c.Set(string(constant.ContextKeyChannelId), 983) + c.Set(string(constant.ContextKeyChannelType), constant.ChannelTypeOpenAI) + c.Set(string(constant.ContextKeyChannelBaseUrl), upstream.URL) + c.Set(string(constant.ContextKeyChannelKey), "synthetic-metadata-upstream") + Relay(c, types.RelayFormatOpenAIResponses) + }) + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"gpt-metadata-controller","input":"test","stream":true}`)) + request.Header.Set("Content-Type", "application/json") + router.ServeHTTP(recorder, request) + require.EqualValues(t, 1, calls.Load()) + require.Equal(t, http.StatusOK, recorder.Code) + require.Contains(t, recorder.Header().Get("Content-Type"), "text/event-stream") + require.Equal(t, 1, strings.Count(recorder.Body.String(), "event: response.failed\n")) + require.NotContains(t, recorder.Body.String(), "response.completed") + model.ProcessPendingBillingSettlementsOnce() + model.ProcessPendingBillingSettlementsOnce() + var settlement model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).First(&settlement).Error) + require.Equal(t, funding.source, settlement.Source) + require.Equal(t, model.BillingSettlementStatusApplied, settlement.Status) + require.EqualValues(t, -100, settlement.AppliedFundingDelta) + require.EqualValues(t, -100, settlement.AppliedTokenDelta) + require.Empty(t, settlement.EffectPayload) + var user model.User + var token model.Token + var channel model.Channel + var sub model.UserSubscription + require.NoError(t, db.First(&user, 981).Error) + require.NoError(t, db.First(&token, 982).Error) + require.NoError(t, db.First(&channel, 983).Error) + require.NoError(t, db.First(&sub, 985).Error) + require.EqualValues(t, funding.wallet, user.Quota) + require.EqualValues(t, funding.subUsed, sub.AmountUsed) + require.EqualValues(t, 1000, token.RemainQuota) + require.Zero(t, token.UsedQuota) + require.Zero(t, user.UsedQuota) + require.Zero(t, user.RequestCount) + require.Zero(t, channel.UsedQuota) + var logs, finalizes int64 + require.NoError(t, db.Model(&model.Log{}).Where("type = ?", model.LogTypeConsume).Count(&logs).Error) + require.Zero(t, logs) + require.NoError(t, db.Model(&model.BillingSettlement{}).Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).Count(&finalizes).Error) + require.EqualValues(t, 1, finalizes) + if funding.source == "subscription" { + var record model.SubscriptionPreConsumeRecord + require.NoError(t, db.Where("request_id = ?", requestID).First(&record).Error) + require.Equal(t, "refunded", record.Status) + require.Equal(t, sub.LastResetTime, record.SubscriptionLastResetTime) + } + }) + } + } + } +} diff --git a/controller/relay_responses_stream_error_test.go b/controller/relay_responses_stream_error_test.go new file mode 100644 index 00000000..d9067319 --- /dev/null +++ b/controller/relay_responses_stream_error_test.go @@ -0,0 +1,281 @@ +package controller + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/constant" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/model" + "github.com/MAX-API-Next/MAX-API/relay/helper" + "github.com/MAX-API-Next/MAX-API/service" + "github.com/MAX-API-Next/MAX-API/setting/model_setting" + "github.com/MAX-API-Next/MAX-API/setting/operation_setting" + "github.com/MAX-API-Next/MAX-API/setting/ratio_setting" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" + "gorm.io/gorm" + "gorm.io/gorm/logger" +) + +// Exercise Relay's real deferred writer after the selected provider adapter, +// rather than reproducing the controller's error branch in a test fixture. +func TestRelayResponsesFailureUsesClientProtocol(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitHttpClient() + service.InitTokenEncoders() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}, &model.Channel{}, &model.Log{}, &model.BillingLogReceipt{}, + &model.BillingSettlement{}, &model.BillingPreConsumeSelection{}, &model.CacheInvalidationTask{})) + oldDB, oldLogDB := model.DB, model.LOG_DB + oldSQLite, oldRedis, oldBatch, oldLogConsume := common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled + oldRatios := ratio_setting.ModelRatio2JSONString() + oldRetry, oldEmpty := common.RetryTimes, common.EmptyCompletionRetryEnabled + oldQuota := *operation_setting.GetQuotaSetting() + oldStreamingTimeout := constant.StreamingTimeout + model.DB, model.LOG_DB = db, db + common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled = true, false, false, false + common.RetryTimes, common.EmptyCompletionRetryEnabled = 2, false + operation_setting.GetQuotaSetting().EnableFreeModelPreConsume = false + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"gpt-terminal-test":0}`)) + t.Cleanup(func() { + model.DB, model.LOG_DB = oldDB, oldLogDB + common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled = oldSQLite, oldRedis, oldBatch, oldLogConsume + common.RetryTimes, common.EmptyCompletionRetryEnabled = oldRetry, oldEmpty + *operation_setting.GetQuotaSetting() = oldQuota + constant.StreamingTimeout = oldStreamingTimeout + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(oldRatios)) + _ = sqlDB.Close() + }) + require.NoError(t, db.Create(&model.User{Id: 701, Username: "terminal-test", Quota: 1000}).Error) + require.NoError(t, db.Create(&model.Channel{Id: 703}).Error) + nativeCreated := `{"type":"response.created","sequence_number":7,"response":{"id":"resp_native_terminal","object":"response","status":"in_progress","model":"gpt-terminal-test","output":[]}}` + nativeText := `{"type":"response.output_text.delta","sequence_number":8,"delta":"partial output"}` + for _, tc := range []struct { + name string + channel int + frames []string + terminal string + json bool + finish string + }{ + {"native_malformed", constant.ChannelTypeOpenAI, []string{nativeCreated, nativeText, `{broken`}, "response.failed", false, ""}, + {"native_eof", constant.ChannelTypeOpenAI, []string{nativeCreated, nativeText}, "response.failed", false, ""}, + {"native_disconnect", constant.ChannelTypeOpenAI, []string{nativeCreated, nativeText}, "response.failed", false, ""}, + {"native_idle_timeout", constant.ChannelTypeOpenAI, []string{nativeCreated, nativeText}, "response.failed", false, ""}, + {"native_provider_failed", constant.ChannelTypeOpenAI, []string{nativeCreated, nativeText, `{"type":"response.failed","sequence_number":9,"response":{"id":"resp_native_terminal","status":"failed","error":{"code":"server_error","message":"provider failure"}}}`}, "response.failed", false, ""}, + {"native_unstarted", constant.ChannelTypeOpenAI, []string{`{broken`}, "", true, ""}, + {"chat_malformed", constant.ChannelTypeAdvancedCustom, []string{`{"choices":[{"delta":{"content":"partial output"}}]}`, `{broken`}, "response.failed", false, ""}, + {"chat_eof", constant.ChannelTypeAdvancedCustom, []string{`{"choices":[{"delta":{"content":"partial output"}}]}`}, "response.incomplete", false, ""}, + {"claude_malformed", constant.ChannelTypeAnthropic, []string{`{"type":"message_start","message":{"id":"msg_test","model":"gpt-terminal-test"}}`, `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial output"}}`, `{broken`}, "response.failed", false, ""}, + {"gemini_eof", constant.ChannelTypeGemini, []string{`{"candidates":[{"content":{"parts":[{"text":"partial output"}]}}]}`}, "response.incomplete", false, ""}, + {"responses_to_chat_malformed", constant.ChannelTypeAdvancedCustom, []string{nativeCreated, nativeText, `{broken`}, "", false, ""}, + {"responses_to_chat_eof", constant.ChannelTypeAdvancedCustom, []string{nativeCreated, nativeText}, "", false, ""}, + {"responses_to_chat_length", constant.ChannelTypeAdvancedCustom, []string{nativeCreated, nativeText, `{"type":"response.incomplete","response":{"id":"resp_native_terminal","status":"incomplete","incomplete_details":{"reason":"max_output_tokens"}}}`}, "", false, "length"}, + {"responses_to_chat_filter", constant.ChannelTypeAdvancedCustom, []string{nativeCreated, nativeText, `{"type":"response.incomplete","response":{"id":"resp_native_terminal","status":"incomplete","incomplete_details":{"reason":"content_filter"}}}`}, "", false, "content_filter"}, + {"responses_to_chat_done", constant.ChannelTypeAdvancedCustom, []string{nativeCreated, nativeText, `{"type":"response.done","response":{"id":"resp_native_terminal","status":"completed"}}`}, "", false, "stop"}, + {"responses_to_chat_global_length", constant.ChannelTypeOpenAI, []string{nativeCreated, nativeText, `{"type":"response.incomplete","response":{"id":"resp_native_terminal","status":"incomplete","incomplete_details":{"reason":"max_output_tokens"}}}`}, "", false, "length"}, + } { + t.Run(tc.name, func(t *testing.T) { + oldGlobal := *model_setting.GetGlobalSettings() + *model_setting.GetGlobalSettings() = model_setting.GlobalSettings{} + t.Cleanup(func() { *model_setting.GetGlobalSettings() = oldGlobal }) + if tc.name == "responses_to_chat_global_length" { + model_setting.GetGlobalSettings().ChatCompletionsToResponsesPolicy = model_setting.ChatCompletionsToResponsesPolicy{Enabled: true, AllChannels: true, ModelPatterns: []string{"^gpt-terminal-test$"}} + } + constant.StreamingTimeout = oldStreamingTimeout + if tc.name == "native_idle_timeout" { + constant.StreamingTimeout = 1 + } + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + if tc.name == "native_disconnect" { + w.Header().Set("Content-Length", "100000") + } + _, _ = fmt.Fprint(w, "data: "+strings.Join(tc.frames, "\n\ndata: ")+"\n\n") + if tc.name == "native_idle_timeout" { + w.(http.Flusher).Flush() + <-r.Context().Done() + } + })) + defer upstream.Close() + router := gin.New() + var usedChannels []string + path, relayFormat := "/v1/responses", types.RelayFormat(types.RelayFormatOpenAIResponses) + requestBody := `{"model":"gpt-terminal-test","input":"test","stream":true}` + if strings.HasPrefix(tc.name, "responses_to_chat_") { + path, relayFormat = "/v1/chat/completions", types.RelayFormatOpenAI + requestBody = `{"model":"gpt-terminal-test","messages":[{"role":"user","content":"test"}],"stream":true}` + } + router.POST(path, func(c *gin.Context) { + c.Set(common.RequestIdKey, "terminal-"+tc.name) + c.Set(string(constant.ContextKeyUserId), 701) + c.Set(string(constant.ContextKeyUserQuota), int64(1000)) + c.Set(string(constant.ContextKeyOriginalModel), "gpt-terminal-test") + c.Set(string(constant.ContextKeyUsingGroup), "default") + c.Set(string(constant.ContextKeyUserGroup), "default") + c.Set(string(constant.ContextKeyChannelId), 703) + c.Set(string(constant.ContextKeyChannelType), tc.channel) + c.Set(string(constant.ContextKeyChannelBaseUrl), upstream.URL) + c.Set(string(constant.ContextKeyChannelKey), "synthetic-terminal-key") + if tc.channel == constant.ChannelTypeAdvancedCustom { + upstreamPath, converter := "/v1/chat/completions", dto.AdvancedCustomConverterOpenAIResponsesToOpenAIChatCompletions + if relayFormat == types.RelayFormatOpenAI { + upstreamPath, converter = "/v1/responses", dto.AdvancedCustomConverterOpenAIChatCompletionsToOpenAIResponses + } + c.Set(string(constant.ContextKeyChannelOtherSetting), dto.ChannelOtherSettings{AdvancedCustom: &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{ + {IncomingPath: path, UpstreamPath: upstreamPath, Converter: converter}, + }}}) + } + Relay(c, relayFormat) + usedChannels = c.GetStringSlice("use_channel") + }) + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, path, strings.NewReader(requestBody)) + request.Header.Set("Content-Type", "application/json") + router.ServeHTTP(recorder, request) + body := recorder.Body.String() + if tc.json { + require.GreaterOrEqual(t, recorder.Code, 400, body) + require.Contains(t, recorder.Header().Get("Content-Type"), "application/json") + require.True(t, gjson.Valid(body), body) + require.NotContains(t, body, "data:") + } else { + require.Equal(t, http.StatusOK, recorder.Code, body) + require.Contains(t, recorder.Header().Get("Content-Type"), "text/event-stream") + require.Contains(t, body, "partial output") + if relayFormat == types.RelayFormatOpenAI { + require.NotContains(t, body, "event: response.") + if tc.finish == "" { + require.Contains(t, body, `data: {"error":`) + } else { + require.NotContains(t, body, `"error":`) + require.Equal(t, 1, strings.Count(body, `"finish_reason":"`+tc.finish+`"`)) + } + require.Equal(t, 1, strings.Count(body, "data: [DONE]")) + for _, frame := range strings.Split(strings.TrimSpace(body), "\n\n") { + require.True(t, strings.HasPrefix(frame, "data: "), frame) + } + require.EqualValues(t, 1, calls.Load()) + return + } + require.Equal(t, 1, strings.Count(body, "event: "+tc.terminal+"\n"), body) + require.NotContains(t, body, "event: response.completed") + require.NotContains(t, body, "[DONE]") + for _, frame := range strings.Split(strings.TrimSpace(body), "\n\n") { + require.True(t, strings.HasPrefix(frame, "event: ") || strings.HasPrefix(frame, ":"), frame) + } + require.EqualValues(t, 1, calls.Load(), "committed output must forbid retry") + require.Len(t, usedChannels, 1) + } + }) + } +} + +func TestStartedResponsesFailureIdentityAndTerminalDedup(t *testing.T) { + for _, terminal := range []string{"", "response.failed", "response.incomplete", "response.completed", "response.done", "response.cancelled", "response.canceled", "error", "response.error"} { + t.Run(terminal, func(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + require.NoError(t, helper.ResponseChunkData(c, dto.ResponsesStreamResponse{Type: "response.created"}, `{"type":"response.created","sequence_number":7,"response":{"id":"resp_identity","model":"gpt-test","output":[],"status":"in_progress"}}`)) + require.NoError(t, helper.ResponseChunkData(c, dto.ResponsesStreamResponse{Type: "response.output_text.delta"}, `{"type":"response.output_text.delta","sequence_number":12,"delta":"partial"}`)) + if terminal != "" { + payload := fmt.Sprintf(`{"type":%q,"sequence_number":13}`, terminal) + require.NoError(t, helper.ResponseChunkData(c, dto.ResponsesStreamResponse{Type: terminal}, payload)) + } + before := recorder.Body.String() + apiErr := types.NewOpenAIError(errors.New("synthetic failure"), types.ErrorCodeBadResponse, 502) + require.True(t, writeStartedStreamError(c, types.RelayFormatOpenAIResponses, apiErr)) + if terminal == "" { + body := recorder.Body.String() + require.Equal(t, 1, strings.Count(body, "event: response.failed\n"), body) + last := strings.Split(strings.TrimSpace(body), "\n\n")[2] + data := strings.SplitN(last, "data: ", 2)[1] + require.Equal(t, "resp_identity", gjson.Get(data, "response.id").String()) + require.Equal(t, "failed", gjson.Get(data, "response.status").String()) + require.EqualValues(t, 13, gjson.Get(data, "sequence_number").Int()) + require.Equal(t, "server_error", gjson.Get(data, "response.error.code").String()) + } else { + require.Equal(t, before, recorder.Body.String()) + } + after := recorder.Body.String() + require.True(t, writeStartedStreamError(c, types.RelayFormatOpenAIResponses, apiErr)) + require.Equal(t, after, recorder.Body.String(), "controller replay must not append a second terminal") + }) + } +} + +type responsesControllerFlushWriter struct { + *httptest.ResponseRecorder + writes int + broken bool +} + +func (w *responsesControllerFlushWriter) Write(p []byte) (int, error) { + w.writes++ + return w.ResponseRecorder.Write(p) +} + +func (w *responsesControllerFlushWriter) Flush() { + if w.broken { + panic("synthetic flush failure") + } + w.ResponseRecorder.Flush() +} + +func TestStartedResponsesFailureDoesNotWriteAfterDownstreamFailure(t *testing.T) { + for _, kind := range []string{"cancel", "short", "pipe", "flush"} { + t.Run(kind, func(t *testing.T) { + base := httptest.NewRecorder() + partial := &partialNativeStreamWriter{ResponseRecorder: base, short: kind == "short"} + flush := &responsesControllerFlushWriter{ResponseRecorder: base} + var writer http.ResponseWriter = partial + if kind == "flush" { + writer = flush + } + c, _ := gin.CreateTestContext(writer) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil).WithContext(ctx) + require.NoError(t, helper.ResponseChunkData(c, dto.ResponsesStreamResponse{Type: "response.output_text.delta"}, `{"type":"response.output_text.delta","delta":"partial"}`)) + if kind == "cancel" { + cancel() + } else { + if kind == "flush" { + flush.broken = true + } + err := helper.ResponseChunkData(c, dto.ResponsesStreamResponse{Type: "response.output_text.delta"}, `{"type":"response.output_text.delta","delta":"broken"}`) + if kind == "short" { + require.ErrorIs(t, err, io.ErrShortWrite) + } else { + require.Error(t, err) + } + } + before, writes := base.Body.String(), partial.writes+flush.writes + apiErr := types.NewError(errors.New("upstream error"), types.ErrorCodeBadResponse) + require.True(t, writeStartedStreamError(c, types.RelayFormatOpenAIResponses, apiErr)) + require.Equal(t, before, base.Body.String()) + _, lateErr := helper.ResponseChunkDataWithDelivery(c, dto.ResponsesStreamResponse{Type: "response.failed"}, `{"type":"response.failed"}`) + require.Error(t, lateErr) + require.Equal(t, writes, partial.writes+flush.writes) + }) + } +} diff --git a/controller/relay_stream_error.go b/controller/relay_stream_error.go index d127bd77..78f31c0e 100644 --- a/controller/relay_stream_error.go +++ b/controller/relay_stream_error.go @@ -11,7 +11,7 @@ import ( // Once SSE is committed, a JSON response cannot change the HTTP status and // would corrupt event framing. The controller remains the single error writer. func writeStartedStreamError(c *gin.Context, format types.RelayFormat, apiErr *types.MaxAPIError) bool { - if format != types.RelayFormatOpenAI || !strings.HasPrefix(c.Writer.Header().Get("Content-Type"), "text/event-stream") { + if (format != types.RelayFormatOpenAI && format != types.RelayFormatOpenAIResponses) || !strings.HasPrefix(c.Writer.Header().Get("Content-Type"), "text/event-stream") { return false } if !c.Writer.Written() { @@ -19,8 +19,16 @@ func writeStartedStreamError(c *gin.Context, format types.RelayFormat, apiErr *t // Gin's ordinary error responder select application/json if no bytes // were committed; otherwise clients would parse raw JSON as SSE. c.Writer.Header().Del("Content-Type") + c.Writer.Header().Del("Cache-Control") + c.Writer.Header().Del("Connection") + c.Writer.Header().Del("Transfer-Encoding") + c.Writer.Header().Del("X-Accel-Buffering") return false } + if format == types.RelayFormatOpenAIResponses { + _ = helper.WriteResponsesStreamFailure(c, apiErr) + return true + } _ = helper.ObjectData(c, gin.H{"error": apiErr.ToOpenAIError()}) helper.Done(c) return true diff --git a/controller/relay_subscription_stream_test.go b/controller/relay_subscription_stream_test.go new file mode 100644 index 00000000..c43ddf9e --- /dev/null +++ b/controller/relay_subscription_stream_test.go @@ -0,0 +1,251 @@ +package controller + +import ( + "crypto/sha256" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/constant" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/model" + "github.com/MAX-API-Next/MAX-API/service" + "github.com/MAX-API-Next/MAX-API/setting/model_setting" + "github.com/MAX-API-Next/MAX-API/setting/ratio_setting" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" + "gorm.io/gorm" + "gorm.io/gorm/logger" +) + +// Paid requests exercise controller admission, adapter selection, partial +// settlement, the deferred refund, and the client-protocol error writer. +func TestRelaySubscriptionStreamFundingAndRefund(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitHttpClient() + service.InitTokenEncoders() + oldRatios, oldCompletion := ratio_setting.ModelRatio2JSONString(), ratio_setting.CompletionRatio2JSONString() + oldFloor, oldRetry, oldCount, oldEmpty := common.PreConsumedQuota, common.RetryTimes, constant.CountToken, common.EmptyCompletionRetryEnabled + common.PreConsumedQuota, constant.CountToken, common.EmptyCompletionRetryEnabled = 100, false, false + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"gpt-funding-controller":1}`)) + require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"gpt-funding-controller":1}`)) + t.Cleanup(func() { + common.PreConsumedQuota, common.RetryTimes, constant.CountToken, common.EmptyCompletionRetryEnabled = oldFloor, oldRetry, oldCount, oldEmpty + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(oldRatios)) + require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(oldCompletion)) + }) + for _, route := range []string{"responses_native", "chat_native", "chat_to_responses", "responses_to_chat"} { + for _, funding := range []struct { + name, preference, source string + wallet, subUsed int64 + allow bool + }{ + {"subscription_only", "subscription_only", "subscription", 1000, 0, false}, + {"wallet_only", "wallet_only", "wallet", 1000, 0, false}, + {"wallet_first_fallback", "wallet_first", "subscription", 0, 0, false}, + {"subscription_first_fallback", "subscription_first", "wallet", 1000, 1000, true}, + } { + for _, usageKind := range []string{"known", "zero", "missing", "above_reservation", "at_reservation"} { + for _, delivered := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/%s/%s/delivered_%t", route, funding.name, usageKind, delivered), func(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + t.Cleanup(func() { _ = sqlDB.Close() }) + require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}, &model.Channel{}, &model.Log{}, &model.BillingLogReceipt{}, + &model.BillingSettlement{}, &model.BillingPreConsumeSelection{}, &model.CacheInvalidationTask{}, + &model.UserSubscription{}, &model.SubscriptionPreConsumeRecord{}, &model.SubscriptionPlan{})) + oldDB, oldLogDB := model.DB, model.LOG_DB + oldSQLite, oldRedis, oldBatch, oldConsume := common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled + oldGlobal := *model_setting.GetGlobalSettings() + model.DB, model.LOG_DB = db, db + common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled = true, false, false, true + *model_setting.GetGlobalSettings() = model_setting.GlobalSettings{} + // Unstarted failures test the selected attempt's refund. Existing + // retry-routing tests separately cover selecting a new channel. + common.RetryTimes = 0 + if delivered { + common.RetryTimes = 2 + } + t.Cleanup(func() { + model.DB, model.LOG_DB = oldDB, oldLogDB + common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled = oldSQLite, oldRedis, oldBatch, oldConsume + *model_setting.GetGlobalSettings() = oldGlobal + }) + require.NoError(t, db.Create(&model.User{Id: 971, Username: "controller-stream-funding", Quota: funding.wallet}).Error) + require.NoError(t, db.Create(&model.Token{Id: 972, UserId: 971, Key: "synthetic-controller-funding", RemainQuota: 1000, Status: common.TokenStatusEnabled}).Error) + require.NoError(t, db.Create(&model.Channel{Id: 973}).Error) + require.NoError(t, db.Create(&model.SubscriptionPlan{Id: 974, Title: "controller-funding-plan", Enabled: true, TotalAmount: 1000, + AllowWalletOverflow: &funding.allow, QuotaResetPeriod: model.SubscriptionResetNever}).Error) + now := time.Now().Unix() + require.NoError(t, db.Create(&model.UserSubscription{Id: 975, UserId: 971, PlanId: 974, Status: "active", AmountTotal: 1000, + AmountUsed: funding.subUsed, AllowWalletOverflow: funding.allow, StartTime: now - 3600, EndTime: now + 3600, LastResetTime: now - 3600}).Error) + prompt, completion := 4, 2 + if usageKind == "zero" { + prompt, completion = 0, 0 + } else if usageKind == "above_reservation" { + prompt, completion = 80, 50 + } else if usageKind == "at_reservation" { + prompt, completion = 80, 20 + } else if usageKind == "missing" { + prompt, completion = 0, service.EstimateTokenByModel("gpt-funding-controller", "partial output") + if route == "responses_native" { + completion = service.CountTextToken("partial output", "gpt-funding-controller") + } + } + chatUpstream := route == "chat_native" || route == "chat_to_responses" + frames := []string{} + if delivered { + if chatUpstream { + frame := map[string]any{"choices": []any{map[string]any{"delta": map[string]any{"content": "partial output"}}}} + if usageKind != "missing" { + frame["usage"] = map[string]any{"prompt_tokens": prompt, "completion_tokens": completion, "total_tokens": prompt + completion} + } + b, marshalErr := common.Marshal(frame) + require.NoError(t, marshalErr) + frames = append(frames, string(b)) + } else { + response := map[string]any{"id": "resp_controller_funding", "model": "gpt-funding-controller", "status": "in_progress", "output": []any{}} + if usageKind != "missing" { + response["usage"] = map[string]any{"input_tokens": prompt, "output_tokens": completion, "total_tokens": prompt + completion} + } + b, marshalErr := common.Marshal(map[string]any{"type": "response.created", "response": response}) + require.NoError(t, marshalErr) + frames = append(frames, string(b), `{"type":"response.output_text.delta","delta":"partial output"}`) + } + } + frames = append(frames, `{broken`) + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, "data: "+strings.Join(frames, "\n\ndata: ")+"\n\n") + })) + t.Cleanup(upstream.Close) + fingerprint := sha256.Sum256([]byte(t.Name())) + requestID := fmt.Sprintf("controller-funding-%x", fingerprint[:16]) + path, format := "/v1/responses", types.RelayFormat(types.RelayFormatOpenAIResponses) + requestBody := `{"model":"gpt-funding-controller","input":"test","stream":true,"max_output_tokens":16}` + if route == "chat_native" || route == "responses_to_chat" { + path, format = "/v1/chat/completions", types.RelayFormatOpenAI + requestBody = `{"model":"gpt-funding-controller","messages":[{"role":"user","content":"test"}],"stream":true,"max_tokens":16}` + } + channelType := constant.ChannelTypeOpenAI + if route == "chat_to_responses" || route == "responses_to_chat" { + channelType = constant.ChannelTypeAdvancedCustom + } + router := gin.New() + router.POST(path, func(c *gin.Context) { + c.Set(common.RequestIdKey, requestID) + c.Set(string(constant.ContextKeyUserId), 971) + c.Set(string(constant.ContextKeyUserQuota), funding.wallet) + c.Set(string(constant.ContextKeyTokenId), 972) + c.Set(string(constant.ContextKeyTokenKey), "synthetic-controller-funding") + c.Set(string(constant.ContextKeyUserSetting), dto.UserSetting{BillingPreference: funding.preference, QuotaWarningThreshold: 1}) + c.Set(string(constant.ContextKeyOriginalModel), "gpt-funding-controller") + c.Set(string(constant.ContextKeyUsingGroup), "default") + c.Set(string(constant.ContextKeyUserGroup), "default") + c.Set(string(constant.ContextKeyChannelId), 973) + c.Set(string(constant.ContextKeyChannelType), channelType) + c.Set(string(constant.ContextKeyChannelBaseUrl), upstream.URL) + c.Set(string(constant.ContextKeyChannelKey), "synthetic-controller-key") + if channelType == constant.ChannelTypeAdvancedCustom { + upstreamPath, converter := "/v1/chat/completions", dto.AdvancedCustomConverterOpenAIResponsesToOpenAIChatCompletions + if route == "responses_to_chat" { + upstreamPath, converter = "/v1/responses", dto.AdvancedCustomConverterOpenAIChatCompletionsToOpenAIResponses + } + c.Set(string(constant.ContextKeyChannelOtherSetting), dto.ChannelOtherSettings{AdvancedCustom: &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{ + {IncomingPath: path, UpstreamPath: upstreamPath, Converter: converter}, + }}}) + } + Relay(c, format) + }) + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, path, strings.NewReader(requestBody)) + request.Header.Set("Content-Type", "application/json") + router.ServeHTTP(recorder, request) + require.EqualValues(t, 1, calls.Load(), recorder.Body.String()) + quota, logs, requests := 0, 0, 0 + if delivered { + quota, logs = prompt+completion, 1 + if quota != 0 { + requests = 1 + } + require.Equal(t, http.StatusOK, recorder.Code) + require.Contains(t, recorder.Header().Get("Content-Type"), "text/event-stream") + require.Contains(t, recorder.Body.String(), "partial output") + if format == types.RelayFormatOpenAIResponses { + require.Equal(t, 1, strings.Count(recorder.Body.String(), "event: response.failed\n")) + require.NotContains(t, recorder.Body.String(), "[DONE]") + } else { + require.Equal(t, 1, strings.Count(recorder.Body.String(), "data: [DONE]")) + require.NotContains(t, recorder.Body.String(), "event: response.") + } + } else { + require.GreaterOrEqual(t, recorder.Code, 400) + require.True(t, gjson.Valid(recorder.Body.String())) + require.Contains(t, recorder.Header().Get("Content-Type"), "application/json") + } + var selection model.BillingPreConsumeSelection + var settlement model.BillingSettlement + require.NoError(t, db.Where("request_id = ?", requestID).First(&selection).Error) + require.Equal(t, funding.source, selection.Source) + require.EqualValues(t, 100, selection.EffectiveQuota) + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).First(&settlement).Error) + require.Equal(t, funding.source, settlement.Source) + require.Equal(t, model.BillingSettlementStatusApplied, settlement.Status) + require.EqualValues(t, quota-100, settlement.AppliedFundingDelta) + require.EqualValues(t, quota-100, settlement.AppliedTokenDelta) + var user model.User + var token model.Token + var channel model.Channel + var sub model.UserSubscription + require.NoError(t, db.First(&user, 971).Error) + require.NoError(t, db.First(&token, 972).Error) + require.NoError(t, db.First(&channel, 973).Error) + require.NoError(t, db.First(&sub, 975).Error) + wallet, subUsed := funding.wallet, funding.subUsed + if funding.source == "subscription" { + subUsed += int64(quota) + var record model.SubscriptionPreConsumeRecord + require.NoError(t, db.Where("request_id = ?", requestID).First(&record).Error) + require.Equal(t, sub.LastResetTime, record.SubscriptionLastResetTime) + } else { + wallet -= int64(quota) + } + require.EqualValues(t, wallet, user.Quota) + require.EqualValues(t, subUsed, sub.AmountUsed) + require.EqualValues(t, quota, user.UsedQuota) + require.EqualValues(t, requests, user.RequestCount) + require.EqualValues(t, 1000-quota, token.RemainQuota) + require.EqualValues(t, quota, token.UsedQuota) + require.EqualValues(t, quota, channel.UsedQuota) + var count int64 + require.NoError(t, db.Model(&model.Log{}).Where("type = ?", model.LogTypeConsume).Count(&count).Error) + require.EqualValues(t, logs, count) + if delivered { + require.Equal(t, model.BillingSettlementEffectApplied, settlement.EffectStatus) + var log model.Log + require.NoError(t, db.Where("type = ?", model.LogTypeConsume).First(&log).Error) + require.Equal(t, quota, log.Quota) + require.Equal(t, prompt, log.PromptTokens) + require.Equal(t, completion, log.CompletionTokens) + } else { + require.Empty(t, settlement.EffectPayload) + } + }) + } + } + } + } +} diff --git a/dto/billing_usage.go b/dto/billing_usage.go index b246e218..6b172251 100644 --- a/dto/billing_usage.go +++ b/dto/billing_usage.go @@ -12,9 +12,12 @@ const ( ) type BillingUsage struct { - Source string `json:"source,omitempty"` - Semantic string `json:"semantic,omitempty"` - Estimated bool `json:"estimated,omitempty"` + Source string `json:"source,omitempty"` + Semantic string `json:"semantic,omitempty"` + Estimated bool `json:"estimated,omitempty"` + // TokenCountsReported distinguishes numeric zero from absent/null counters. + // Set only from provider wire fields, never inferred from Go zero values. + TokenCountsReported bool `json:"token_counts_reported,omitempty"` OpenAIUsage *Usage `json:"openai_usage,omitempty"` ClaudeUsage *ClaudeUsage `json:"claude_usage,omitempty"` GeminiUsageMetadata *GeminiUsageMetadata `json:"gemini_usage_metadata,omitempty"` @@ -58,6 +61,23 @@ func NewOpenAIResponsesBillingUsage(usage *Usage) *BillingUsage { return newOpenAIBillingUsage(BillingUsageSourceOAIResponses, usage) } +// Reported usage constructors retain an explicit zero pair for settlement. +func NewReportedOpenAIChatBillingUsage(usage *Usage) *BillingUsage { + return newReportedOpenAIBillingUsage(BillingUsageSourceOAIChat, usage) +} + +func NewReportedOpenAIResponsesBillingUsage(usage *Usage) *BillingUsage { + return newReportedOpenAIBillingUsage(BillingUsageSourceOAIResponses, usage) +} + +func newReportedOpenAIBillingUsage(source string, usage *Usage) *BillingUsage { + if usage == nil { + return nil + } + return &BillingUsage{Source: source, Semantic: BillingUsageSemanticOpenAI, + TokenCountsReported: true, OpenAIUsage: cloneOpenAIUsage(usage)} +} + func newOpenAIBillingUsage(source string, usage *Usage) *BillingUsage { if !HasOpenAIUsageTokens(usage) { return nil diff --git a/dto/openai_response.go b/dto/openai_response.go index 1cebfacf..d5c260df 100644 --- a/dto/openai_response.go +++ b/dto/openai_response.go @@ -407,6 +407,20 @@ type ResponsesOutput struct { Input string `json:"input,omitempty"` } +// A custom tool's empty input is still a supplied protocol value. Preserve it +// without adding an input field to ordinary Responses output items. +func (r ResponsesOutput) MarshalJSON() ([]byte, error) { + type outputAlias ResponsesOutput + var input *string + if r.Type == BuildInCallCustomToolCall || r.Input != "" { + input = &r.Input + } + return common.Marshal(struct { + outputAlias + Input *string `json:"input,omitempty"` + }{outputAlias: outputAlias(r), Input: input}) +} + // ArgumentsString returns function arguments or custom-tool input in the string form expected by Chat Completions. func (r *ResponsesOutput) ArgumentsString() string { if r == nil { @@ -458,13 +472,15 @@ const ( // ResponsesStreamResponse 用于处理 /v1/responses 流式响应 type ResponsesStreamResponse struct { - Type string `json:"type"` - Response *OpenAIResponsesResponse `json:"response,omitempty"` - Code any `json:"code,omitempty"` - Message string `json:"message,omitempty"` - Param string `json:"param,omitempty"` - Delta string `json:"delta,omitempty"` - Item *ResponsesOutput `json:"item,omitempty"` + Type string `json:"type"` + Response *OpenAIResponsesResponse `json:"response,omitempty"` + Code any `json:"code,omitempty"` + Message string `json:"message,omitempty"` + Param string `json:"param,omitempty"` + Delta string `json:"delta,omitempty"` + Input *string `json:"input,omitempty"` + Arguments *string `json:"arguments,omitempty"` + Item *ResponsesOutput `json:"item,omitempty"` // - response.function_call_arguments.delta // - response.function_call_arguments.done OutputIndex *int `json:"output_index,omitempty"` diff --git a/model/subscription.go b/model/subscription.go index 15680a31..6135c3d3 100644 --- a/model/subscription.go +++ b/model/subscription.go @@ -34,9 +34,10 @@ const ( ) var ( - ErrSubscriptionOrderNotFound = errors.New("subscription order not found") - ErrSubscriptionOrderStatusInvalid = errors.New("subscription order status invalid") - ErrSubscriptionQuotaInsufficient = errors.New("subscription quota insufficient") + ErrSubscriptionOrderNotFound = errors.New("subscription order not found") + ErrSubscriptionOrderStatusInvalid = errors.New("subscription order status invalid") + ErrSubscriptionQuotaInsufficient = errors.New("subscription quota insufficient") + ErrSubscriptionResetPendingBilling = errors.New("subscription has unsettled requests; retry quota reset after settlement") ) func invalidateSubscriptionUserCache(userId int, operation string) { @@ -1073,19 +1074,86 @@ func resetUserSubscriptionTx(tx *gorm.DB, sub *UserSubscription, plan *Subscript if tx == nil || sub == nil || plan == nil { return errors.New("invalid reset args") } - sub.AmountUsed = 0 + updated := *sub + updated.AmountUsed = 0 if advanceResetTime { nextReset := calcNextResetTime(time.Unix(now, 0), plan, sub.EndTime) - sub.NextResetTime = nextReset + updated.NextResetTime = nextReset if nextReset > 0 { - sub.LastResetTime = now + updated.LastResetTime = now } else { - sub.LastResetTime = 0 + updated.LastResetTime = 0 + } + } + if updated.LastResetTime == sub.LastResetTime { + if err := validateSubscriptionResetBillingTx(tx, *sub); err != nil { + return err } } + *sub = updated return tx.Save(sub).Error } +// The caller holds the subscription row lock. An unchanged period anchor +// cannot distinguish a cleared reservation from later usage in that period. +// Only a recorded funding application makes such a reservation safe to reset; +// effect-only replay does not debit the subscription again. +func validateSubscriptionResetBillingTx(tx *gorm.DB, sub UserSubscription) error { + var unresolved BillingSettlement + query := tx.Select("id").Where("source = ? AND subscription_id = ?", BillingSettlementSourceSubscription, sub.Id). + Where("status NOT IN ?", []string{BillingSettlementStatusApplied, ""}).Limit(1).Find(&unresolved) + if query.Error != nil { + return query.Error + } + if query.RowsAffected > 0 { + return ErrSubscriptionResetPendingBilling + } + + const batchSize = 200 + lastID := 0 + for { + var records []SubscriptionPreConsumeRecord + if err := tx.Select("id", "request_id", "user_id"). + Where("user_subscription_id = ? AND subscription_last_reset_time = ? AND status = ?", sub.Id, sub.LastResetTime, "consumed"). + Where("id > ?", lastID).Order("id asc").Limit(batchSize).Find(&records).Error; err != nil { + return err + } + if len(records) == 0 { + return nil + } + keys := make([]string, len(records)) + requestIDs := make([]string, len(records)) + for i, record := range records { + keys[i] = BillingRequestFinalizeOperationKey(record.RequestId) + requestIDs[i] = record.RequestId + } + var settlements []BillingSettlement + if err := tx.Select("operation_key", "source", "status", "user_id", "subscription_id", "subscription_pre_consume_request_id", "task_id"). + Where("operation_key IN ? OR subscription_pre_consume_request_id IN ?", keys, requestIDs).Find(&settlements).Error; err != nil { + return err + } + applied := make(map[string]bool, len(settlements)) + for _, settlement := range settlements { + canonicalFinalize := settlement.OperationKey == BillingRequestFinalizeOperationKey(settlement.SubscriptionPreConsumeRequestID) + if settlement.TaskID > 0 { + canonicalFinalize = canonicalFinalize || settlement.OperationKey == BillingTaskFinalizeOperationKey(settlement.TaskID) || + settlement.OperationKey == BillingTaskManualCompletionOperationKey(settlement.TaskID) + } + // Empty status is an applied legacy record, as in ApplyBillingSettlementOnce. + if canonicalFinalize && (settlement.Status == BillingSettlementStatusApplied || settlement.Status == "") && + settlement.Source == BillingSettlementSourceSubscription && settlement.UserID == sub.UserId && settlement.SubscriptionID == sub.Id { + applied[settlement.SubscriptionPreConsumeRequestID] = true + } + } + for _, record := range records { + if record.UserId != sub.UserId || !applied[record.RequestId] { + return ErrSubscriptionResetPendingBilling + } + } + lastID = records[len(records)-1].Id + } +} + func buildSubscriptionResetResult(plan *SubscriptionPlan, subs []UserSubscription, advanceResetTime bool) *SubscriptionResetResult { userIds := make([]int, 0, len(subs)) seenUsers := make(map[int]struct{}, len(subs)) diff --git a/model/subscription_reset_billing_test.go b/model/subscription_reset_billing_test.go new file mode 100644 index 00000000..4aac189a --- /dev/null +++ b/model/subscription_reset_billing_test.go @@ -0,0 +1,259 @@ +package model + +import ( + "errors" + "fmt" + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func seedResetBillingReservation(t *testing.T) (SubscriptionPlan, UserSubscription, string) { + t.Helper() + setupUserUpdateTestState(t) + now := GetDBTimestamp() + plan := SubscriptionPlan{Id: 4503, Title: "reset-billing", Enabled: true, TotalAmount: 1000, QuotaResetPeriod: SubscriptionResetDaily} + sub := UserSubscription{Id: 4504, UserId: 4501, PlanId: plan.Id, Status: "active", AmountTotal: 1000, + StartTime: now - 3600, EndTime: now + 172800, LastResetTime: now, NextResetTime: now + 86400} + require.NoError(t, DB.Create(&User{Id: sub.UserId, Username: "reset-billing-user", Quota: 1000, Status: common.UserStatusEnabled}).Error) + require.NoError(t, DB.Create(&Token{Id: 4502, UserId: sub.UserId, Key: "synthetic-reset-billing", RemainQuota: 1000, Status: common.TokenStatusEnabled}).Error) + require.NoError(t, DB.Create(&plan).Error) + require.NoError(t, DB.Create(&sub).Error) + requestID := "reset-billing-reservation" + _, err := PreConsumeTokenAndUserSubscription(requestID, sub.UserId, 4502, "synthetic-reset-billing", "gpt-test", 0, 100) + require.NoError(t, err) + return plan, sub, requestID +} + +func resetBillingFinalize(sub UserSubscription, requestID string) BillingSettlementInput { + return BillingSettlementInput{OperationKey: BillingRequestFinalizeOperationKey(requestID), Source: BillingSettlementSourceSubscription, + UserID: sub.UserId, SubscriptionID: sub.Id, TokenID: 4502, TokenKey: "synthetic-reset-billing", FundingDelta: -94, TokenDelta: -94, + SubscriptionPreConsumeRequestID: requestID} +} + +func TestSubscriptionResetValidatesFundingEvidence(t *testing.T) { + for _, tc := range []struct { + name string + updates map[string]any + missing bool + allowed bool + }{ + {name: "missing_finalize", missing: true}, + {name: "funding_pending", updates: map[string]any{"status": BillingSettlementStatusPending}}, + {name: "funding_manual", updates: map[string]any{"status": BillingSettlementStatusManual}}, + {name: "funding_unknown", updates: map[string]any{"status": "outcome_unknown"}}, + {name: "applied", allowed: true}, + {name: "legacy_applied", updates: map[string]any{"status": ""}, allowed: true}, + {name: "effect_pending", updates: map[string]any{"effect_status": BillingSettlementEffectPending}, allowed: true}, + {name: "wrong_source", updates: map[string]any{"source": BillingSettlementSourceWallet}}, + {name: "wrong_user", updates: map[string]any{"user_id": 9999}}, + {name: "wrong_subscription", updates: map[string]any{"subscription_id": 9999}}, + {name: "wrong_request", updates: map[string]any{"subscription_pre_consume_request_id": "another-request"}}, + {name: "wrong_operation", updates: map[string]any{"operation_key": BillingRequestFinalizeOperationKey("another-request")}}, + } { + t.Run(tc.name, func(t *testing.T) { + plan, sub, requestID := seedResetBillingReservation(t) + wantUsed := int64(100) + if !tc.missing { + _, _, err := ApplyBillingSettlementOnce(resetBillingFinalize(sub, requestID)) + require.NoError(t, err) + wantUsed = 6 + if tc.updates != nil { + require.NoError(t, DB.Model(&BillingSettlement{}).Where("operation_key = ?", BillingRequestFinalizeOperationKey(requestID)).Updates(tc.updates).Error) + } + } + _, err := AdminResetUserSubscriptionsByPlan(sub.UserId, plan.Id, false) + if tc.allowed { + require.NoError(t, err) + wantUsed = 0 + } else { + require.ErrorIs(t, err, ErrSubscriptionResetPendingBilling) + } + var got UserSubscription + require.NoError(t, DB.First(&got, sub.Id).Error) + require.Equal(t, wantUsed, got.AmountUsed) + require.Equal(t, sub.LastResetTime, got.LastResetTime) + require.Equal(t, sub.NextResetTime, got.NextResetTime) + var token Token + require.NoError(t, DB.First(&token, 4502).Error) + wantToken := 994 + if tc.missing { + wantToken = 900 + } + require.EqualValues(t, wantToken, token.RemainQuota) + }) + } +} + +func TestSubscriptionResetSameSecondRefusesUnsettledReservation(t *testing.T) { + plan, sub, _ := seedResetBillingReservation(t) + err := DB.Transaction(func(tx *gorm.DB) error { + var locked UserSubscription + if err := withRowLock(tx).First(&locked, sub.Id).Error; err != nil { + return err + } + return resetUserSubscriptionTx(tx, &locked, &plan, sub.LastResetTime, true) + }) + require.ErrorIs(t, err, ErrSubscriptionResetPendingBilling) + var got UserSubscription + require.NoError(t, DB.First(&got, sub.Id).Error) + require.EqualValues(t, 100, got.AmountUsed) + require.Equal(t, sub.LastResetTime, got.LastResetTime) +} + +func TestSubscriptionPlanResetRollsBackEarlierSubscriptions(t *testing.T) { + plan, sub, _ := seedResetBillingReservation(t) + first := sub + first.Id, first.AmountUsed, first.EndTime = 4505, 25, sub.EndTime-1 + require.NoError(t, DB.Create(&first).Error) + _, err := AdminResetPlanSubscriptions(plan.Id, false) + require.ErrorIs(t, err, ErrSubscriptionResetPendingBilling) + var got UserSubscription + require.NoError(t, DB.First(&got, first.Id).Error) + require.EqualValues(t, 25, got.AmountUsed, "an earlier reset must roll back when a later subscription is blocked") + got = UserSubscription{} + require.NoError(t, DB.First(&got, sub.Id).Error) + require.EqualValues(t, 100, got.AmountUsed) +} + +func TestSubscriptionResetLookupFailureDoesNotEraseReservation(t *testing.T) { + plan, sub, _ := seedResetBillingReservation(t) + injected := errors.New("reset funding lookup unavailable") + const callback = "test:reset_billing_lookup_failure" + require.NoError(t, DB.Callback().Query().Before("gorm:query").Register(callback, func(tx *gorm.DB) { + if tx.Statement.Table == "billing_settlements" { + tx.AddError(injected) + } + })) + t.Cleanup(func() { require.NoError(t, DB.Callback().Query().Remove(callback)) }) + _, err := AdminResetUserSubscriptionsByPlan(sub.UserId, plan.Id, false) + require.ErrorIs(t, err, injected) + var got UserSubscription + require.NoError(t, DB.First(&got, sub.Id).Error) + require.EqualValues(t, 100, got.AmountUsed) +} + +func TestSubscriptionResetRetainsManualFundingGuardAfterTombstoneCleanup(t *testing.T) { + plan, sub, requestID := seedResetBillingReservation(t) + input := resetBillingFinalize(sub, requestID) + input.FundingDelta, input.TokenDelta = 2000, 2000 + _, _, err := ApplyBillingSettlementOnce(input) + require.ErrorIs(t, err, ErrSubscriptionQuotaInsufficient) + require.NoError(t, DB.Model(&SubscriptionPreConsumeRecord{}).Where("request_id = ?", requestID). + Update("created_at", GetDBTimestamp()-86400).Error) + deleted, err := CleanupSubscriptionPreConsumeRecords(3600) + require.NoError(t, err) + require.EqualValues(t, 1, deleted) + _, err = AdminResetUserSubscriptionsByPlan(sub.UserId, plan.Id, false) + require.ErrorIs(t, err, ErrSubscriptionResetPendingBilling) + var got UserSubscription + require.NoError(t, DB.First(&got, sub.Id).Error) + require.EqualValues(t, 100, got.AmountUsed) +} + +func TestSubscriptionResetChecksAllReservationsBeyondOneBatch(t *testing.T) { + plan, sub, requestID := seedResetBillingReservation(t) + require.NoError(t, DB.Model(&SubscriptionPreConsumeRecord{}).Where("request_id = ?", requestID).Update("id", 2000).Error) + var records []SubscriptionPreConsumeRecord + var applied []BillingSettlement + for i := 1; i <= 1050; i++ { + id := fmt.Sprintf("reset-history-%d", i) + records = append(records, SubscriptionPreConsumeRecord{Id: i, RequestId: id, UserId: sub.UserId, UserSubscriptionId: sub.Id, + SubscriptionLastResetTime: sub.LastResetTime, PreConsumed: 1, Status: "consumed"}) + applied = append(applied, BillingSettlement{OperationKey: BillingRequestFinalizeOperationKey(id), Source: BillingSettlementSourceSubscription, + UserID: sub.UserId, SubscriptionID: sub.Id, SubscriptionPreConsumeRequestID: id, Status: BillingSettlementStatusApplied}) + } + require.NoError(t, DB.CreateInBatches(&records, 100).Error) + require.NoError(t, DB.CreateInBatches(&applied, 100).Error) + _, err := AdminResetUserSubscriptionsByPlan(sub.UserId, plan.Id, false) + require.ErrorIs(t, err, ErrSubscriptionResetPendingBilling) + var got UserSubscription + require.NoError(t, DB.First(&got, sub.Id).Error) + require.EqualValues(t, 100, got.AmountUsed) +} + +func TestSubscriptionResetAllowsEffectOnlyRecovery(t *testing.T) { + plan, sub, requestID := seedResetBillingReservation(t) + input := resetBillingFinalize(sub, requestID) + input.Effect = &BillingSettlementEffect{LogType: LogTypeConsume, TokenID: 4502, ModelName: "gpt-test", UpdateUsage: true, + PromptTokens: 4, CompletionTokens: 2, Quota: 6, QuotaIsActual: true} + _, _, err := ApplyBillingSettlementOnce(input) + require.NoError(t, err) + _, err = AdminResetUserSubscriptionsByPlan(sub.UserId, plan.Id, false) + require.NoError(t, err) + for range 2 { + require.NoError(t, ProcessBillingSettlementEffect(input.OperationKey)) + _, already, err := ApplyBillingSettlementOnce(input) + require.NoError(t, err) + require.True(t, already) + } + var got UserSubscription + require.NoError(t, DB.First(&got, sub.Id).Error) + require.Zero(t, got.AmountUsed) + var token Token + require.NoError(t, DB.First(&token, 4502).Error) + require.EqualValues(t, 994, token.RemainQuota) + var count int64 + require.NoError(t, DB.Model(&Log{}).Where("type = ?", LogTypeConsume).Count(&count).Error) + require.EqualValues(t, 1, count) +} + +func TestSubscriptionResetAllowsAppliedTaskFunding(t *testing.T) { + plan, sub, requestID := seedResetBillingReservation(t) + require.NoError(t, DB.Create(&Task{ID: 4506, Quota: 100}).Error) + input := resetBillingFinalize(sub, requestID) + input.TaskID, input.TaskQuota, input.TaskQuotaTarget = 4506, 100, 6 + input.OperationKey = BillingTaskFinalizeOperationKey(input.TaskID) + _, _, err := ApplyBillingSettlementOnce(input) + require.NoError(t, err) + _, err = AdminResetUserSubscriptionsByPlan(sub.UserId, plan.Id, false) + require.NoError(t, err, "an applied task finalize owns its original subscription reservation") + var got UserSubscription + require.NoError(t, DB.First(&got, sub.Id).Error) + require.Zero(t, got.AmountUsed) +} + +func TestSubscriptionResetConcurrentFinalize(t *testing.T) { + for round := range 12 { + t.Run(fmt.Sprintf("round_%02d", round), func(t *testing.T) { + plan, sub, requestID := seedResetBillingReservation(t) + input := resetBillingFinalize(sub, requestID) + start := make(chan struct{}) + resetDone, settleDone := make(chan error, 1), make(chan error, 1) + go func() { + <-start + _, err := AdminResetUserSubscriptionsByPlan(sub.UserId, plan.Id, false) + resetDone <- err + }() + go func() { + <-start + _, _, err := ApplyBillingSettlementOnce(input) + settleDone <- err + }() + close(start) + settleErr, resetErr := <-settleDone, <-resetDone + require.NoError(t, settleErr) + wantUsed := int64(0) + if resetErr != nil { + require.ErrorIs(t, resetErr, ErrSubscriptionResetPendingBilling) + wantUsed = 6 + } + var got UserSubscription + require.NoError(t, DB.First(&got, sub.Id).Error) + require.Equal(t, wantUsed, got.AmountUsed) + var token Token + require.NoError(t, DB.First(&token, 4502).Error) + require.EqualValues(t, 994, token.RemainQuota) + delta, already, err := ApplyBillingSettlementOnce(input) + require.NoError(t, err) + require.True(t, already) + require.EqualValues(t, -94, delta) + _, err = AdminResetUserSubscriptionsByPlan(sub.UserId, plan.Id, false) + require.NoError(t, err) + require.NoError(t, DB.First(&got, sub.Id).Error) + require.Zero(t, got.AmountUsed) + }) + } +} diff --git a/relay/channel/baidu/relay-baidu.go b/relay/channel/baidu/relay-baidu.go index a1bf9d19..666422c0 100644 --- a/relay/channel/baidu/relay-baidu.go +++ b/relay/channel/baidu/relay-baidu.go @@ -130,7 +130,7 @@ func baiduStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http. } }) service.CloseResponseBodyGracefully(resp) - if apiErr := helper.FirstResultTimeoutError(info); apiErr != nil { + if apiErr := helper.FirstResultTimeoutError(c, info); apiErr != nil { return apiErr, nil } return nil, usage diff --git a/relay/channel/claude/adaptor.go b/relay/channel/claude/adaptor.go index 0ae63c68..e42c37ed 100644 --- a/relay/channel/claude/adaptor.go +++ b/relay/channel/claude/adaptor.go @@ -11,6 +11,7 @@ import ( "github.com/MAX-API-Next/MAX-API/relay/channel" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" "github.com/MAX-API-Next/MAX-API/relay/reasoningcompat" + "github.com/MAX-API-Next/MAX-API/service" "github.com/MAX-API-Next/MAX-API/setting/model_setting" "github.com/MAX-API-Next/MAX-API/types" @@ -120,8 +121,11 @@ func (a *Adaptor) ConvertEmbeddingRequest(c *gin.Context, info *relaycommon.Rela } func (a *Adaptor) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) { - // TODO implement me - return nil, errors.New("not implemented") + chatRequest, err := service.ResponsesRequestToChatCompletionsRequest(&request) + if err != nil { + return nil, err + } + return a.ConvertOpenAIRequest(c, info, chatRequest) } func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (any, error) { diff --git a/relay/channel/claude/relay-claude.go b/relay/channel/claude/relay-claude.go index db8c00bf..833b2210 100644 --- a/relay/channel/claude/relay-claude.go +++ b/relay/channel/claude/relay-claude.go @@ -18,6 +18,7 @@ import ( "github.com/MAX-API-Next/MAX-API/relay/reasoningcompat" "github.com/MAX-API-Next/MAX-API/relay/reasonmap" "github.com/MAX-API-Next/MAX-API/service" + "github.com/MAX-API-Next/MAX-API/service/openaicompat" "github.com/MAX-API-Next/MAX-API/setting/model_setting" "github.com/MAX-API-Next/MAX-API/types" @@ -527,10 +528,9 @@ func ResponseClaude2OpenAI(claudeResponse *dto.ClaudeResponse) *dto.OpenAITextRe Object: "chat.completion", Created: common.GetTimestamp(), } - var responseText string + var responseText strings.Builder var responseThinking string if len(claudeResponse.Content) > 0 { - responseText = claudeResponse.Content[0].GetText() if claudeResponse.Content[0].Thinking != nil { responseThinking = *claudeResponse.Content[0].Thinking } @@ -557,7 +557,7 @@ func ResponseClaude2OpenAI(claudeResponse *dto.ClaudeResponse) *dto.OpenAITextRe thinkingContent = *message.Thinking } case "text": - responseText = message.GetText() + responseText.WriteString(message.GetText()) } } choice := dto.OpenAITextResponseChoice{ @@ -567,7 +567,7 @@ func ResponseClaude2OpenAI(claudeResponse *dto.ClaudeResponse) *dto.OpenAITextRe }, FinishReason: stopReasonClaude2OpenAI(claudeResponse.StopReason), } - choice.SetStringContent(responseText) + choice.SetStringContent(responseText.String()) if len(responseThinking) > 0 { choice.ReasoningContent = &responseThinking } @@ -584,12 +584,19 @@ func ResponseClaude2OpenAI(claudeResponse *dto.ClaudeResponse) *dto.OpenAITextRe } type ClaudeResponseInfo struct { - ResponseId string - Created int64 - Model string - ResponseText strings.Builder - Usage *dto.Usage - Done bool + ResponseId string + Created int64 + Model string + ResponseText strings.Builder + Usage *dto.Usage + Done bool + ResponsesState *openaicompat.ChatToResponsesStreamState + ResponsesTerminalSeen bool + responsesOutputDelivered bool + responsesInputUsageReported bool + responsesFinalOutputUsageReported bool + responsesUsageFinalized bool + responsesDeliveredText strings.Builder pendingToolUses []claudePendingToolUse } @@ -876,27 +883,105 @@ func HandleStreamResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud if err != nil { logger.LogError(c, "send_stream_response_failed: "+err.Error()) } + } else if info.RelayFormat == types.RelayFormatOpenAIResponses { + if claudeResponse.Type == "message_stop" || + (claudeResponse.Delta != nil && claudeResponse.Delta.StopReason != nil && strings.TrimSpace(*claudeResponse.Delta.StopReason) != "") { + claudeInfo.ResponsesTerminalSeen = true + } + FormatClaudeResponseInfo(&claudeResponse, nil, claudeInfo) + if claudeResponse.Type == "message_start" && gjson.Get(data, "message.usage.input_tokens").Type == gjson.Number { + claudeInfo.responsesInputUsageReported = true + } + if claudeResponse.Type == "message_delta" && gjson.Get(data, "usage.output_tokens").Type == gjson.Number { + claudeInfo.responsesFinalOutputUsageReported = true + claudeInfo.Usage.CompletionTokens = claudeResponse.Usage.OutputTokens + } + if claudeResponse.Type == "message_start" && claudeResponse.Message != nil { + info.UpstreamModelName = claudeResponse.Message.Model + } + if claudeInfo.ResponsesState == nil { + claudeInfo.ResponsesState = openaicompat.NewChatToResponsesStreamState(claudeInfo.ResponseId, info.UpstreamModelName) + claudeInfo.ResponsesState.CustomToolNames = service.ResponsesCustomToolNames(info) + } + responsesUsage := buildOpenAIStyleUsageFromClaudeUsage(claudeInfo.Usage) + claudeInfo.ResponsesState.Usage = openaicompat.UsageFromChatUsage(&responsesUsage) + response := StreamResponseClaude2OpenAI(&claudeResponse) + if response == nil { + return nil + } + events, convertErr := openaicompat.ChatCompletionsStreamChunkToResponsesEvents(response, claudeInfo.ResponsesState) + if convertErr != nil { + return types.NewError(convertErr, types.ErrorCodeBadResponseBody) + } + return writeClaudeResponsesEvents(c, info, claudeInfo, events) } return nil } -func HandleStreamFinalResponse(c *gin.Context, info *relaycommon.RelayInfo, claudeInfo *ClaudeResponseInfo) { - defer claudeInfo.discardPendingToolUses() - if claudeInfo.Usage.PromptTokens == 0 { - //上游出错 +func writeClaudeResponsesEvents(c *gin.Context, info *relaycommon.RelayInfo, claudeInfo *ClaudeResponseInfo, events []openaicompat.ChatToResponsesStreamEvent) *types.MaxAPIError { + for _, event := range events { + payload, marshalErr := common.Marshal(event.Payload) + if marshalErr != nil { + return types.NewError(marshalErr, types.ErrorCodeJsonMarshalFailed) + } + delivered, writeErr := helper.ResponseChunkDataWithDelivery(c, dto.ResponsesStreamResponse{Type: event.Type}, string(payload)) + if delivered && helper.ResponsesStreamEventHasOutput(event.Payload) { + claudeInfo.responsesOutputDelivered = true + if event.Payload.Delta != "" { + claudeInfo.responsesDeliveredText.WriteString(event.Payload.Delta) + if claudeInfo.responsesUsageFinalized { + claudeInfo.responsesUsageFinalized = false + // Refresh only the locally synthesized sidecar, never provider usage. + if claudeInfo.Usage.BillingUsage != nil && claudeInfo.Usage.BillingUsage.Estimated { + claudeInfo.Usage.BillingUsage = nil + } + } + } + } + if writeErr != nil { + return types.NewError(writeErr, types.ErrorCodeBadResponse) + } + if delivered && helper.ResponsesStreamEventHasFirstResult(event.Payload) { + info.SetFirstResultTime() + } + } + return nil +} + +// Finalize usage independently of downstream writes so interrupted Responses +// streams retain the same provider and fallback accounting evidence as EOF. +func finalizeClaudeStreamUsage(c *gin.Context, info *relaycommon.RelayInfo, claudeInfo *ClaudeResponseInfo) { + if info.RelayFormat == types.RelayFormatOpenAIResponses && claudeInfo.responsesUsageFinalized { + return + } + complete := claudeInfo.Done + if info.RelayFormat == types.RelayFormatOpenAIResponses { + complete = claudeInfo.ResponsesTerminalSeen + } + explicitZero := info.RelayFormat == types.RelayFormatOpenAIResponses && complete && claudeInfo.responsesFinalOutputUsageReported + estimated := false + if info.RelayFormat == types.RelayFormatOpenAIResponses && claudeInfo.Usage.PromptTokens == 0 && !claudeInfo.responsesInputUsageReported { + claudeInfo.Usage.PromptTokens = info.GetEstimatePromptTokens() + estimated = claudeInfo.Usage.PromptTokens > 0 } - if claudeInfo.Usage.CompletionTokens == 0 || !claudeInfo.Done { + if (claudeInfo.Usage.CompletionTokens == 0 && !explicitZero) || !complete { if common.DebugEnabled { common.SysLog("claude response usage is not complete, maybe upstream error") } // 只补缺失字段，不整份覆盖——保留 message_start 已拿到的 cache 字段 - fallback := service.ResponseText2Usage(c, claudeInfo.ResponseText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) - if claudeInfo.Usage.CompletionTokens == 0 || - (!claudeInfo.Done && fallback.CompletionTokens > claudeInfo.Usage.CompletionTokens) { + text := claudeInfo.ResponseText.String() + if info.RelayFormat == types.RelayFormatOpenAIResponses { + text = claudeInfo.responsesDeliveredText.String() + } + fallback := service.ResponseText2Usage(c, text, info.UpstreamModelName, info.GetEstimatePromptTokens()) + if (claudeInfo.Usage.CompletionTokens == 0 && !explicitZero) || + (!complete && fallback.CompletionTokens > claudeInfo.Usage.CompletionTokens) { claudeInfo.Usage.CompletionTokens = fallback.CompletionTokens + estimated = true } - if claudeInfo.Usage.PromptTokens == 0 { + if claudeInfo.Usage.PromptTokens == 0 && !claudeInfo.responsesInputUsageReported { claudeInfo.Usage.PromptTokens = fallback.PromptTokens + estimated = true } claudeInfo.Usage.TotalTokens = claudeInfo.Usage.PromptTokens + claudeInfo.Usage.CompletionTokens } @@ -904,8 +989,38 @@ func HandleStreamFinalResponse(c *gin.Context, info *relaycommon.RelayInfo, clau claudeInfo.Usage.UsageSemantic = "anthropic" } if claudeInfo.Usage != nil && claudeInfo.Usage.BillingUsage == nil { - claudeInfo.Usage.BillingUsage = dto.NewClaudeMessagesBillingUsage(buildMessageDeltaPatchUsage(nil, claudeInfo)) + providerUsage := buildMessageDeltaPatchUsage(nil, claudeInfo) + claudeInfo.Usage.BillingUsage = dto.NewClaudeMessagesBillingUsage(providerUsage) + if claudeInfo.Usage.BillingUsage == nil && info.RelayFormat == types.RelayFormatOpenAIResponses && (claudeInfo.responsesInputUsageReported || claudeInfo.responsesFinalOutputUsageReported) { + claudeInfo.Usage.BillingUsage = &dto.BillingUsage{Source: dto.BillingUsageSourceClaudeMessages, Semantic: dto.BillingUsageSemanticAnthropic, ClaudeUsage: providerUsage} + } + if claudeInfo.Usage.BillingUsage != nil && info.RelayFormat == types.RelayFormatOpenAIResponses { + claudeInfo.Usage.BillingUsage.Estimated = estimated + claudeInfo.Usage.BillingUsage.TokenCountsReported = claudeInfo.responsesInputUsageReported && claudeInfo.responsesFinalOutputUsageReported + } } + claudeInfo.Usage.TotalTokens = claudeInfo.Usage.PromptTokens + claudeInfo.Usage.CompletionTokens + if info.RelayFormat == types.RelayFormatOpenAIResponses { + claudeInfo.responsesUsageFinalized = true + } +} + +func HandleStreamFinalResponse(c *gin.Context, info *relaycommon.RelayInfo, claudeInfo *ClaudeResponseInfo) *types.MaxAPIError { + defer claudeInfo.discardPendingToolUses() + if info.RelayFormat == types.RelayFormatOpenAIResponses && claudeInfo.ResponsesState != nil { + flushErr := writeClaudeResponsesEvents(c, info, claudeInfo, openaicompat.FlushChatCompletionsStreamToResponsesOutput(claudeInfo.ResponsesState)) + finalizeClaudeStreamUsage(c, info, claudeInfo) + if flushErr != nil { + return types.NewError(flushErr, flushErr.GetErrorCode(), types.ErrOptionWithSkipRetry()) + } + responsesUsage := buildOpenAIStyleUsageFromClaudeUsage(claudeInfo.Usage) + claudeInfo.ResponsesState.Usage = openaicompat.UsageFromChatUsage(&responsesUsage) + if finalErr := writeClaudeResponsesEvents(c, info, claudeInfo, openaicompat.FinalizeChatCompletionsStreamToResponses(claudeInfo.ResponsesState)); finalErr != nil { + return types.NewError(finalErr, finalErr.GetErrorCode(), types.ErrOptionWithSkipRetry()) + } + return nil + } + finalizeClaudeStreamUsage(c, info, claudeInfo) if info.RelayFormat == types.RelayFormatClaude { // @@ -920,6 +1035,7 @@ func HandleStreamFinalResponse(c *gin.Context, info *relaycommon.RelayInfo, clau } helper.Done(c) } + return nil } func ClaudeStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*dto.Usage, *types.MaxAPIError) { @@ -930,6 +1046,11 @@ func ClaudeStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon. ResponseText: strings.Builder{}, Usage: &dto.Usage{}, } + if info.RelayFormat == types.RelayFormatOpenAIResponses { + info.EnableFirstResultTracking() + claudeInfo.ResponsesState = openaicompat.NewChatToResponsesStreamState(claudeInfo.ResponseId, info.UpstreamModelName) + claudeInfo.ResponsesState.CustomToolNames = service.ResponsesCustomToolNames(info) + } defer claudeInfo.discardPendingToolUses() var err *types.MaxAPIError helper.StreamScannerHandler(c, resp, info, func(data string, sr *helper.StreamResult) { @@ -939,13 +1060,56 @@ func ClaudeStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon. } }) if err != nil { + if info.RelayFormat == types.RelayFormatOpenAIResponses && c.Writer.Written() { + err = types.NewError(err, err.GetErrorCode(), types.ErrOptionWithSkipRetry()) + if !claudeInfo.responsesOutputDelivered { + return nil, err + } + finalizeClaudeStreamUsage(c, info, claudeInfo) + claudeInfo.ResponsesState.MarkIncomplete("upstream_error") + return claudeInfo.Usage, types.NewError(err, err.GetErrorCode(), types.ErrOptionWithSkipRetry()) + } return nil, err } - if apiErr := helper.FirstResultTimeoutError(info); apiErr != nil { + if apiErr := helper.FirstResultTimeoutError(c, info); apiErr != nil { + if info.RelayFormat == types.RelayFormatOpenAIResponses && claudeInfo.responsesOutputDelivered { + finalizeClaudeStreamUsage(c, info, claudeInfo) + return claudeInfo.Usage, apiErr + } return nil, apiErr } + if info.RelayFormat == types.RelayFormatOpenAIResponses { + finalizeClaudeStreamUsage(c, info, claudeInfo) + streamStatus := info.StreamStatus + if !claudeInfo.ResponsesTerminalSeen { + reason := "upstream_eof" + if streamStatus != nil && streamStatus.IsAbnormalEnd() { + reason = string(streamStatus.EndReason) + } + if claudeInfo.ResponsesState != nil { + claudeInfo.ResponsesState.MarkIncomplete(reason) + // EOF means the client can still receive a terminal Responses + // event. For cancellation or transport failure, do not write + // after the client has gone away. + if streamStatus == nil || streamStatus.EndReason == relaycommon.StreamEndReasonEOF { + if finalErr := HandleStreamFinalResponse(c, info, claudeInfo); finalErr != nil { + return claudeInfo.Usage, finalErr + } + } + } + if !claudeInfo.responsesOutputDelivered { + return nil, openaicompat.NewResponsesStreamIncompleteError(reason) + } + return claudeInfo.Usage, openaicompat.NewResponsesStreamIncompleteError(reason) + } + if streamStatus != nil && streamStatus.IsAbnormalEnd() { + return claudeInfo.Usage, openaicompat.NewResponsesStreamIncompleteError(string(streamStatus.EndReason)) + } + } - HandleStreamFinalResponse(c, info, claudeInfo) + if finalErr := HandleStreamFinalResponse(c, info, claudeInfo); finalErr != nil { + return claudeInfo.Usage, finalErr + } if service.ResponseAuditEnabled() { service.SetRelayResponseAuditContent(info, claudeInfo.ResponseText.String()) } @@ -992,6 +1156,17 @@ func HandleClaudeResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud if service.ResponseAuditEnabled() { auditResponse = ResponseClaude2OpenAI(&claudeResponse) } + case types.RelayFormatOpenAIResponses: + auditResponse = ResponseClaude2OpenAI(&claudeResponse) + auditResponse.Usage = buildOpenAIStyleUsageFromClaudeUsage(claudeInfo.Usage) + responsesResponse, _, convertErr := service.ChatCompletionsResponseToResponsesResponseWithCustomTools(auditResponse, claudeInfo.ResponseId, service.ResponsesCustomToolNames(info)) + if convertErr != nil { + return types.NewError(convertErr, types.ErrorCodeBadResponseBody) + } + responseData, err = common.Marshal(responsesResponse) + if err != nil { + return types.NewError(err, types.ErrorCodeJsonMarshalFailed) + } } if claudeResponse.Usage != nil && claudeResponse.Usage.ServerToolUse != nil && claudeResponse.Usage.ServerToolUse.WebSearchRequests > 0 { diff --git a/relay/channel/claude/relay_claude_test.go b/relay/channel/claude/relay_claude_test.go index 0a45fb30..450cd3db 100644 --- a/relay/channel/claude/relay_claude_test.go +++ b/relay/channel/claude/relay_claude_test.go @@ -11,10 +11,12 @@ import ( "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/dto" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/service/openaicompat" "github.com/MAX-API-Next/MAX-API/setting/config" "github.com/MAX-API-Next/MAX-API/setting/operation_setting" "github.com/MAX-API-Next/MAX-API/types" "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -43,6 +45,52 @@ func TestClaudeStreamHandlerReturnsRetryableErrorBeforeFirstResponseTimeout(t *t require.False(t, types.IsSkipRetryError(err)) } +func TestConvertOpenAIResponsesRequestPreservesCustomToolHistory(t *testing.T) { + input, err := common.Marshal([]map[string]any{ + {"role": "user", "content": "apply the patch"}, + {"type": "custom_tool_call", "call_id": "call_custom", "name": "apply_patch", "input": "*** Begin Patch"}, + {"type": "custom_tool_call_output", "call_id": "call_custom", "output": "ok"}, + }) + require.NoError(t, err) + tools, err := common.Marshal([]map[string]any{{ + "type": "custom", + "name": "apply_patch", + "description": "Apply a patch", + }}) + require.NoError(t, err) + + converted, err := (&Adaptor{}).ConvertOpenAIResponsesRequest(nil, nil, dto.OpenAIResponsesRequest{ + Model: "claude-test", + Input: input, + Tools: tools, + }) + require.NoError(t, err) + claudeRequest, ok := converted.(*dto.ClaudeRequest) + require.True(t, ok) + require.Len(t, claudeRequest.Messages, 3) + + toolUse, err := claudeRequest.Messages[1].ParseContent() + require.NoError(t, err) + require.Len(t, toolUse, 2) + var toolUseBlock *dto.ClaudeMediaMessage + for i := range toolUse { + if toolUse[i].Type == "tool_use" { + toolUseBlock = &toolUse[i] + break + } + } + require.NotNil(t, toolUseBlock) + assert.Equal(t, "call_custom", toolUseBlock.Id) + assert.Equal(t, "apply_patch", toolUseBlock.Name) + assert.Equal(t, "*** Begin Patch", toolUseBlock.Input.(map[string]any)["input"]) + + toolResult, err := claudeRequest.Messages[2].ParseContent() + require.NoError(t, err) + require.Len(t, toolResult, 1) + assert.Equal(t, "tool_result", toolResult[0].Type) + assert.Equal(t, "call_custom", toolResult[0].ToolUseId) +} + func setClaudeToolPricesForTest(t *testing.T, additions map[string]float64) { t.Helper() setting := config.GlobalConfig.Get("tool_price_setting").(*operation_setting.ToolPriceSetting) @@ -136,6 +184,105 @@ func TestHandleClaudeResponseDataRecordsActualCustomToolUse(t *testing.T) { require.Equal(t, []relaycommon.ToolUsageItem{{Name: "lookup", CallCount: 1, PricePer1K: 5}}, info.ToolUsageSnapshot().Items) } +func TestHandleClaudeResponseDataConvertsResponsesFormat(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + info := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatOpenAIResponses, + OriginModelName: "claude-test", + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "claude-test"}, + } + response := dto.ClaudeResponse{ + Type: "message", + Id: "msg_responses", + Model: "claude-test", + StopReason: "end_turn", + Content: []dto.ClaudeMediaMessage{{ + Type: "text", + Text: common.GetPointer("hello"), + }}, + Usage: &dto.ClaudeUsage{InputTokens: 4, OutputTokens: 2}, + } + data, err := common.Marshal(response) + require.NoError(t, err) + + maxErr := HandleClaudeResponseData(c, info, &ClaudeResponseInfo{ResponseId: "msg_responses", Usage: &dto.Usage{}}, nil, data) + require.Nil(t, maxErr) + var converted dto.OpenAIResponsesResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &converted)) + require.Len(t, converted.Output, 1) + assert.Equal(t, "message", converted.Output[0].Type) + require.Len(t, converted.Output[0].Content, 1) + assert.Equal(t, "hello", converted.Output[0].Content[0].Text) + assert.Equal(t, 6, converted.Usage.TotalTokens) +} + +func TestHandleStreamResponseDataConvertsResponsesFormat(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + info := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatOpenAIResponses, + OriginModelName: "claude-test", + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "claude-test"}, + } + claudeInfo := &ClaudeResponseInfo{ + ResponseId: "msg_stream", + Usage: &dto.Usage{}, + ResponsesState: openaicompat.NewChatToResponsesStreamState("msg_stream", "claude-test"), + } + streamData := []string{ + `{"type":"message_start","message":{"id":"msg_stream","model":"claude-test","usage":{"input_tokens":4,"output_tokens":0}}}`, + `{"type":"content_block_start","index":0,"content_block":{"type":"text","text":"hello"}}`, + `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" world"}}`, + `{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":2}}`, + } + for _, data := range streamData { + require.Nil(t, HandleStreamResponseData(c, info, claudeInfo, data)) + } + HandleStreamFinalResponse(c, info, claudeInfo) + + body := recorder.Body.String() + assert.Contains(t, body, "response.created") + assert.Contains(t, body, "response.output_text.delta") + assert.Contains(t, body, "response.completed") +} + +func TestClaudeResponsesStreamEOFReturnsIncompleteWithoutFalseCompletion(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + info := &relaycommon.RelayInfo{ + IsStream: true, + RelayFormat: types.RelayFormatOpenAIResponses, + OriginModelName: "claude-test", + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "claude-test"}, + DisablePing: true, + } + body := "data: " + strings.Join([]string{ + `{"type":"message_start","message":{"id":"msg_eof","model":"claude-test","usage":{"input_tokens":4,"output_tokens":0}}}`, + `{"type":"content_block_start","index":0,"content_block":{"type":"text","text":"partial"}}`, + `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" output"}}`, + }, "\n\ndata: ") + "\n\n" + + usage, maxAPIError := ClaudeStreamHandler(c, &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(body)), + }, info) + require.NotNil(t, usage) + require.NotNil(t, maxAPIError) + require.True(t, types.IsSkipRetryError(maxAPIError)) + require.Equal(t, 3, info.ReceivedResponseCount) + require.Equal(t, 4, usage.PromptTokens) + require.Positive(t, usage.CompletionTokens) + require.NotNil(t, usage.BillingUsage) + require.Contains(t, recorder.Body.String(), " output") + require.Contains(t, recorder.Body.String(), "response.incomplete") + require.NotContains(t, recorder.Body.String(), "response.completed") +} + func TestHandleStreamResponseDataBillsCustomToolUseOnlyAfterToolUseTerminal(t *testing.T) { setClaudeToolPricesForTest(t, map[string]float64{"lookup": 5}) gin.SetMode(gin.TestMode) diff --git a/relay/channel/claude/responses_repair_test.go b/relay/channel/claude/responses_repair_test.go new file mode 100644 index 00000000..cdb9591b --- /dev/null +++ b/relay/channel/claude/responses_repair_test.go @@ -0,0 +1,104 @@ +package claude + +import ( + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/dto" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +const claudeResponsesPartialFrames = "data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_partial\",\"model\":\"claude-test\",\"usage\":{\"input_tokens\":4}}}\n\n" + + "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"partial output\"}}\n\n" + +type claudeResponsesBrokenBody struct{ *strings.Reader } + +func (b *claudeResponsesBrokenBody) Read(p []byte) (int, error) { + n, err := b.Reader.Read(p) + if err == io.EOF { + return 0, errors.New("synthetic upstream disconnect") + } + return n, err +} + +func (*claudeResponsesBrokenBody) Close() error { return nil } + +type claudeResponsesTerminalFailureWriter struct{ *httptest.ResponseRecorder } + +func (w *claudeResponsesTerminalFailureWriter) Write(p []byte) (int, error) { + if strings.Contains(string(p), "response.completed") { + return 0, errors.New("synthetic terminal write failure") + } + return w.ResponseRecorder.Write(p) +} + +type claudeResponsesCancelWriter struct { + *httptest.ResponseRecorder + cancel context.CancelFunc +} + +func (w *claudeResponsesCancelWriter) Write(p []byte) (int, error) { + n, err := w.ResponseRecorder.Write(p) + if strings.Contains(string(p), `"delta":"partial output"`) { + w.cancel() + } + return n, err +} + +func TestClaudeResponsesRepairPreservesPartialUsageOnFailure(t *testing.T) { + for _, failure := range []string{"malformed", "provider_error", "disconnect", "terminal_write", "cancel"} { + t.Run(failure, func(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + var writer http.ResponseWriter = recorder + body := claudeResponsesPartialFrames + requestContext, cancel := context.WithCancel(context.Background()) + defer cancel() + switch failure { + case "malformed": + body += "data: {broken\n\n" + case "provider_error": + body += "data: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"synthetic provider failure\"}}\n\n" + case "terminal_write": + writer = &claudeResponsesTerminalFailureWriter{recorder} + body += "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":2}}\n\n" + case "cancel": + writer = &claudeResponsesCancelWriter{recorder, cancel} + } + var reader io.ReadCloser = io.NopCloser(strings.NewReader(body)) + if failure == "disconnect" { + reader = &claudeResponsesBrokenBody{strings.NewReader(body)} + } + c, _ := gin.CreateTestContext(writer) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil).WithContext(requestContext) + info := &relaycommon.RelayInfo{IsStream: true, RelayFormat: types.RelayFormatOpenAIResponses, + OriginModelName: "claude-test", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "claude-test"}, DisablePing: true} + usage, apiErr := ClaudeStreamHandler(c, &http.Response{StatusCode: http.StatusOK, Body: reader}, info) + require.NotNil(t, apiErr) + require.True(t, types.IsSkipRetryError(apiErr)) + require.NotNil(t, usage) + require.Equal(t, 4, usage.PromptTokens) + require.Positive(t, usage.CompletionTokens) + require.NotNil(t, usage.BillingUsage) + require.Contains(t, recorder.Body.String(), "partial output") + require.NotContains(t, recorder.Body.String(), "event: response.completed") + }) + } +} + +func TestClaudeResponsesRepairPreservesAllTextBlocks(t *testing.T) { + resp := ResponseClaude2OpenAI(&dto.ClaudeResponse{Content: []dto.ClaudeMediaMessage{ + {Type: "text", Text: common.GetPointer("first ")}, + {Type: "text", Text: common.GetPointer("second")}, + }}) + require.Equal(t, "first second", resp.Choices[0].Message.StringContent()) +} diff --git a/relay/channel/claude/responses_revalidation_test.go b/relay/channel/claude/responses_revalidation_test.go new file mode 100644 index 00000000..ab318df7 --- /dev/null +++ b/relay/channel/claude/responses_revalidation_test.go @@ -0,0 +1,60 @@ +package claude + +import ( + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestClaudeResponsesUsageDistinguishesNullAndNumericZero(t *testing.T) { + for _, usageCase := range []struct { + name, input, output string + prompt int + estimatedOutput bool + }{ + {"null_output", "4", "null", 4, true}, + {"missing_output", "4", "", 4, true}, + {"zero_output", "4", "0", 4, false}, + {"null_input", "null", "2", 10, false}, + {"zero_input", "0", "2", 0, false}, + {"all_zero", "0", "0", 0, false}, + } { + t.Run(usageCase.name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + info := &relaycommon.RelayInfo{IsStream: true, DisablePing: true, RelayFormat: types.RelayFormatOpenAIResponses, + OriginModelName: "claude-test", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "claude-test"}} + info.SetEstimatePromptTokens(10) + outputField := "" + if usageCase.output != "" { + outputField = `"output_tokens":` + usageCase.output + } + frames := []string{ + fmt.Sprintf(`{"type":"message_start","message":{"id":"msg_test","model":"claude-test","usage":{"input_tokens":%s}}}`, usageCase.input), + `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial output"}}`, + `{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{` + outputField + `}}`, + } + usage, err := ClaudeStreamHandler(c, &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader("data: " + strings.Join(frames, "\n\ndata: ") + "\n\n"))}, info) + require.Nil(t, err) + require.Equal(t, usageCase.prompt, usage.PromptTokens) + if usageCase.estimatedOutput { + require.Positive(t, usage.CompletionTokens) + } else if usageCase.output == "0" { + require.Zero(t, usage.CompletionTokens) + } else { + require.Equal(t, 2, usage.CompletionTokens) + } + require.NotNil(t, usage.BillingUsage) + require.Equal(t, usageCase.estimatedOutput || usageCase.input == "null", usage.BillingUsage.Estimated) + }) + } +} diff --git a/relay/channel/dify/relay-dify.go b/relay/channel/dify/relay-dify.go index 83d15d5d..be3e6976 100644 --- a/relay/channel/dify/relay-dify.go +++ b/relay/channel/dify/relay-dify.go @@ -252,7 +252,7 @@ func difyStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R sr.Error(err) } }) - if apiErr := helper.FirstResultTimeoutError(info); apiErr != nil { + if apiErr := helper.FirstResultTimeoutError(c, info); apiErr != nil { return nil, apiErr } helper.Done(c) diff --git a/relay/channel/gemini/adaptor_responses.go b/relay/channel/gemini/adaptor_responses.go index 87c3cb58..f8fbcb6c 100644 --- a/relay/channel/gemini/adaptor_responses.go +++ b/relay/channel/gemini/adaptor_responses.go @@ -8,12 +8,6 @@ import ( "github.com/MAX-API-Next/MAX-API/dto" ) -const ( - geminiResponsesInputTypeCustomToolCall = "custom_tool_call" - geminiResponsesInputTypeCustomToolCallOutput = "custom_tool_call_output" - geminiResponsesInputTypeFunctionCallOutput = "function_call_output" -) - func preprocessGeminiOpenAIResponsesRequest(request dto.OpenAIResponsesRequest) (dto.OpenAIResponsesRequest, error) { tools, err := filterGeminiResponsesTools(request.Tools) if err != nil { @@ -43,7 +37,7 @@ func filterGeminiResponsesTools(raw []byte) ([]byte, error) { filtered := make([]map[string]any, 0, len(tools)) for _, tool := range tools { toolType := strings.TrimSpace(common.Interface2String(tool["type"])) - if toolType != "function" { + if toolType != "function" && toolType != "custom" { return nil, fmt.Errorf("gemini responses conversion does not support tool type %q", toolType) } filtered = append(filtered, tool) @@ -64,31 +58,10 @@ func filterGeminiResponsesInput(raw []byte) ([]byte, error) { return nil, err } - skippedCustomCallIDs := make(map[string]struct{}) - for _, item := range items { - if strings.TrimSpace(common.Interface2String(item["type"])) != geminiResponsesInputTypeCustomToolCall { - continue - } - if callID := strings.TrimSpace(common.Interface2String(item["call_id"])); callID != "" { - skippedCustomCallIDs[callID] = struct{}{} - } - } - - filtered := make([]map[string]any, 0, len(items)) - for _, item := range items { - itemType := strings.TrimSpace(common.Interface2String(item["type"])) - switch itemType { - case geminiResponsesInputTypeCustomToolCall, geminiResponsesInputTypeCustomToolCallOutput: - continue - case geminiResponsesInputTypeFunctionCallOutput: - if _, ok := skippedCustomCallIDs[strings.TrimSpace(common.Interface2String(item["call_id"]))]; ok { - continue - } - } - filtered = append(filtered, item) - } - - return common.Marshal(filtered) + // Custom tool history is normalized by the shared Responses-to-Chat + // converter. Keep both the call and its output so Gemini receives the same + // function-call context as Claude and OpenAI-compatible upstreams. + return common.Marshal(items) } func geminiRawJSONPresent(raw []byte) bool { diff --git a/relay/channel/gemini/adaptor_responses_test.go b/relay/channel/gemini/adaptor_responses_test.go index 73dd2647..5b577ad3 100644 --- a/relay/channel/gemini/adaptor_responses_test.go +++ b/relay/channel/gemini/adaptor_responses_test.go @@ -27,7 +27,7 @@ func TestConvertOpenAIResponsesRequestToGeminiInstructionsAndInput(t *testing.T) assert.Equal(t, "hello", got.Contents[0].Parts[0].Text) } -func TestConvertOpenAIResponsesRequestToGeminiSkipsCustomToolCalls(t *testing.T) { +func TestConvertOpenAIResponsesRequestToGeminiPreservesCustomToolCalls(t *testing.T) { got := mustConvertResponsesToGemini(t, dto.OpenAIResponsesRequest{ Model: "gemini-test", Input: mustGeminiRawMessage(t, []map[string]any{ @@ -67,20 +67,45 @@ func TestConvertOpenAIResponsesRequestToGeminiSkipsCustomToolCalls(t *testing.T) "type": "object", }, }, + { + "type": "custom", + "name": "apply_patch", + "description": "Apply a patch", + }, }), }) require.Len(t, got.GetTools(), 1) - require.Len(t, got.Contents, 2) + require.Len(t, got.Contents, 3) + functions, err := common.Marshal(got.GetTools()[0].FunctionDeclarations) + require.NoError(t, err) + assert.Contains(t, string(functions), "apply_patch") assert.Equal(t, "model", got.Contents[0].Role) - require.Len(t, got.Contents[0].Parts, 1) - assert.Equal(t, "before custom", got.Contents[0].Parts[0].Text) - assert.Nil(t, got.Contents[0].Parts[0].FunctionCall) + require.Len(t, got.Contents[0].Parts, 2) + var functionCall *dto.FunctionCall + var textPart string + for _, part := range got.Contents[0].Parts { + if part.FunctionCall != nil { + functionCall = part.FunctionCall + } + if part.Text != "" { + textPart = part.Text + } + } + require.NotNil(t, functionCall) + assert.Equal(t, "before custom", textPart) + assert.Equal(t, "call_custom", functionCall.ID) + assert.Equal(t, "apply_patch", functionCall.FunctionName) + arguments, ok := functionCall.Arguments.(map[string]any) + require.True(t, ok) + assert.Equal(t, "patch body", arguments["input"]) assert.Equal(t, "user", got.Contents[1].Role) - require.Len(t, got.Contents[1].Parts, 1) - assert.Equal(t, "next turn", got.Contents[1].Parts[0].Text) - assert.Nil(t, got.Contents[1].Parts[0].FunctionResponse) + require.Len(t, got.Contents[1].Parts, 2) + assert.Equal(t, "ok", got.Contents[1].Parts[0].FunctionResponse.Response["content"]) + assert.Equal(t, "legacy custom output", got.Contents[1].Parts[1].FunctionResponse.Response["content"]) + assert.Equal(t, "user", got.Contents[2].Role) + assert.Equal(t, "next turn", got.Contents[2].Parts[0].Text) } func TestConvertOpenAIResponsesRequestToGeminiRejectsUnsupportedTools(t *testing.T) { diff --git a/relay/channel/gemini/relay-gemini.go b/relay/channel/gemini/relay-gemini.go index 42f2c3eb..36d42d7b 100644 --- a/relay/channel/gemini/relay-gemini.go +++ b/relay/channel/gemini/relay-gemini.go @@ -410,6 +410,7 @@ func CovertOpenAI2Gemini(c *gin.Context, textRequest dto.GeneralOpenAIRequest, i } toolCall := dto.GeminiPart{ FunctionCall: &dto.FunctionCall{ + ID: call.ID, FunctionName: call.Function.Name, Arguments: args, }, @@ -903,8 +904,12 @@ func getResponseToolCall(item *dto.GeminiPart) *dto.ToolCallResponse { if err != nil { return nil } + callID := strings.TrimSpace(item.FunctionCall.ID) + if callID == "" { + callID = fmt.Sprintf("call_%s", common.GetUUID()) + } return &dto.ToolCallResponse{ - ID: fmt.Sprintf("call_%s", common.GetUUID()), + ID: callID, Type: "function", Function: dto.FunctionResponse{ Arguments: string(argsBytes), @@ -1356,7 +1361,6 @@ func streamResponseGeminiChat2OpenAI(geminiResponse *dto.GeminiChatResponse) (*d } appended++ } - isTools := false isThought := false if candidate.FinishReason != nil { // Map Gemini FinishReason to OpenAI finish_reason @@ -1401,7 +1405,6 @@ func streamResponseGeminiChat2OpenAI(geminiResponse *dto.GeminiChatResponse) (*d content.WriteByte(')') } } else if part.FunctionCall != nil { - isTools = true if call := getResponseToolCall(&part); call != nil { call.SetIndex(len(choice.Delta.ToolCalls)) choice.Delta.ToolCalls = append(choice.Delta.ToolCalls, *call) @@ -1437,9 +1440,8 @@ func streamResponseGeminiChat2OpenAI(geminiResponse *dto.GeminiChatResponse) (*d } else { choice.Delta.SetContentString(content.String()) } - if isTools { - choice.FinishReason = &constant.FinishReasonToolCalls - } + // Call arguments do not finish a candidate. The stream handler emits + // the terminal chunk once after the provider's STOP event. choices = append(choices, choice) } @@ -1449,6 +1451,18 @@ func streamResponseGeminiChat2OpenAI(geminiResponse *dto.GeminiChatResponse) (*d return &response, isStop } +func hasGeminiTerminalFinishReason(geminiResponse *dto.GeminiChatResponse) bool { + if geminiResponse == nil { + return false + } + for _, candidate := range geminiResponse.Candidates { + if candidate.FinishReason != nil && strings.TrimSpace(*candidate.FinishReason) != "" { + return true + } + } + return false +} + func handleStream(c *gin.Context, info *relaycommon.RelayInfo, resp *dto.ChatCompletionsStreamResponse) error { streamData, err := common.Marshal(resp) if err != nil { @@ -1520,7 +1534,7 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http sr.Stop(fmt.Errorf("gemini callback stopped")) } }) - if apiErr := helper.FirstResultTimeoutError(info); apiErr != nil { + if apiErr := helper.FirstResultTimeoutError(c, info); apiErr != nil { return nil, apiErr } @@ -1555,8 +1569,14 @@ func GeminiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp * finishReason := constant.FinishReasonStop toolCallIndexByChoice := make(map[int]map[string]int) nextToolCallIndexByChoice := make(map[int]int) + functionCalls := newGeminiStreamFunctionCalls() + var conversionErr *types.MaxAPIError usage, err := geminiStreamHandler(c, info, resp, func(data string, geminiResponse *dto.GeminiChatResponse) bool { + if prepareErr := functionCalls.prepare(geminiResponse); prepareErr != nil { + conversionErr = types.NewError(prepareErr, types.ErrorCodeBadResponseBody) + return false + } response, isStop := streamResponseGeminiChat2OpenAI(geminiResponse) response.Id = id @@ -1640,6 +1660,15 @@ func GeminiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp * if err != nil { return usage, err } + if conversionErr == nil && len(functionCalls.pending) > 0 { + conversionErr = types.NewError(fmt.Errorf("Gemini stream ended with an incomplete function call"), types.ErrorCodeBadResponseBody) + } + if conversionErr != nil { + if c.Writer.Written() { + conversionErr = types.NewError(conversionErr, conversionErr.GetErrorCode(), types.ErrOptionWithSkipRetry()) + } + return usage, conversionErr + } response := helper.GenerateFinalUsageResponse(id, createAt, info.UpstreamModelName, *usage) if info.RelayFormat == types.RelayFormatClaude && info.ClaudeConvertInfo != nil && !info.ClaudeConvertInfo.Done { diff --git a/relay/channel/gemini/relay_gemini_usage_test.go b/relay/channel/gemini/relay_gemini_usage_test.go index 7d810bf1..bec99b17 100644 --- a/relay/channel/gemini/relay_gemini_usage_test.go +++ b/relay/channel/gemini/relay_gemini_usage_test.go @@ -3,6 +3,7 @@ package gemini import ( "bytes" "context" + "fmt" "io" "net/http" "net/http/httptest" @@ -468,6 +469,7 @@ func TestGeminiResponsesStreamHandlerEstimatesUsageAndAuditWhenUpstreamUsageMiss chunk := dto.GeminiChatResponse{ Candidates: []dto.GeminiChatCandidate{ { + FinishReason: common.GetPointer("STOP"), Content: dto.GeminiChatContent{ Role: "model", Parts: []dto.GeminiPart{ @@ -477,7 +479,8 @@ func TestGeminiResponsesStreamHandlerEstimatesUsageAndAuditWhenUpstreamUsageMiss }, }, } - chunkData, err := common.Marshal(chunk) + // The DTO serializes zero usageMetadata; omit the field to model absence. + chunkData, err := common.Marshal(map[string]any{"candidates": chunk.Candidates}) require.NoError(t, err) streamBody := []byte("data: " + string(chunkData) + "\n" + "data: [DONE]\n") @@ -502,6 +505,141 @@ func TestGeminiResponsesStreamHandlerEstimatesUsageAndAuditWhenUpstreamUsageMiss require.Positive(t, completed.Response.Usage.TotalTokens) } +func TestGeminiResponsesStreamEOFReturnsIncompleteWithoutFalseCompletion(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + info := &relaycommon.RelayInfo{ + IsStream: true, + RelayFormat: types.RelayFormatOpenAIResponses, + OriginModelName: "gemini-test", + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gemini-test"}, + DisablePing: true, + } + chunk := `{"candidates":[{"index":0,"content":{"role":"model","parts":[{"text":"partial"}]}}]}` + + usage, maxAPIError := GeminiResponsesStreamHandler(c, info, &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader("data: " + chunk + "\n")), + }) + require.NotNil(t, usage) + require.NotNil(t, maxAPIError) + require.True(t, types.IsSkipRetryError(maxAPIError)) + require.Contains(t, recorder.Body.String(), "response.incomplete") + require.NotContains(t, recorder.Body.String(), "response.completed") +} + +func TestGeminiResponsesStreamRecognizesNonStopTerminalReasons(t *testing.T) { + tests := []struct { + name string + finishReason string + incompleteHint string + }{ + {name: "max tokens", finishReason: "MAX_TOKENS", incompleteHint: "max_output_tokens"}, + {name: "safety", finishReason: "SAFETY", incompleteHint: "content_filter"}, + {name: "recitation", finishReason: "RECITATION", incompleteHint: "content_filter"}, + {name: "blocklist", finishReason: "BLOCKLIST", incompleteHint: "content_filter"}, + {name: "prohibited content", finishReason: "PROHIBITED_CONTENT", incompleteHint: "content_filter"}, + {name: "sensitive pii", finishReason: "SPII", incompleteHint: "content_filter"}, + {name: "other provider stop", finishReason: "OTHER", incompleteHint: "content_filter"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + info := &relaycommon.RelayInfo{ + IsStream: true, + RelayFormat: types.RelayFormatOpenAIResponses, + OriginModelName: "gemini-test", + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gemini-test"}, + DisablePing: true, + } + chunk := fmt.Sprintf(`{"candidates":[{"index":0,"finishReason":%q,"content":{"role":"model","parts":[{"text":"partial"}]} }],"usageMetadata":{"promptTokenCount":4,"candidatesTokenCount":2,"totalTokenCount":6}}`, tt.finishReason) + + usage, maxAPIError := GeminiResponsesStreamHandler(c, info, &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader("data: " + chunk + "\n")), + }) + require.Nil(t, maxAPIError) + require.NotNil(t, usage) + require.Equal(t, 4, usage.PromptTokens) + require.Equal(t, 2, usage.CompletionTokens) + require.Equal(t, 6, usage.TotalTokens) + require.Contains(t, recorder.Body.String(), "response.incomplete") + require.Contains(t, recorder.Body.String(), tt.incompleteHint) + require.NotContains(t, recorder.Body.String(), "responses stream ended before a terminal event") + }) + } +} + +func TestGeminiResponsesStreamIncompletePreservesProviderUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + info := &relaycommon.RelayInfo{ + IsStream: true, + RelayFormat: types.RelayFormatOpenAIResponses, + OriginModelName: "gemini-test", + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gemini-test"}, + DisablePing: true, + } + chunk := `{"candidates":[{"index":0,"content":{"role":"model","parts":[{"text":"partial"}]}}],"usageMetadata":{"promptTokenCount":4,"candidatesTokenCount":2,"totalTokenCount":6}}` + + usage, maxAPIError := GeminiResponsesStreamHandler(c, info, &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader("data: " + chunk + "\n")), + }) + require.NotNil(t, maxAPIError) + require.True(t, types.IsSkipRetryError(maxAPIError)) + require.Equal(t, 4, usage.PromptTokens) + var incomplete dto.ResponsesStreamResponse + require.NoError(t, common.UnmarshalJsonStr(sseDataForEvent(t, recorder.Body.String(), "response.incomplete"), &incomplete)) + require.NotNil(t, incomplete.Response) + require.NotNil(t, incomplete.Response.Usage) + require.Equal(t, 4, incomplete.Response.Usage.InputTokens) + require.Equal(t, 2, incomplete.Response.Usage.OutputTokens) + require.Equal(t, 6, incomplete.Response.Usage.TotalTokens) +} + +func TestGeminiResponsesStreamMalformedFrameIsAbnormal(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + info := &relaycommon.RelayInfo{ + IsStream: true, + RelayFormat: types.RelayFormatOpenAIResponses, + OriginModelName: "gemini-test", + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gemini-test"}, + DisablePing: true, + } + valid := `data: {"candidates":[{"index":0,"content":{"role":"model","parts":[{"text":"partial"}]}}]}` + "\n" + reader, writer := io.Pipe() + finished := make(chan struct{}) + go func() { + _, _ = writer.Write([]byte(valid)) + _, _ = writer.Write([]byte("data: {not-json}\n")) + <-finished + }() + defer func() { + close(finished) + _ = writer.Close() + }() + + usage, maxAPIError := GeminiResponsesStreamHandler(c, info, &http.Response{ + StatusCode: http.StatusOK, + Body: reader, + }) + require.NotNil(t, usage) + require.NotNil(t, maxAPIError) + require.True(t, types.IsSkipRetryError(maxAPIError)) + require.NotContains(t, recorder.Body.String(), "response.completed") +} + func TestGeminiTextGenerationHandlerUsesEstimatedPromptTokensWhenUsagePromptMissing(t *testing.T) { t.Parallel() diff --git a/relay/channel/gemini/relay_responses.go b/relay/channel/gemini/relay_responses.go index a5762e46..9342c07b 100644 --- a/relay/channel/gemini/relay_responses.go +++ b/relay/channel/gemini/relay_responses.go @@ -5,6 +5,7 @@ import ( "fmt" "io" "net/http" + "strings" "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/constant" @@ -16,6 +17,7 @@ import ( "github.com/MAX-API-Next/MAX-API/service/openaicompat" "github.com/MAX-API-Next/MAX-API/types" "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" ) func GeminiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.MaxAPIError) { @@ -54,7 +56,7 @@ func GeminiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *h usage := buildUsageFromGeminiResponse(c, info, &geminiResponse) chatResp.Usage = usage - responsesResp, responsesUsage, err := service.ChatCompletionsResponseToResponsesResponse(chatResp, helper.GetResponseID(c)) + responsesResp, responsesUsage, err := service.ChatCompletionsResponseToResponsesResponseWithCustomTools(chatResp, helper.GetResponseID(c), service.ResponsesCustomToolNames(info)) if err != nil { return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) } @@ -74,14 +76,22 @@ func GeminiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *h } func GeminiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.MaxAPIError) { + info.EnableFirstResultTracking() responseID := helper.GetResponseID(c) created := common.GetTimestamp() state := openaicompat.NewChatToResponsesStreamState(responseID, info.UpstreamModelName) + state.CustomToolNames = service.ResponsesCustomToolNames(info) state.Created = created finishReason := constant.FinishReasonStop toolCallIndexByChoice := make(map[int]map[string]int) nextToolCallIndexByChoice := make(map[int]int) var streamErr *types.MaxAPIError + terminalResponseSeen := false + outputDelivered := false + var deliveredText strings.Builder + var reportedUsage *dto.Usage + reportedOutputUsage := false + functionCalls := newGeminiStreamFunctionCalls() sendEvent := func(event openaicompat.ChatToResponsesStreamEvent) bool { data, err := common.Marshal(event.Payload) @@ -89,10 +99,18 @@ func GeminiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, r streamErr = types.NewOpenAIError(err, types.ErrorCodeJsonMarshalFailed, http.StatusInternalServerError) return false } - if err := helper.ResponseChunkData(c, dto.ResponsesStreamResponse{Type: event.Type}, string(data)); err != nil { - streamErr = types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError) + delivered, writeErr := helper.ResponseChunkDataWithDelivery(c, dto.ResponsesStreamResponse{Type: event.Type}, string(data)) + if delivered && helper.ResponsesStreamEventHasOutput(event.Payload) { + outputDelivered = true + deliveredText.WriteString(event.Payload.Delta) + } + if writeErr != nil { + streamErr = types.NewOpenAIError(writeErr, types.ErrorCodeBadResponse, http.StatusInternalServerError) return false } + if delivered && helper.ResponsesStreamEventHasFirstResult(event.Payload) { + info.SetFirstResultTime() + } return true } sendChunk := func(chunk *dto.ChatCompletionsStreamResponse) bool { @@ -110,6 +128,13 @@ func GeminiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, r } usage, err := geminiStreamHandler(c, info, resp, func(data string, geminiResponse *dto.GeminiChatResponse) bool { + if prepareErr := functionCalls.prepare(geminiResponse); prepareErr != nil { + streamErr = types.NewError(prepareErr, types.ErrorCodeBadResponseBody) + return false + } + if hasGeminiTerminalFinishReason(geminiResponse) { + terminalResponseSeen = true + } response, isStop := streamResponseGeminiChat2OpenAI(geminiResponse) response.Id = responseID response.Created = created @@ -118,6 +143,9 @@ func GeminiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, r if response.IsToolCall() { finishReason = constant.FinishReasonToolCalls } + if isStop { + terminalResponseSeen = true + } for choiceIdx := range response.Choices { choiceKey := response.Choices[choiceIdx].Index for toolIdx := range response.Choices[choiceIdx].Delta.ToolCalls { @@ -144,32 +172,129 @@ func GeminiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, r if !sendChunk(response) { return false } + inputReported := gjson.Get(data, "usageMetadata.promptTokenCount").Type == gjson.Number + outputReported := gjson.Get(data, "usageMetadata.candidatesTokenCount").Type == gjson.Number + if inputReported || outputReported || dto.HasGeminiUsageMetadataTokens(geminiResponse.GetUsageMetadata()) { + if metadata := geminiResponse.GetUsageMetadata(); metadata != nil { + fallbackPrompt := 0 + if !inputReported { + fallbackPrompt = info.GetEstimatePromptTokens() + } + mapped := buildUsageFromGeminiMetadata(*metadata, fallbackPrompt) + // A numeric zero completion count must not be inferred from total. + if outputReported { + mapped.CompletionTokens = metadata.CandidatesTokenCount + metadata.ThoughtsTokenCount + } + if !outputReported { + promptTokens := mapped.PromptTokens + patchGeminiZeroCompletionUsage(c, info, &mapped, deliveredText.String(), 0) + if inputReported { + mapped.PromptTokens = promptTokens + mapped.TotalTokens = mapped.PromptTokens + mapped.CompletionTokens + if mapped.BillingUsage != nil && mapped.BillingUsage.Estimated { + mapped.BillingUsage.GeminiUsageMetadata.PromptTokenCount = metadata.PromptTokenCount + mapped.BillingUsage = dto.NewEstimatedGeminiChatBillingUsage(&mapped) + } + } + } + if mapped.BillingUsage == nil { + copyMetadata := *metadata + mapped.BillingUsage = &dto.BillingUsage{Source: dto.BillingUsageSourceGeminiChat, Semantic: dto.BillingUsageSemanticGemini, GeminiUsageMetadata: ©Metadata} + } + if mapped.PromptTokens != metadata.PromptTokenCount+metadata.ToolUsePromptTokenCount { + mapped.TotalTokens = mapped.PromptTokens + mapped.CompletionTokens + mapped.BillingUsage = dto.NewEstimatedGeminiChatBillingUsage(&mapped) + } + if mapped.BillingUsage != nil { + mapped.BillingUsage.TokenCountsReported = inputReported && outputReported + } + reportedUsage = &mapped + reportedOutputUsage = outputReported + } + } if isStop { return sendChunk(helper.GenerateStopResponse(responseID, created, info.UpstreamModelName, finishReason)) } return true }) if err != nil { - return usage, err + // The shared scanner's timeout may have no usage return even though + // this bridge delivered whitespace and retained provider usage below. + streamErr = err } - if streamErr != nil { - return nil, streamErr + if reportedUsage != nil { + usage = reportedUsage } - if usage != nil { state.Usage = openaicompat.UsageFromChatUsage(usage) } - if state.Usage == nil || state.Usage.TotalTokens == 0 { - estimated := service.ResponseText2Usage(c, state.UsageText(), info.UpstreamModelName, info.GetEstimatePromptTokens()) - state.Usage = openaicompat.UsageFromChatUsage(estimated) - if usage == nil || usage.TotalTokens == 0 { - usage = estimated + refreshUsage := func() { + if reportedUsage == nil { + usage = service.ResponseText2Usage(c, deliveredText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) + attachEstimatedGeminiBillingUsage(usage) + } else if !reportedOutputUsage { + // Partial metadata can report input (including zero) while leaving + // output unknown. Preserve that input and estimate the delivered output. + mapped := *reportedUsage + estimated := service.ResponseText2Usage(c, deliveredText.String(), info.UpstreamModelName, mapped.PromptTokens) + if estimated.CompletionTokens > mapped.CompletionTokens { + mapped.CompletionTokens = estimated.CompletionTokens + } + mapped.TotalTokens = mapped.PromptTokens + mapped.CompletionTokens + mapped.BillingUsage = dto.NewEstimatedGeminiChatBillingUsage(&mapped) + usage = &mapped + } + if usage != nil { + state.Usage = openaicompat.UsageFromChatUsage(usage) } } - for _, event := range openaicompat.FinalizeChatCompletionsStreamToResponses(state) { - if !sendEvent(event) { + if streamErr != nil { + if c.Writer.Written() { + streamErr = types.NewError(streamErr, streamErr.GetErrorCode(), types.ErrOptionWithSkipRetry()) + } + if !outputDelivered { return nil, streamErr } + refreshUsage() + return usage, streamErr + } + finalizeStream := func() bool { + for _, event := range openaicompat.FlushChatCompletionsStreamToResponsesOutput(state) { + if !sendEvent(event) { + refreshUsage() + return false + } + } + refreshUsage() + for _, event := range openaicompat.FinalizeChatCompletionsStreamToResponses(state) { + if !sendEvent(event) { + return false + } + } + return true + } + if !terminalResponseSeen || len(functionCalls.pending) > 0 { + reason := "upstream_eof" + if info.StreamStatus != nil && info.StreamStatus.IsAbnormalEnd() { + reason = string(info.StreamStatus.EndReason) + } + state.MarkIncomplete(reason) + if info.StreamStatus == nil || info.StreamStatus.EndReason == relaycommon.StreamEndReasonEOF { + if !finalizeStream() { + return usage, types.NewError(streamErr, streamErr.GetErrorCode(), types.ErrOptionWithSkipRetry()) + } + } + if !outputDelivered { + return nil, openaicompat.NewResponsesStreamIncompleteError(reason) + } + return usage, openaicompat.NewResponsesStreamIncompleteError(reason) + } + if info.StreamStatus != nil && info.StreamStatus.IsAbnormalEnd() { + return usage, openaicompat.NewResponsesStreamIncompleteError(string(info.StreamStatus.EndReason)) + } + + if !finalizeStream() { + return usage, types.NewError(streamErr, streamErr.GetErrorCode(), types.ErrOptionWithSkipRetry()) } if service.ResponseAuditEnabled() { service.SetRelayResponseAuditContent(info, state.UsageText()) diff --git a/relay/channel/gemini/responses_revalidation_test.go b/relay/channel/gemini/responses_revalidation_test.go new file mode 100644 index 00000000..87401d30 --- /dev/null +++ b/relay/channel/gemini/responses_revalidation_test.go @@ -0,0 +1,213 @@ +package gemini + +import ( + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/MAX-API-Next/MAX-API/dto" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func runGeminiResponsesRevalidation(t *testing.T, frames []string, custom bool) (*dto.Usage, *types.MaxAPIError, string) { + t.Helper() + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + info := &relaycommon.RelayInfo{IsStream: true, DisablePing: true, OriginModelName: "gemini-test", + RelayFormat: types.RelayFormatOpenAIResponses, ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gemini-test"}} + info.SetEstimatePromptTokens(10) + if custom { + info.Request = &dto.OpenAIResponsesRequest{Tools: []byte(`[{"type":"custom","name":"lookup"}]`)} + } + resp := &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader("data: " + strings.Join(frames, "\n\ndata: ") + "\n\n"))} + usage, err := GeminiResponsesStreamHandler(c, info, resp) + return usage, err, w.Body.String() +} + +func TestGeminiResponsesToolSnapshotsEmitCompleteArgumentsOnce(t *testing.T) { + for _, custom := range []bool{false, true} { + for _, emptyTerminal := range []bool{false, true} { + t.Run(fmt.Sprintf("custom_%t/empty_terminal_%t", custom, emptyTerminal), func(t *testing.T) { + args := `{"q":"max"}` + if custom { + args = `{"input":"patch body"}` + } + completeCall := fmt.Sprintf(`{"id":"call-1","name":"lookup","args":%s,"willContinue":false}`, args) + terminalCall := completeCall + if emptyTerminal { + terminalCall = `{"willContinue":false}` + } + frames := []string{ + fmt.Sprintf(`{"candidates":[{"index":0,"content":{"parts":[{"functionCall":{"id":"call-1","name":"lookup","args":%s,"willContinue":true}}]}}]}`, args), + fmt.Sprintf(`{"candidates":[{"index":0,"content":{"parts":[{"functionCall":%s}]}}]}`, terminalCall), + fmt.Sprintf(`{"candidates":[{"index":0,"content":{"parts":[{"functionCall":%s}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":2,"totalTokenCount":12}}`, completeCall), + } + _, apiErr, body := runGeminiResponsesRevalidation(t, frames, custom) + require.Nil(t, apiErr) + var completed gjson.Result + var doneCount int + for _, line := range strings.Split(body, "\n") { + if !strings.HasPrefix(line, "data: ") { + continue + } + data := strings.TrimPrefix(line, "data: ") + if gjson.Get(data, "type").String() == "response.completed" { + completed = gjson.Get(data, "response") + } + if strings.HasSuffix(gjson.Get(data, "type").String(), "arguments.done") || strings.HasSuffix(gjson.Get(data, "type").String(), "input.done") { + doneCount++ + } + } + require.True(t, completed.Exists(), body) + require.Equal(t, 1, doneCount, body) + require.Len(t, completed.Get("output").Array(), 1) + item := completed.Get("output.0") + require.Equal(t, "call-1", item.Get("call_id").String()) + if custom { + require.Equal(t, "patch body", item.Get("input").String()) + } else { + require.JSONEq(t, args, item.Get("arguments").String()) + } + }) + } + } +} + +func TestGeminiResponsesUsageDistinguishesMissingNullAndZero(t *testing.T) { + for _, input := range []struct { + name, field string + expected int + }{ + {"missing", "", 10}, {"null", `"promptTokenCount":null,`, 10}, {"zero", `"promptTokenCount":0,`, 0}, + } { + t.Run(input.name, func(t *testing.T) { + usage, err, _ := runGeminiResponsesRevalidation(t, []string{ + `{"candidates":[{"content":{"parts":[{"text":"hello"}]},"finishReason":"STOP"}],"usageMetadata":{` + input.field + `"candidatesTokenCount":2}}`, + }, false) + require.Nil(t, err) + require.Equal(t, input.expected, usage.PromptTokens) + require.Equal(t, 2, usage.CompletionTokens) + require.NotNil(t, usage.BillingUsage) + require.Equal(t, input.expected, usage.BillingUsage.GeminiUsageMetadata.PromptTokenCount) + require.Equal(t, input.expected != 0, usage.BillingUsage.Estimated) + }) + } +} + +func TestGeminiResponsesToolSnapshotsRejectConflictsAndIncompleteCalls(t *testing.T) { + for _, scenario := range []struct { + name string + frames []string + }{ + {"changed_arguments", []string{ + `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-1","name":"lookup","args":{"q":"a"}}}]}}]}`, + `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-1","name":"lookup","args":{"q":"b"}}}]},"finishReason":"STOP"}]}`, + }}, + {"changed_name", []string{ + `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-1","name":"lookup","args":{},"willContinue":true}}]}}]}`, + `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-1","name":"other","args":{},"willContinue":false}}]},"finishReason":"STOP"}]}`, + }}, + {"unfinished_at_stop", []string{`{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-1","name":"lookup","args":{},"willContinue":true}}]},"finishReason":"STOP"}]}`}}, + {"unfinished_at_eof", []string{`{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-1","name":"lookup","args":{},"willContinue":true}}]}}]}`}}, + {"unfinished_after_other_candidate_stop", []string{ + `{"candidates":[{"index":0,"content":{"parts":[{"text":"partial"}]},"finishReason":"STOP"}]}`, + `{"candidates":[{"index":1,"content":{"parts":[{"functionCall":{"id":"call-1","name":"lookup","args":{},"willContinue":true}}]}}]}`, + }}, + } { + t.Run(scenario.name, func(t *testing.T) { + _, err, body := runGeminiResponsesRevalidation(t, scenario.frames, false) + require.NotNil(t, err) + require.NotContains(t, body, "event: response.completed") + }) + } +} + +func TestGeminiResponsesIndependentCallsAndLatestPendingSnapshot(t *testing.T) { + _, err, body := runGeminiResponsesRevalidation(t, []string{ + `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-1","name":"lookup","args":{"q":"old"},"willContinue":true}}]}}]}`, + `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-1","name":"lookup","args":{"q":"new"},"willContinue":false}}]}}]}`, + `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-2","name":"lookup","args":{"q":"second"}}}]},"finishReason":"STOP"}]}`, + }, false) + require.Nil(t, err) + for _, line := range strings.Split(body, "\n") { + data := strings.TrimPrefix(line, "data: ") + if gjson.Get(data, "type").String() != "response.completed" { + continue + } + output := gjson.Get(data, "response.output").Array() + require.Len(t, output, 2) + require.Equal(t, "call-1", output[0].Get("call_id").String()) + require.JSONEq(t, `{"q":"new"}`, output[0].Get("arguments").String()) + require.Equal(t, "call-2", output[1].Get("call_id").String()) + require.JSONEq(t, `{"q":"second"}`, output[1].Get("arguments").String()) + return + } + t.Fatal("missing response.completed: " + body) +} + +func TestGeminiResponsesMissingOutputDoesNotReplaceNumericZeroInput(t *testing.T) { + for _, field := range []string{"", `,"candidatesTokenCount":null`} { + usage, err, _ := runGeminiResponsesRevalidation(t, []string{ + `{"candidates":[{"content":{"parts":[{"text":"hello"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":0` + field + `}}`, + }, false) + require.Nil(t, err) + require.Zero(t, usage.PromptTokens) + require.Positive(t, usage.CompletionTokens) + require.True(t, usage.BillingUsage.Estimated) + require.Zero(t, usage.BillingUsage.GeminiUsageMetadata.PromptTokenCount) + } +} + +func TestGeminiChatToolSnapshotsFinishOnceAndRejectPendingEOF(t *testing.T) { + for _, scenario := range []struct{ complete, combined bool }{{false, false}, {true, false}, {true, true}} { + complete := scenario.complete + t.Run(fmt.Sprintf("complete_%t/combined_%t", complete, scenario.combined), func(t *testing.T) { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + info := &relaycommon.RelayInfo{IsStream: true, DisablePing: true, RelayFormat: types.RelayFormatOpenAI, OriginModelName: "gemini-test", + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gemini-test"}} + frames := []string{`{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-1","name":"lookup","args":{"q":"max"},"willContinue":true}}]}}]}`} + if complete { + frames = append(frames, + `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-1","name":"lookup","args":{"q":"max"},"willContinue":false}}]}}]}`, + `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-1","name":"lookup","args":{"q":"max"}}}]},"finishReason":"STOP"}]}`) + } + if scenario.combined { + frames[1] = strings.Replace(frames[1], `"content":`, `"finishReason":"STOP","content":`, 1) + frames = frames[:2] + } + _, apiErr := GeminiChatStreamHandler(c, info, &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader("data: " + strings.Join(frames, "\n\ndata: ") + "\n\n"))}) + if !complete { + require.NotNil(t, apiErr) + return + } + require.Nil(t, apiErr) + var arguments strings.Builder + finishCount := 0 + for _, line := range strings.Split(w.Body.String(), "\n") { + data := strings.TrimPrefix(line, "data: ") + for _, choice := range gjson.Get(data, "choices").Array() { + for _, call := range choice.Get("delta.tool_calls").Array() { + arguments.WriteString(call.Get("function.arguments").String()) + } + if choice.Get("finish_reason").Type == gjson.String { + finishCount++ + } + } + } + require.JSONEq(t, `{"q":"max"}`, arguments.String()) + require.Equal(t, 1, finishCount, w.Body.String()) + }) + } +} diff --git a/relay/channel/gemini/stream_function_calls.go b/relay/channel/gemini/stream_function_calls.go new file mode 100644 index 00000000..1ecd3bc4 --- /dev/null +++ b/relay/channel/gemini/stream_function_calls.go @@ -0,0 +1,108 @@ +package gemini + +import ( + "fmt" + "strings" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/dto" +) + +type geminiCallPosition struct { + candidate int64 + part int +} +type geminiCallIdentity struct { + candidate int64 + id string +} + +// Gemini args are complete JSON snapshots, unlike Chat arguments deltas. +// Buffer willContinue sequences and emit each stable call ID only once. +type geminiStreamFunctionCalls struct { + pending map[geminiCallPosition]dto.FunctionCall + completed map[geminiCallIdentity]dto.FunctionCall +} + +func newGeminiStreamFunctionCalls() *geminiStreamFunctionCalls { + return &geminiStreamFunctionCalls{pending: make(map[geminiCallPosition]dto.FunctionCall), completed: make(map[geminiCallIdentity]dto.FunctionCall)} +} + +func (s *geminiStreamFunctionCalls) prepare(response *dto.GeminiChatResponse) error { + for candidateIndex := range response.Candidates { + candidate := &response.Candidates[candidateIndex] + parts := make([]dto.GeminiPart, 0, len(candidate.Content.Parts)) + for partIndex, part := range candidate.Content.Parts { + if part.FunctionCall == nil { + parts = append(parts, part) + continue + } + position := geminiCallPosition{candidate.Index, partIndex} + call := *part.FunctionCall + call.ID = strings.TrimSpace(call.ID) + call.FunctionName = strings.TrimSpace(call.FunctionName) + if previous, exists := s.pending[position]; exists { + if call.ID != "" && previous.ID != "" && call.ID != previous.ID { + return fmt.Errorf("Gemini function call changed id while incomplete") + } + if call.FunctionName != "" && previous.FunctionName != "" && call.FunctionName != previous.FunctionName { + return fmt.Errorf("Gemini function call changed name while incomplete") + } + if call.ID == "" { + call.ID = previous.ID + } + if call.FunctionName == "" { + call.FunctionName = previous.FunctionName + } + if call.Arguments == nil { + call.Arguments = previous.Arguments + } + } + identity := geminiCallIdentity{candidate.Index, call.ID} + if previous, exists := s.completed[identity]; call.ID != "" && exists { + if call.FunctionName != "" && call.FunctionName != previous.FunctionName { + return fmt.Errorf("Gemini completed function call changed name") + } + if call.Arguments != nil { + oldArgs, err := common.Marshal(previous.Arguments) + if err != nil { + return err + } + newArgs, err := common.Marshal(call.Arguments) + if err != nil { + return err + } + if string(oldArgs) != string(newArgs) { + return fmt.Errorf("Gemini completed function call changed arguments") + } + } + continue + } + if call.WillContinue != nil && *call.WillContinue { + s.pending[position] = call + continue + } + if call.FunctionName == "" { + return fmt.Errorf("Gemini completed function call is missing name") + } + delete(s.pending, position) + if call.Arguments == nil { + call.Arguments = map[string]any{} + } + if call.ID != "" { + s.completed[identity] = call + } + part.FunctionCall = &call + parts = append(parts, part) + } + candidate.Content.Parts = parts + if candidate.FinishReason != nil && strings.TrimSpace(*candidate.FinishReason) != "" { + for position := range s.pending { + if position.candidate == candidate.Index { + return fmt.Errorf("Gemini finished with an incomplete function call") + } + } + } + } + return nil +} diff --git a/relay/channel/openai/chat_stream.go b/relay/channel/openai/chat_stream.go index cfc02fa5..1075b68c 100644 --- a/relay/channel/openai/chat_stream.go +++ b/relay/channel/openai/chat_stream.go @@ -91,6 +91,71 @@ func (frame *chatStreamFrame) hasPayload() bool { return false } +// Keep timing separate from the conservative empty-retry predicate. Provider +// metadata and fields dropped by typed formatting are not delivered results. +func (frame *chatStreamFrame) hasFirstResult(formatted bool) bool { + for _, choice := range frame.Choices { + if choice.FinishReason == "content_filter" { + return true + } + for _, key := range []string{"content", "reasoning", "reasoning_content"} { + value := choice.Delta[key] + if text, ok := value.(string); ok && strings.TrimSpace(text) != "" { + return true + } + if parts, ok := value.([]any); ok && key == "content" && !formatted { + for _, part := range parts { + if content, ok := part.(map[string]any); ok { + for _, field := range []string{"text", "refusal", "image_url", "data"} { + if chatStreamValueHasPayload(content[field]) { + return true + } + } + } + } + } + } + if calls, ok := choice.Delta["tool_calls"].([]any); ok { + for _, call := range calls { + if chatStreamCallHasFirstResult(call, formatted) { + return true + } + } + } + if !formatted { + if chatStreamValueHasPayload(choice.Delta["refusal"]) || chatStreamCallHasFirstResult(choice.Delta["function_call"], false) { + return true + } + if audio, ok := choice.Delta["audio"].(map[string]any); ok && + (chatStreamValueHasPayload(audio["data"]) || chatStreamValueHasPayload(audio["transcript"])) { + return true + } + } + } + return false +} + +func chatStreamCallHasFirstResult(value any, formatted bool) bool { + call, ok := value.(map[string]any) + if !ok { + return false + } + for _, field := range []string{"id", "name", "arguments"} { + if chatStreamValueHasPayload(call[field]) { + return true + } + } + if function, ok := call["function"].(map[string]any); ok && chatStreamCallHasFirstResult(function, false) { + return true + } + if !formatted { + if custom, ok := call["custom"].(map[string]any); ok { + return chatStreamValueHasPayload(custom["name"]) || chatStreamValueHasPayload(custom["input"]) + } + } + return false +} + func (frame *chatStreamFrame) upstreamError() *types.MaxAPIError { if len(frame.Error) == 0 || strings.TrimSpace(string(frame.Error)) == "null" { return nil @@ -130,6 +195,7 @@ func chatStreamRequestErrorIsPermanent(upstream types.OpenAIError) bool { // policy. A partial response returns usage AND a SkipRetry error; TextHelper // settles that usage before returning the terminal error to the controller. func oaiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.MaxAPIError) { + info.EnableFirstResultTracking() if info.RetryIndex > 0 && !c.Writer.Written() { info.StreamStatus = nil info.SendResponseCount, info.ReceivedResponseCount = 0, 0 @@ -221,6 +287,9 @@ func oaiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *htt sr.Stop(streamErr) return } + if frame.hasFirstResult(info.ChannelSetting.ForceFormat || info.ChannelSetting.ThinkingToContent) { + info.SetFirstResultTime() + } hasPayload = hasPayload || visible recordUsage(&frame, data) observeOpenAIStreamToolCalls(observer, data) @@ -233,7 +302,10 @@ func oaiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *htt // race the final retry decision. Cancellation is never an empty completion. if c.Request.Context().Err() != nil { streamErr = types.NewOpenAIError(c.Request.Context().Err(), types.ErrorCodeBadResponse, http.StatusBadGateway, types.ErrOptionWithSkipRetry()) - } else if streamErr == nil && !info.StreamStatus.IsNormalEnd() { + } else if streamErr == nil { + streamErr = helper.FirstResultTimeoutError(c, info) + } + if streamErr == nil && !info.StreamStatus.IsNormalEnd() { streamErr = types.NewOpenAIError(fmt.Errorf("upstream chat stream ended: %s", info.StreamStatus.Summary()), types.ErrorCodeBadResponse, http.StatusBadGateway) } if streamErr != nil { diff --git a/relay/channel/openai/chat_via_responses.go b/relay/channel/openai/chat_via_responses.go index 73eb61b4..a233046f 100644 --- a/relay/channel/openai/chat_via_responses.go +++ b/relay/channel/openai/chat_via_responses.go @@ -111,11 +111,19 @@ func OaiResponsesToChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp } func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.MaxAPIError) { + info.EnableFirstResultTracking() if resp == nil || resp.Body == nil { return nil, types.NewOpenAIError(fmt.Errorf("invalid response"), types.ErrorCodeBadResponse, http.StatusInternalServerError) } defer service.CloseResponseBodyGracefully(resp) + writeObserver := &chatStreamWriteObserver{ResponseWriter: c.Writer} + c.Writer = writeObserver + defer func() { + if writeObserver.err == nil { + c.Writer = writeObserver.ResponseWriter + } + }() responseId := helper.GetResponseID(c) createAt := time.Now().Unix() @@ -130,6 +138,9 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo sentStop bool sawToolCall bool hasVisiblePayload bool + outputDelivered bool + deliveredUsageText strings.Builder + usageReported bool emptyCompletionRecorded bool streamErr *types.MaxAPIError ) @@ -151,12 +162,39 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo if chunk == nil { return true } + recordFirstResult := func() { + firstResult := false + for _, choice := range chunk.Choices { + content, reasoning := choice.Delta.GetContentString(), choice.Delta.GetReasoningContent() + if content != "" || reasoning != "" || len(choice.Delta.ToolCalls) > 0 { + outputDelivered = true + deliveredUsageText.WriteString(content) + deliveredUsageText.WriteString(reasoning) + for _, tool := range choice.Delta.ToolCalls { + deliveredUsageText.WriteString(tool.Function.Name) + deliveredUsageText.WriteString(tool.Function.Arguments) + } + } + if strings.TrimSpace(choice.Delta.GetContentString()) != "" || + strings.TrimSpace(choice.Delta.GetReasoningContent()) != "" || len(choice.Delta.ToolCalls) > 0 { + firstResult = true + } + } + if firstResult { + info.SetFirstResultTime() + } + } if info.RelayFormat == types.RelayFormatOpenAI { info.SendResponseCount++ if err := helper.ObjectData(c, chunk); err != nil { streamErr = types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError) return false } + if writeObserver.err != nil { + streamErr = types.NewOpenAIError(writeObserver.err, types.ErrorCodeBadResponse, http.StatusInternalServerError) + return false + } + recordFirstResult() return true } @@ -169,6 +207,11 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo streamErr = types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError) return false } + if writeObserver.err != nil { + streamErr = types.NewOpenAIError(writeObserver.err, types.ErrorCodeBadResponse, http.StatusInternalServerError) + return false + } + recordFirstResult() return true } @@ -327,9 +370,16 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo var streamResp dto.ResponsesStreamResponse if err := common.UnmarshalJsonStr(data, &streamResp); err != nil { logger.LogError(c, "failed to unmarshal responses stream event: "+err.Error()) - sr.Error(err) + streamErr = types.NewError(err, types.ErrorCodeBadResponseBody) + sr.Stop(streamErr) return } + if streamResp.Response != nil && (helper.HasUsageTokenFields(data, "response.usage", "input_tokens", "output_tokens") || + helper.HasUsageTokenFields(data, "response.usage", "prompt_tokens", "completion_tokens")) { + applyResponsesUsage(usage, streamResp.Response.Usage) + usage.BillingUsage = dto.NewReportedOpenAIResponsesBillingUsage(usage) + usageReported = true + } switch streamResp.Type { case "response.created": @@ -475,7 +525,15 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo case "response.function_call_arguments.done", "response.custom_tool_call_input.done": - case "response.completed": + case "response.completed", "response.done", "response.incomplete": + if streamResp.Type == "response.incomplete" { + if streamResp.Response == nil { + streamResp.Response = &dto.OpenAIResponsesResponse{} + } + if len(streamResp.Response.Status) == 0 { + streamResp.Response.Status = []byte(`"incomplete"`) + } + } completedOutputTypes := []string(nil) if streamResp.Response != nil { if terminalErr := responsesTerminalError(streamResp.Response, http.StatusInternalServerError, info.SendResponseCount > 0); terminalErr != nil { @@ -491,7 +549,7 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo if streamResp.Response.CreatedAt != 0 { createAt = int64(streamResp.Response.CreatedAt) } - if streamResp.Response.Usage != nil { + if streamResp.Response.Usage != nil && !usageReported { usage = openaicompat.UsageFromResponsesUsage(streamResp.Response.Usage) } if outputText.Len() == 0 { @@ -565,7 +623,9 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo info.ClaudeConvertInfo.Usage = usage } finishReason := "stop" - if sawToolCall { + if mappedReason, ok := openaicompat.ResponsesFinishReasonFromStatus(streamResp.Response); ok { + finishReason = mappedReason + } else if sawToolCall { finishReason = "tool_calls" } stop := helper.GenerateStopResponse(responseId, createAt, model, finishReason) @@ -586,11 +646,26 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo } }) - if apiErr := helper.FirstResultTimeoutError(info); apiErr != nil { + if apiErr := helper.FirstResultTimeoutError(c, info); apiErr != nil { + if outputDelivered && info.RelayFormat == types.RelayFormatOpenAI { + if !usageReported { + usage = service.ResponseText2Usage(c, deliveredUsageText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) + } + return usage, apiErr + } return nil, apiErr } if streamErr != nil { + if c.Writer.Written() { + types.ErrOptionWithSkipRetry()(streamErr) + } + if outputDelivered && info.RelayFormat == types.RelayFormatOpenAI { + if !usageReported { + usage = service.ResponseText2Usage(c, deliveredUsageText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) + } + return usage, streamErr + } return nil, streamErr } if !hasVisiblePayload && !emptyCompletionRecorded { @@ -599,8 +674,22 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo return nil, retryErr } } + if !sentStop { + reason := "upstream_eof" + if info.StreamStatus != nil && info.StreamStatus.IsAbnormalEnd() { + reason = string(info.StreamStatus.EndReason) + } + apiErr := openaicompat.NewResponsesStreamIncompleteError(reason) + if outputDelivered && info.RelayFormat == types.RelayFormatOpenAI { + if !usageReported { + usage = service.ResponseText2Usage(c, deliveredUsageText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) + } + return usage, apiErr + } + return nil, apiErr + } - if usage.TotalTokens == 0 { + if !usageReported && usage.TotalTokens == 0 { usage = service.ResponseText2Usage(c, usageText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) } diff --git a/relay/channel/openai/empty_completion.go b/relay/channel/openai/empty_completion.go index 161d0bba..bb7eec8e 100644 --- a/relay/channel/openai/empty_completion.go +++ b/relay/channel/openai/empty_completion.go @@ -9,6 +9,7 @@ import ( "github.com/MAX-API-Next/MAX-API/constant" "github.com/MAX-API-Next/MAX-API/dto" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/relay/helper" "github.com/MAX-API-Next/MAX-API/service" "github.com/MAX-API-Next/MAX-API/types" @@ -232,7 +233,12 @@ func responsesStreamHasVisibleOutput(streamResp *dto.ResponsesStreamResponse) bo "response.reasoning_summary_text.delta", "response.function_call_arguments.delta": return strings.TrimSpace(streamResp.Delta) != "" + case "response.custom_tool_call_input.delta": + return strings.TrimSpace(streamResp.Delta) != "" case dto.ResponsesOutputTypeItemAdded, dto.ResponsesOutputTypeItemDone: + if streamResp.Item != nil && streamResp.Item.Type == dto.BuildInCallFileSearchCall { + return helper.ResponsesStreamEventHasOutput(*streamResp) + } return responsesOutputHasVisiblePayload(streamResp.Item) case "response.completed": return responsesHasVisibleOutput(streamResp.Response) @@ -246,7 +252,7 @@ func responsesOutputHasVisiblePayload(output *dto.ResponsesOutput) bool { return false } switch output.Type { - case "function_call", dto.ResponsesOutputTypeImageGenerationCall, dto.BuildInCallWebSearchCall: + case "function_call", dto.BuildInCallCustomToolCall, dto.ResponsesOutputTypeImageGenerationCall, dto.BuildInCallWebSearchCall: return true case "message": if output.Role != "" && output.Role != "assistant" { diff --git a/relay/channel/openai/empty_completion_test.go b/relay/channel/openai/empty_completion_test.go index c780b62d..3db7fbfc 100644 --- a/relay/channel/openai/empty_completion_test.go +++ b/relay/channel/openai/empty_completion_test.go @@ -258,9 +258,9 @@ func TestOaiResponsesStreamHandlerFlushesImmediatelyWhenEmptyRetryDisabled(t *te }, } - _, err := OaiResponsesStreamHandler(c, info, resp) - if err != nil { - t.Fatalf("unexpected error: %v", err) + usage, err := OaiResponsesStreamHandler(c, info, resp) + if err == nil || !types.IsSkipRetryError(err) || usage != nil { + t.Fatalf("metadata-only EOF must return no billable partial usage and forbid replay: usage=%+v, err=%v", usage, err) } if body := recorder.Body.String(); !strings.Contains(body, "response.created") { t.Fatalf("response body = %q, want forwarded pre-content event", body) @@ -299,9 +299,9 @@ func TestOaiResponsesStreamHandlerCapsPendingEvents(t *testing.T) { }, } - _, err := OaiResponsesStreamHandler(c, info, resp) - if err != nil { - t.Fatalf("unexpected error after pending cap forces forwarding: %v", err) + usage, err := OaiResponsesStreamHandler(c, info, resp) + if err == nil || !types.IsSkipRetryError(err) || usage != nil { + t.Fatalf("flushing pending metadata must not turn an incomplete stream into a billable result: usage=%+v, err=%v", usage, err) } if body := recorder.Body.String(); !strings.Contains(body, "response.created") { t.Fatalf("response body = %q, want forwarded events after cap", body) diff --git a/relay/channel/openai/file_search_stream_delivery_test.go b/relay/channel/openai/file_search_stream_delivery_test.go new file mode 100644 index 00000000..200d07e6 --- /dev/null +++ b/relay/channel/openai/file_search_stream_delivery_test.go @@ -0,0 +1,87 @@ +package openai + +import ( + "fmt" + "io" + "net/http" + "strings" + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/dto" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/relay/helper" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/stretchr/testify/require" +) + +func TestOaiFileSearchCompletedItemsRetainPartialUsage(t *testing.T) { + oldEmpty, oldRetry := common.EmptyCompletionRetryEnabled, common.RetryTimes + common.RetryTimes = 2 + t.Cleanup(func() { common.EmptyCompletionRetryEnabled, common.RetryTimes = oldEmpty, oldRetry }) + for _, retry := range []bool{false, true} { + for _, status := range []string{"completed", ""} { + for _, malformed := range []bool{false, true} { + name := "legacy_status" + if status != "" { + name = status + } + t.Run(fmt.Sprintf("retry_%t/%s/malformed_%t", retry, name, malformed), func(t *testing.T) { + common.EmptyCompletionRetryEnabled = retry + c, info := newResponsesUsageTestContext() + info.ResponsesUsageInfo = &relaycommon.ResponsesUsageInfo{BuiltInTools: map[string]*relaycommon.BuildInToolInfo{ + dto.BuildInToolFileSearch: {ToolName: dto.BuildInToolFileSearch}, + }} + item := map[string]any{"type": dto.BuildInCallFileSearchCall, "id": "fs_test", "status": status, + "results": []any{map[string]any{"text": "synthetic-search-result"}}} + data, err := common.Marshal(map[string]any{"type": dto.ResponsesOutputTypeItemDone, "output_index": 0, "item": item}) + require.NoError(t, err) + body := "data: " + `{"type":"response.created","response":{"id":"resp_file_search","usage":{"input_tokens":4,"output_tokens":2}}}` + "\n\n" + body += "data: " + string(data) + "\n\ndata: " + string(data) + "\n\n" + if malformed { + body += "data: {broken\n\n" + } + usage, apiErr := OaiResponsesStreamHandler(c, info, &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(body))}) + require.NotNil(t, apiErr) + require.True(t, types.IsSkipRetryError(apiErr)) + require.NotNil(t, usage) + require.Equal(t, 4, usage.PromptTokens) + require.Equal(t, 2, usage.CompletionTokens) + require.True(t, info.HasRecordedChannelFirstResult()) + require.Equal(t, 1, info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolFileSearch].CallCount) + }) + } + } + } +} + +func TestOaiFileSearchLifecycleItemsRemainMetadata(t *testing.T) { + oldEmpty := common.EmptyCompletionRetryEnabled + common.EmptyCompletionRetryEnabled = false + t.Cleanup(func() { common.EmptyCompletionRetryEnabled = oldEmpty }) + for _, tc := range []struct{ name, eventType, status string }{ + {"added_running", dto.ResponsesOutputTypeItemAdded, "in_progress"}, + {"added_completed", dto.ResponsesOutputTypeItemAdded, "completed"}, + {"done_running", dto.ResponsesOutputTypeItemDone, "in_progress"}, + {"done_failed", dto.ResponsesOutputTypeItemDone, "failed"}, + {"done_cancelled", dto.ResponsesOutputTypeItemDone, "cancelled"}, + {"done_unknown", dto.ResponsesOutputTypeItemDone, "future_status"}, + {"nil_item", dto.ResponsesOutputTypeItemDone, ""}, + } { + t.Run(tc.name, func(t *testing.T) { + event := dto.ResponsesStreamResponse{Type: tc.eventType, Item: &dto.ResponsesOutput{Type: dto.BuildInCallFileSearchCall, Status: tc.status}} + if tc.name == "nil_item" { + event.Item = nil + } + require.False(t, helper.ResponsesStreamEventHasOutput(event)) + require.False(t, helper.ResponsesStreamEventHasFirstResult(event)) + data, err := common.Marshal(event) + require.NoError(t, err) + c, info := newResponsesUsageTestContext() + usage, apiErr := OaiResponsesStreamHandler(c, info, &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader("data: " + string(data) + "\n\n"))}) + require.NotNil(t, apiErr) + require.Nil(t, usage) + require.False(t, info.HasRecordedChannelFirstResult()) + }) + } +} diff --git a/relay/channel/openai/first_result_test.go b/relay/channel/openai/first_result_test.go new file mode 100644 index 00000000..7e3224f7 --- /dev/null +++ b/relay/channel/openai/first_result_test.go @@ -0,0 +1,143 @@ +package openai + +import ( + "io" + "net/http" + "strings" + "testing" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/dto" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/setting/operation_setting" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +// Use the real adapters and keep the provider open after metadata. The close +// is a bounded fallback: a broken watchdog must fail rather than hang the suite. +func stalledFirstResultStream(t *testing.T, frames []string) *http.Response { + t.Helper() + r, w := io.Pipe() + finished := make(chan struct{}) + t.Cleanup(func() { _ = r.Close(); _ = w.Close(); <-finished }) + go func() { + defer close(finished) + defer w.Close() + for _, frame := range frames { + if _, err := io.WriteString(w, "data: "+frame+"\n\n"); err != nil { + return + } + } + time.Sleep(1500 * time.Millisecond) + }() + return &http.Response{StatusCode: http.StatusOK, Body: r} +} + +func TestAdaptersMetadataDoesNotEndFirstResultWait(t *testing.T) { + created := `{"type":"response.created","response":{"id":"resp_test","status":"in_progress","output":[]}}` + tests := []struct { + name string + frames []string + buffer bool + handler func(*gin.Context, *relaycommon.RelayInfo, *http.Response) (*dto.Usage, *types.MaxAPIError) + }{ + {"chat_role", []string{chatRoleFrame}, false, OaiStreamHandler}, + {"chat_usage", []string{chatUsageFrame}, false, OaiStreamHandler}, + {"chat_buffer_limit", stringsToFrames(chatRoleFrame, maxPendingChatStreamEvents+1), true, OaiStreamHandler}, + {"responses_created", []string{created}, false, OaiResponsesStreamHandler}, + {"responses_buffer_limit", stringsToFrames(created, maxPendingResponsesStreamEvents), true, OaiResponsesStreamHandler}, + {"chat_to_responses_role", []string{chatRoleFrame}, false, OaiChatToResponsesStreamHandler}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + oldSetting := *operation_setting.GetMonitorSetting() + *operation_setting.GetMonitorSetting() = operation_setting.MonitorSetting{StreamingFirstResultTimeoutSeconds: 1} + t.Cleanup(func() { *operation_setting.GetMonitorSetting() = oldSetting }) + c, recorder, info := newChatStreamTest(t) + info.InitChannelMeta(c) + info.UpstreamModelName = "gpt-test" + common.EmptyCompletionRetryEnabled = test.buffer + usage, apiErr := test.handler(c, info, stalledFirstResultStream(t, test.frames)) + require.Equal(t, relaycommon.StreamEndReasonTimeout, info.StreamStatus.EndReason, + "metadata must not disable the first-result deadline") + require.NotNil(t, apiErr) + require.Equal(t, types.ErrorCodeChannelResponseTimeExceeded, apiErr.GetErrorCode()) + require.True(t, types.IsSkipRetryError(apiErr), "forwarded metadata closes the replay window") + require.Nil(t, usage, "metadata must not create partial output usage") + require.NotEmpty(t, recorder.Body.String(), "metadata forwarding remains compatible") + }) + } +} + +func stringsToFrames(frame string, count int) []string { + frames := make([]string, count) + for i := range frames { + frames[i] = frame + } + return frames +} + +func TestNativeChatFirstResultAcceptsSupportedOutput(t *testing.T) { + for name, delta := range map[string]string{ + "text": `{"content":"hello"}`, "reasoning": `{"reasoning_content":"thinking"}`, + "refusal": `{"refusal":"cannot comply"}`, "audio": `{"audio":{"data":"YWJj"}}`, + "image": `{"content":[{"type":"image_url","image_url":{"url":"test-image"}}]}`, + "tool": `{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"lookup","arguments":"{}"}}]}`, + } { + t.Run(name, func(t *testing.T) { + c, _, info := newChatStreamTest(t) + _, apiErr := replayChatStream(c, info, chatRoleFrame, `{"choices":[{"delta":`+delta+`}]}`, chatUsageFrame, "[DONE]") + require.Nil(t, apiErr) + require.True(t, info.HasRecordedChannelFirstResult()) + require.False(t, info.FirstResultTime.IsZero()) + }) + } +} + +func TestFirstResultWriteFailuresDoNotCountAsOutput(t *testing.T) { + for name, handler := range map[string]func(*gin.Context, *relaycommon.RelayInfo, *http.Response) (*dto.Usage, *types.MaxAPIError){ + "chat": OaiStreamHandler, "responses_to_chat": OaiResponsesToChatStreamHandler, + "responses": OaiResponsesStreamHandler, "chat_to_responses": OaiChatToResponsesStreamHandler, + } { + t.Run(name, func(t *testing.T) { + c, _, info := newChatStreamTest(t) + common.EmptyCompletionRetryEnabled = false + failed, _ := gin.CreateTestContext(&rejectedChatResponseWriter{header: make(http.Header)}) + c.Writer = failed.Writer + frame := `{"choices":[{"delta":{"content":"never delivered"}}]}` + if name == "responses" || name == "responses_to_chat" { + frame = `{"type":"response.output_text.delta","delta":"never delivered"}` + } + _, apiErr := handler(c, info, &http.Response{StatusCode: 200, Body: io.NopCloser(strings.NewReader("data: " + frame + "\n\n"))}) + require.NotNil(t, apiErr) + require.True(t, types.IsSkipRetryError(apiErr), "committed write failures must not replay") + require.False(t, info.HasRecordedChannelFirstResult()) + require.True(t, info.FirstResultTime.IsZero()) + }) + } +} + +func TestChatFirstResultExcludesMetadataAndDroppedExtensions(t *testing.T) { + for _, test := range []struct { + name, delta string + force bool + }{ + {"provider_metadata", `{"provider_stats":{"tokens":0}}`, false}, + {"tool_index", `{"tool_calls":[{"index":0}]}`, false}, + {"audio_id", `{"audio":{"id":"audio-1"}}`, false}, + {"formatted_audio", `{"audio":{"data":"YWJj"}}`, true}, + {"formatted_refusal", `{"refusal":"discarded by the existing typed format"}`, true}, + } { + t.Run(test.name, func(t *testing.T) { + c, _, info := newChatStreamTest(t) + common.EmptyCompletionRetryEnabled = false + info.ChannelSetting.ForceFormat = test.force + _, _ = replayChatStream(c, info, `{"choices":[{"delta":`+test.delta+`}]}`, "[DONE]") + require.True(t, info.HasRecordedChannelFirstResponse()) + require.False(t, info.HasRecordedChannelFirstResult(), "only delivered generation fields may close the output deadline") + }) + } +} diff --git a/relay/channel/openai/helper.go b/relay/channel/openai/helper.go index 2b897716..4c136e99 100644 --- a/relay/channel/openai/helper.go +++ b/relay/channel/openai/helper.go @@ -202,9 +202,9 @@ func HandleFinalResponse(c *gin.Context, info *relaycommon.RelayInfo, lastStream } } -func sendResponsesStreamData(c *gin.Context, streamResponse dto.ResponsesStreamResponse, data string) { +func sendResponsesStreamData(c *gin.Context, streamResponse dto.ResponsesStreamResponse, data string) (bool, error) { if data == "" { - return + return false, nil } - _ = helper.ResponseChunkData(c, streamResponse, data) + return helper.ResponseChunkDataWithDelivery(c, streamResponse, data) } diff --git a/relay/channel/openai/relay-openai.go b/relay/channel/openai/relay-openai.go index 42fca63f..0b1e9165 100644 --- a/relay/channel/openai/relay-openai.go +++ b/relay/channel/openai/relay-openai.go @@ -159,7 +159,7 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re }) // 对音频模型，从倒数第二个stream data中提取usage信息 - if apiErr := helper.FirstResultTimeoutError(info); apiErr != nil { + if apiErr := helper.FirstResultTimeoutError(c, info); apiErr != nil { return nil, apiErr } diff --git a/relay/channel/openai/relay_responses.go b/relay/channel/openai/relay_responses.go index 17f2206b..379e7c15 100644 --- a/relay/channel/openai/relay_responses.go +++ b/relay/channel/openai/relay_responses.go @@ -12,6 +12,7 @@ import ( relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" "github.com/MAX-API-Next/MAX-API/relay/helper" "github.com/MAX-API-Next/MAX-API/service" + "github.com/MAX-API-Next/MAX-API/service/openaicompat" "github.com/MAX-API-Next/MAX-API/types" "github.com/gin-gonic/gin" @@ -237,6 +238,7 @@ func responsesWebSearchToolName(info *relaycommon.RelayInfo) string { } func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.MaxAPIError) { + info.EnableFirstResultTracking() if resp == nil || resp.Body == nil { logger.LogError(c, "invalid response or response body") return nil, types.NewError(fmt.Errorf("invalid response"), types.ErrorCodeBadResponse) @@ -255,13 +257,49 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp streamForwarded := false hasVisiblePayload := false emptyCompletionRecorded := false - - flushPendingEvents := func() { + terminalEventSeen := false + usageReported := false + outputDelivered := false + + sendEvent := func(event pendingResponsesStreamEvent) bool { + delivered, writeErr := sendResponsesStreamData(c, event.response, event.data) + if delivered { + streamForwarded = true + if helper.ResponsesStreamEventHasOutput(event.response) || responsesStreamHasVisibleOutput(&event.response) { + outputDelivered = true + } + if event.response.Type == "response.output_text.delta" { + responseTextBuilder.WriteString(event.response.Delta) + } + } + if writeErr != nil { + streamErr = types.NewError(writeErr, types.ErrorCodeBadResponse) + return false + } + if delivered && responsesStreamHasVisibleOutput(&event.response) { + info.SetFirstResultTime() + } + return true + } + finalizeUsage := func() { + if !usageReported { + if usage.CompletionTokens == 0 && responseTextBuilder.Len() > 0 { + usage.CompletionTokens = service.CountTextToken(responseTextBuilder.String(), info.UpstreamModelName) + } + if usage.PromptTokens == 0 && usage.CompletionTokens != 0 { + usage.PromptTokens = info.GetEstimatePromptTokens() + } + } + usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens + } + flushPendingEvents := func() bool { for _, event := range pendingEvents { - sendResponsesStreamData(c, event.response, event.data) + if !sendEvent(event) { + return false + } } pendingEvents = pendingEvents[:0] - streamForwarded = true + return true } shouldBufferForEmptyRetry := func() bool { return !streamForwarded && shouldRetryEmptyCompletion(c, info) @@ -277,21 +315,32 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp var streamResponse dto.ResponsesStreamResponse if err := common.UnmarshalJsonStr(data, &streamResponse); err != nil { logger.LogError(c, "failed to unmarshal stream response: "+err.Error()) - sr.Error(err) + streamErr = types.NewError(err, types.ErrorCodeBadResponseBody) + sr.Stop(streamErr) return } + if streamResponse.Response != nil && (dto.HasOpenAIUsageTokens(streamResponse.Response.Usage) || + helper.HasUsageTokenFields(data, "response.usage", "input_tokens", "output_tokens") || + helper.HasUsageTokenFields(data, "response.usage", "prompt_tokens", "completion_tokens")) { + usageReported = true + applyResponsesUsage(usage, streamResponse.Response.Usage) + if helper.HasUsageTokenFields(data, "response.usage", "input_tokens", "output_tokens") || + helper.HasUsageTokenFields(data, "response.usage", "prompt_tokens", "completion_tokens") { + usage.BillingUsage = dto.NewReportedOpenAIResponsesBillingUsage(usage) + } + } if responsesStreamHasVisibleOutput(&streamResponse) { hasVisiblePayload = true } var terminalEventErr *types.MaxAPIError switch streamResponse.Type { case "response.completed": + terminalEventSeen = true if streamResponse.Response != nil { if terminalErr := responsesTerminalError(streamResponse.Response, http.StatusInternalServerError, streamForwarded); terminalErr != nil { terminalEventErr = terminalErr break } - applyResponsesUsage(usage, streamResponse.Response.Usage) observeResponsesOutputs(info, streamResponse.Response.Output) if streamResponse.Response.HasImageGenerationCall() { c.Set("image_generation_call", true) @@ -307,9 +356,12 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp return } } + case "response.incomplete", "response.done": + terminalEventSeen = true + if streamResponse.Response != nil { + observeResponsesOutputs(info, streamResponse.Response.Output) + } case "response.output_text.delta": - // 处理输出文本 - responseTextBuilder.WriteString(streamResponse.Delta) case dto.ResponsesOutputTypeItemDone: // 函数调用处理 if streamResponse.Item != nil { @@ -324,7 +376,10 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp if terminalEventErr != nil { if streamForwarded { - sendResponsesStreamData(c, streamResponse, data) + if !sendEvent(pendingResponsesStreamEvent{response: streamResponse, data: data}) { + sr.Stop(streamErr) + return + } } streamErr = terminalEventErr sr.Stop(streamErr) @@ -332,7 +387,9 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp } if streamForwarded { - sendResponsesStreamData(c, streamResponse, data) + if !sendEvent(pendingResponsesStreamEvent{response: streamResponse, data: data}) { + sr.Stop(streamErr) + } return } @@ -342,19 +399,34 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp }) switch streamResponse.Type { case "response.completed": - flushPendingEvents() + if !flushPendingEvents() { + sr.Stop(streamErr) + } default: if hasVisiblePayload || len(pendingEvents) >= maxPendingResponsesStreamEvents || !shouldBufferForEmptyRetry() { - flushPendingEvents() + if !flushPendingEvents() { + sr.Stop(streamErr) + } } } }) - if apiErr := helper.FirstResultTimeoutError(info); apiErr != nil { + if apiErr := helper.FirstResultTimeoutError(c, info); apiErr != nil { + if outputDelivered { + finalizeUsage() + return usage, apiErr + } return nil, apiErr } if streamErr != nil { + if streamForwarded || c.Writer.Written() { + streamErr = types.NewError(streamErr, streamErr.GetErrorCode(), types.ErrOptionWithSkipRetry()) + } + if outputDelivered { + finalizeUsage() + return usage, streamErr + } return nil, streamErr } if !hasVisiblePayload && !emptyCompletionRecorded { @@ -364,30 +436,35 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp } } if !streamForwarded && len(pendingEvents) > 0 { - flushPendingEvents() - } - - if usage.CompletionTokens == 0 { - // 计算输出文本的 token 数量 - tempStr := responseTextBuilder.String() - if len(tempStr) > 0 { - // 非正常结束，使用输出文本的 token 数量 - completionTokens := service.CountTextToken(tempStr, info.UpstreamModelName) - usage.CompletionTokens = completionTokens + if !flushPendingEvents() { + if outputDelivered { + finalizeUsage() + return usage, types.NewError(streamErr, streamErr.GetErrorCode(), types.ErrOptionWithSkipRetry()) + } + return nil, types.NewError(streamErr, streamErr.GetErrorCode(), types.ErrOptionWithSkipRetry()) } } - if usage.PromptTokens == 0 && usage.CompletionTokens != 0 { - usage.PromptTokens = info.GetEstimatePromptTokens() - } - - usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens + finalizeUsage() if service.ResponseAuditEnabled() { service.SetRelayResponseAuditContent(info, responseTextBuilder.String()) } if hasVisiblePayload { recordEmptyCompletionRetrySuccess(c, info, usage) } + if !terminalEventSeen { + reason := "upstream_eof" + if info.StreamStatus != nil && info.StreamStatus.IsAbnormalEnd() { + reason = string(info.StreamStatus.EndReason) + } + // Lifecycle/usage metadata is not delivered generation. Returning it + // as partial usage would start a finalize before the failure refund, + // or make an empty zero-usage stream look successfully completed. + if !outputDelivered { + return nil, openaicompat.NewResponsesStreamIncompleteError(reason) + } + return usage, openaicompat.NewResponsesStreamIncompleteError(reason) + } return usage, nil } diff --git a/relay/channel/openai/responses_terminal_repair_test.go b/relay/channel/openai/responses_terminal_repair_test.go new file mode 100644 index 00000000..168b96bb --- /dev/null +++ b/relay/channel/openai/responses_terminal_repair_test.go @@ -0,0 +1,195 @@ +package openai + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/dto" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/service" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestResponsesToChatLegalTerminalEvents(t *testing.T) { + for _, terminal := range []struct { + name, event, status, reason, finish string + }{ + {"length", "response.incomplete", "incomplete", "max_output_tokens", "length"}, + {"filter", "response.incomplete", "incomplete", "content_filter", "content_filter"}, + {"omitted_status", "response.incomplete", "", "max_output_tokens", "length"}, + {"done", "response.done", "completed", "", "stop"}, + {"done_incomplete", "response.done", "incomplete", "max_output_tokens", "length"}, + {"completed_incomplete", "response.completed", "incomplete", "content_filter", "content_filter"}, + } { + for _, zero := range []bool{false, true} { + for _, terminalOnly := range []bool{false, true} { + t.Run(terminal.name+map[bool]string{false: "/known", true: "/zero"}[zero]+map[bool]string{false: "/deltas", true: "/terminal_only"}[terminalOnly], func(t *testing.T) { + recorder := httptest.NewRecorder() + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + info := &relaycommon.RelayInfo{RelayFormat: types.RelayFormatOpenAI, OriginModelName: "gpt-test", DisablePing: true, ShouldIncludeUsage: true, + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-test"}} + info.SetEstimatePromptTokens(10) + prompt, completion := 10, 2 + if zero { + prompt, completion = 0, 0 + } + response := map[string]any{"id": "resp_terminal", "model": "gpt-test", + "output": []any{map[string]any{"type": "message", "role": "assistant", "content": []any{map[string]any{"type": "output_text", "text": "partial output"}}}}, + "usage": map[string]any{"input_tokens": prompt, "output_tokens": completion, "total_tokens": prompt + completion}} + if terminal.status != "" { + response["status"] = terminal.status + } + if terminal.reason != "" { + response["incomplete_details"] = map[string]any{"reason": terminal.reason} + } + data, err := common.Marshal(map[string]any{"type": terminal.event, "response": response}) + require.NoError(t, err) + frames := []string{`{"type":"response.created","response":{"id":"resp_terminal","model":"gpt-test"}}`} + if !terminalOnly { + frames = append(frames, `{"type":"response.output_text.delta","delta":"partial output"}`) + } + frames = append(frames, string(data)) + usage, apiErr := OaiResponsesToChatStreamHandler(c, info, &http.Response{StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader("data: " + strings.Join(frames, "\n\ndata: ") + "\n\n"))}) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Equal(t, prompt, usage.PromptTokens) + require.Equal(t, completion, usage.CompletionTokens) + require.Equal(t, prompt+completion, usage.TotalTokens) + require.Contains(t, recorder.Body.String(), `"finish_reason":"`+terminal.finish+`"`) + require.Equal(t, 1, strings.Count(recorder.Body.String(), "partial output")) + require.Equal(t, 1, strings.Count(recorder.Body.String(), "data: [DONE]")) + require.NotContains(t, recorder.Body.String(), `"error":`) + }) + } + } + } +} + +func TestResponsesToChatTerminalOnlyToolCalls(t *testing.T) { + for _, event := range []string{"response.completed", "response.done", "response.incomplete"} { + for _, toolType := range []string{dto.BuildInCallFunctionCall, dto.BuildInCallCustomToolCall} { + t.Run(event+"/"+toolType, func(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + info := &relaycommon.RelayInfo{RelayFormat: types.RelayFormatOpenAI, OriginModelName: "gpt-test", DisablePing: true, + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-test"}} + item := map[string]any{"type": toolType, "id": "item_terminal", "call_id": "call_terminal", "name": "lookup"} + if toolType == dto.BuildInCallCustomToolCall { + item["input"] = "terminal payload" + } else { + item["arguments"] = "terminal payload" + } + status, finish := "completed", "tool_calls" + if event == "response.incomplete" { + status, finish = "incomplete", "length" + } + data, err := common.Marshal(map[string]any{"type": event, "response": map[string]any{"id": "resp_tool_terminal", "status": status, "output": []any{item}, + "usage": map[string]any{"input_tokens": 10, "output_tokens": 2, "total_tokens": 12}}}) + require.NoError(t, err) + usage, apiErr := OaiResponsesToChatStreamHandler(c, info, &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader("data: " + string(data) + "\n\n"))}) + require.Nil(t, apiErr) + require.Equal(t, 12, usage.TotalTokens) + require.Contains(t, recorder.Body.String(), `"arguments":"terminal payload"`) + require.Contains(t, recorder.Body.String(), `"name":"lookup"`) + require.Contains(t, recorder.Body.String(), `"id":"call_terminal"`) + require.Contains(t, recorder.Body.String(), `"finish_reason":"`+finish+`"`) + require.Equal(t, 1, strings.Count(recorder.Body.String(), "data: [DONE]")) + }) + } + } +} + +type responsesToChatToolFailureWriter struct { + *httptest.ResponseRecorder +} + +func (w *responsesToChatToolFailureWriter) Write(p []byte) (int, error) { + if strings.Contains(string(p), "undelivered") { + return 0, io.ErrClosedPipe + } + return w.ResponseRecorder.Write(p) +} + +func TestResponsesToChatPartialEstimateIncludesDeliveredTools(t *testing.T) { + for _, toolType := range []string{dto.BuildInCallFunctionCall, dto.BuildInCallCustomToolCall} { + for _, mixed := range []bool{false, true} { + for _, ending := range []string{"eof", "malformed", "write_failure"} { + t.Run(toolType+map[bool]string{false: "/tool_only/", true: "/mixed/"}[mixed]+ending, func(t *testing.T) { + recorder := httptest.NewRecorder() + var writer http.ResponseWriter = recorder + if ending == "write_failure" { + writer = &responsesToChatToolFailureWriter{recorder} + } + c, _ := gin.CreateTestContext(writer) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + info := &relaycommon.RelayInfo{RelayFormat: types.RelayFormatOpenAI, OriginModelName: "gpt-test", DisablePing: true, + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-test"}} + info.SetEstimatePromptTokens(10) + frames, deliveredText := []string{}, "" + if mixed { + frames = append(frames, `{"type":"response.output_text.delta","delta":"visible text"}`) + frames = append(frames, `{"type":"response.reasoning_summary_text.delta","delta":"reasoning"}`) + deliveredText = "visible textreasoning" + } + item := map[string]any{"type": toolType, "id": "item_partial", "call_id": "call_partial", "name": "lookup"} + itemData, err := common.Marshal(map[string]any{"type": "response.output_item.added", "item": item}) + require.NoError(t, err) + frames = append(frames, string(itemData)) + argumentEvent := "response.function_call_arguments.delta" + if toolType == dto.BuildInCallCustomToolCall { + argumentEvent = "response.custom_tool_call_input.delta" + } + argumentData, err := common.Marshal(map[string]any{"type": argumentEvent, "item_id": "item_partial", "delta": "delivered tool arguments"}) + require.NoError(t, err) + frames = append(frames, string(argumentData)) + deliveredText += "lookupdelivered tool arguments" + if ending == "malformed" { + frames = append(frames, `{broken`) + } + if ending == "write_failure" { + failed, err := common.Marshal(map[string]any{"type": argumentEvent, "item_id": "item_partial", "delta": "undelivered"}) + require.NoError(t, err) + frames = append(frames, string(failed)) + } + usage, apiErr := OaiResponsesToChatStreamHandler(c, info, &http.Response{StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader("data: " + strings.Join(frames, "\n\ndata: ") + "\n\n"))}) + require.NotNil(t, apiErr) + require.True(t, types.IsSkipRetryError(apiErr)) + require.NotNil(t, usage) + require.Equal(t, 10, usage.PromptTokens) + require.Equal(t, service.EstimateTokenByModel("gpt-test", deliveredText), usage.CompletionTokens) + require.Contains(t, recorder.Body.String(), "delivered tool arguments") + require.NotContains(t, recorder.Body.String(), "undelivered") + }) + } + } + } +} + +func TestResponsesToChatLegalTerminalsPreserveProviderFailure(t *testing.T) { + for _, event := range []string{"response.completed", "response.done", "response.incomplete"} { + t.Run(event, func(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + info := &relaycommon.RelayInfo{RelayFormat: types.RelayFormatOpenAI, OriginModelName: "gpt-test", DisablePing: true, + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-test"}} + data := `{"type":"response.output_text.delta","delta":"partial"}` + "\n\ndata: " + `{"type":"` + event + `","response":{"status":"failed","error":{"message":"provider failure"}}}` + _, apiErr := OaiResponsesToChatStreamHandler(c, info, &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader("data: " + data + "\n\n"))}) + require.NotNil(t, apiErr) + require.Equal(t, "provider failure", apiErr.Error()) + require.True(t, types.IsSkipRetryError(apiErr)) + require.NotContains(t, recorder.Body.String(), `"finish_reason":"stop"`) + }) + } +} diff --git a/relay/channel/openai/responses_usage_test.go b/relay/channel/openai/responses_usage_test.go index 4d55b77e..04d3492d 100644 --- a/relay/channel/openai/responses_usage_test.go +++ b/relay/channel/openai/responses_usage_test.go @@ -520,6 +520,29 @@ func TestOaiResponsesStreamHandlerAllowsRetryWhenFirstEventIsError(t *testing.T) } } +func TestOaiResponsesStreamHandlerEOFAfterOutputReturnsIncompleteError(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + info := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatOpenAIResponses, + OriginModelName: "gpt-4o", + DisablePing: true, + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-4o"}, + } + + usage, maxAPIError := OaiResponsesStreamHandler(c, info, &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(`data: {"type":"response.output_text.delta","delta":"partial"}` + "\n")), + }) + require.NotNil(t, usage) + require.NotNil(t, maxAPIError) + require.True(t, types.IsSkipRetryError(maxAPIError)) + require.Contains(t, recorder.Body.String(), "response.output_text.delta") +} + func TestOaiResponsesToChatStreamHandlerPreservesMessageOnlyTerminalError(t *testing.T) { setOpenAIToolPricesForTest(t, map[string]float64{"lookup": 5}) c, info := newOpenAIToolBillingContext("gpt-test") diff --git a/relay/channel/openai/responses_via_chat.go b/relay/channel/openai/responses_via_chat.go index 7d4c0655..0b1089bc 100644 --- a/relay/channel/openai/responses_via_chat.go +++ b/relay/channel/openai/responses_via_chat.go @@ -38,11 +38,11 @@ func OaiChatToResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp observeOpenAITextToolCalls(info, chatResp.Choices) responseID := helper.GetResponseID(c) - responsesResp, usage, err := service.ChatCompletionsResponseToResponsesResponse(&chatResp, responseID) + responsesResp, usage, err := service.ChatCompletionsResponseToResponsesResponseWithCustomTools(&chatResp, responseID, service.ResponsesCustomToolNames(info)) if err != nil { return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) } - if usage == nil || usage.TotalTokens == 0 { + if (usage == nil || usage.TotalTokens == 0) && !helper.HasUsageTokenFields(string(body), "usage", "prompt_tokens", "completion_tokens") { text := service.ExtractOutputTextFromResponses(responsesResp) usage = service.ResponseText2Usage(c, text, info.UpstreamModelName, info.GetEstimatePromptTokens()) responsesResp.Usage = openaicompat.UsageFromChatUsage(usage) @@ -60,6 +60,7 @@ func OaiChatToResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp } func OaiChatToResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.MaxAPIError) { + info.EnableFirstResultTracking() if resp == nil || resp.Body == nil { return nil, types.NewOpenAIError(fmt.Errorf("invalid response"), types.ErrorCodeBadResponse, http.StatusInternalServerError) } @@ -67,8 +68,14 @@ func OaiChatToResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo responseID := helper.GetResponseID(c) state := openaicompat.NewChatToResponsesStreamState(responseID, info.UpstreamModelName) + state.CustomToolNames = service.ResponsesCustomToolNames(info) toolCallObserver := newOpenAIStreamToolCallObserver(info) var streamErr *types.MaxAPIError + terminalResponseSeen := false + usageReported := false + outputDelivered := false + var deliveredText strings.Builder + var deliveredUsage *dto.Usage sendEvent := func(event openaicompat.ChatToResponsesStreamEvent) bool { data, err := common.Marshal(event.Payload) @@ -76,10 +83,18 @@ func OaiChatToResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo streamErr = types.NewOpenAIError(err, types.ErrorCodeJsonMarshalFailed, http.StatusInternalServerError) return false } - if err := helper.ResponseChunkData(c, dto.ResponsesStreamResponse{Type: event.Type}, string(data)); err != nil { - streamErr = types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError) + delivered, writeErr := helper.ResponseChunkDataWithDelivery(c, dto.ResponsesStreamResponse{Type: event.Type}, string(data)) + if delivered && helper.ResponsesStreamEventHasOutput(event.Payload) { + outputDelivered = true + deliveredText.WriteString(event.Payload.Delta) + } + if writeErr != nil { + streamErr = types.NewOpenAIError(writeErr, types.ErrorCodeBadResponse, http.StatusInternalServerError) return false } + if delivered && helper.ResponsesStreamEventHasFirstResult(event.Payload) { + info.SetFirstResultTime() + } return true } @@ -103,10 +118,22 @@ func OaiChatToResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo var chunk dto.ChatCompletionsStreamResponse if err := common.UnmarshalJsonStr(data, &chunk); err != nil { logger.LogError(c, "failed to unmarshal chat stream response: "+err.Error()) - sr.Error(err) + streamErr = types.NewError(err, types.ErrorCodeBadResponseBody) + sr.Stop(streamErr) return } + reported := helper.HasUsageTokenFields(data, "usage", "prompt_tokens", "completion_tokens") || dto.HasOpenAIUsageTokens(chunk.Usage) + if !reported { + chunk.Usage = nil + } else if helper.HasUsageTokenFields(data, "usage", "prompt_tokens", "completion_tokens") { + chunk.Usage.BillingUsage = dto.NewReportedOpenAIChatBillingUsage(chunk.Usage) + } toolCallObserver.Observe(&chunk) + for _, choice := range chunk.Choices { + if choice.FinishReason != nil && strings.TrimSpace(*choice.FinishReason) != "" { + terminalResponseSeen = true + } + } events, err := openaicompat.ChatCompletionsStreamChunkToResponsesEvents(&chunk, state) if err != nil { @@ -120,26 +147,77 @@ func OaiChatToResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo return } } + deliveredUsage = openaicompat.UsageFromChatUsage(state.Usage) + usageReported = usageReported || reported }) - if apiErr := helper.FirstResultTimeoutError(info); apiErr != nil { + usage := deliveredUsage + if usage == nil { + usage = &dto.Usage{} + } + usageTextSize := -1 + refreshUsage := func() { + if !usageReported && usageTextSize != deliveredText.Len() { + usage = service.ResponseText2Usage(c, deliveredText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens()) + state.Usage = openaicompat.UsageFromChatUsage(usage) + usageTextSize = deliveredText.Len() + } + } + refreshUsage() + finalizeStream := func() bool { + for _, event := range openaicompat.FlushChatCompletionsStreamToResponsesOutput(state) { + if !sendEvent(event) { + refreshUsage() + return false + } + } + refreshUsage() + for _, event := range openaicompat.FinalizeChatCompletionsStreamToResponses(state) { + if !sendEvent(event) { + return false + } + } + return true + } + if apiErr := helper.FirstResultTimeoutError(c, info); apiErr != nil { + // A whitespace delta is delivered output even while the meaningful + // first-result deadline is still running. Preserve partial settlement. + if outputDelivered { + return usage, apiErr + } return nil, apiErr } - if streamErr != nil { + if c.Writer.Written() { + streamErr = types.NewError(streamErr, streamErr.GetErrorCode(), types.ErrOptionWithSkipRetry()) + } + if outputDelivered { + return usage, streamErr + } return nil, streamErr } - - usage := state.Usage - if usage == nil || usage.TotalTokens == 0 { - usage = service.ResponseText2Usage(c, state.UsageText(), info.UpstreamModelName, info.GetEstimatePromptTokens()) - state.Usage = openaicompat.UsageFromChatUsage(usage) + if !terminalResponseSeen { + reason := "upstream_eof" + if info.StreamStatus != nil && info.StreamStatus.IsAbnormalEnd() { + reason = string(info.StreamStatus.EndReason) + } + state.MarkIncomplete(reason) + if info.StreamStatus == nil || info.StreamStatus.EndReason == relaycommon.StreamEndReasonEOF { + if !finalizeStream() { + return usage, types.NewError(streamErr, streamErr.GetErrorCode(), types.ErrOptionWithSkipRetry()) + } + } + if !outputDelivered { + return nil, openaicompat.NewResponsesStreamIncompleteError(reason) + } + return usage, openaicompat.NewResponsesStreamIncompleteError(reason) + } + if info.StreamStatus != nil && info.StreamStatus.IsAbnormalEnd() { + return usage, openaicompat.NewResponsesStreamIncompleteError(string(info.StreamStatus.EndReason)) } - for _, event := range openaicompat.FinalizeChatCompletionsStreamToResponses(state) { - if !sendEvent(event) { - return nil, streamErr - } + if !finalizeStream() { + return usage, types.NewError(streamErr, streamErr.GetErrorCode(), types.ErrOptionWithSkipRetry()) } if service.ResponseAuditEnabled() { service.SetRelayResponseAuditContent(info, state.UsageText()) diff --git a/relay/channel/xai/text.go b/relay/channel/xai/text.go index bb7467f9..6ea52db7 100644 --- a/relay/channel/xai/text.go +++ b/relay/channel/xai/text.go @@ -66,7 +66,7 @@ func xAIStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re sr.Error(err) } }) - if apiErr := helper.FirstResultTimeoutError(info); apiErr != nil { + if apiErr := helper.FirstResultTimeoutError(c, info); apiErr != nil { return nil, apiErr } diff --git a/relay/chat_completions_via_responses.go b/relay/chat_completions_via_responses.go index 41a22427..ec8d79df 100644 --- a/relay/chat_completions_via_responses.go +++ b/relay/chat_completions_via_responses.go @@ -160,7 +160,7 @@ func chatCompletionsViaResponses(c *gin.Context, info *relaycommon.RelayInfo, ad usage, maxApiErr := openaichannel.OaiResponsesToChatStreamHandler(c, info, httpResp) if maxApiErr != nil { service.ResetStatusCode(maxApiErr, statusCodeMappingStr) - return nil, maxApiErr + return usage, maxApiErr } return usage, nil } diff --git a/relay/chat_stream_billing_test.go b/relay/chat_stream_billing_test.go index c47c51a4..2bb1d770 100644 --- a/relay/chat_stream_billing_test.go +++ b/relay/chat_stream_billing_test.go @@ -262,3 +262,59 @@ func TestNativeChatPartialBillingDoesNotRecordSuccess(t *testing.T) { time.Sleep(10 * time.Millisecond) } } + +func TestResponsesPartialBillingRecorderPreservesFailureMetricAndRefundBoundary(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.PerfMetric{}, &model.PerfMetricFlushReceipt{})) + previousDB, previousRedis := model.DB, common.RedisEnabled + setting := config.GlobalConfig.Get("perf_metrics_setting").(*perf_metrics_setting.PerfMetricsSetting) + previousSetting := *setting + model.DB, common.RedisEnabled, setting.Enabled = db, false, true + t.Cleanup(func() { + model.DB, common.RedisEnabled = previousDB, previousRedis + *setting = previousSetting + sqlDB, _ := db.DB() + _ = sqlDB.Close() + }) + + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + billing := &chatStreamBillingRecorder{} + sequence := chatStreamMetricSequence.Add(1) + info := &relaycommon.RelayInfo{ + RequestId: fmt.Sprintf("responses-partial-canary-%d", sequence), + OriginModelName: fmt.Sprintf("responses-partial-canary-model-%d", sequence), + RelayMode: relayconstant.RelayModeResponses, + RelayFormat: types.RelayFormatOpenAIResponses, + StartTime: time.Now(), + IsStream: true, + Billing: billing, + UserQuota: 1000, + UserSetting: dto.UserSetting{QuotaWarningThreshold: 1}, + ChannelMeta: &relaycommon.ChannelMeta{}, + PriceData: types.PriceData{ModelRatio: 1, CompletionRatio: 1, GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 1}}, + } + + service.PostPartialConsumeQuota(c, info, &dto.Usage{PromptTokens: 10, CompletionTokens: 2, TotalTokens: 12}) + require.Equal(t, []int{12}, billing.quotas) + require.Len(t, billing.effects, 1) + + apiErr := types.NewOpenAIError(fmt.Errorf("synthetic responses incomplete"), types.ErrorCodeBadResponse, http.StatusBadGateway, types.ErrOptionWithSkipRetry()) + service.HandleFailedBilling(c, info, apiErr) + require.Zero(t, billing.refunds, "partial settlement must prevent a full refund") + + perfmetrics.RecordRelaySample(info, false, 0) + result, queryErr := perfmetrics.QuerySummaryAll(1, nil) + require.NoError(t, queryErr) + found := false + for _, summary := range result.Models { + if summary.ModelName == info.OriginModelName { + found = true + require.EqualValues(t, 1, summary.RequestCount) + require.Zero(t, summary.SuccessRate) + break + } + } + require.True(t, found, "missing partial Responses failure metric") +} diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go index e1d3e43e..60b208cb 100644 --- a/relay/common/relay_info.go +++ b/relay/common/relay_info.go @@ -101,6 +101,7 @@ type RelayInfo struct { TokenUnlimited bool StartTime time.Time FirstResponseTime time.Time + FirstResultTime time.Time isFirstResponse bool channelAttemptStartTime time.Time channelFirstResponseMu sync.Mutex @@ -108,6 +109,9 @@ type RelayInfo struct { channelFirstResponseEvaluating bool channelFirstResponseRecorded atomic.Bool channelFirstResponseSignal chan struct{} + channelFirstResultTracking bool + channelFirstResultRecorded atomic.Bool + channelFirstResultSignal chan struct{} //SendLastReasoningResponse bool IsStream bool IsGeminiBatchEmbedding bool @@ -271,6 +275,11 @@ func (info *RelayInfo) InitChannelMeta(c *gin.Context) { info.channelAttemptStartTime = time.Now() info.channelFirstResponseRecorded.Store(false) info.channelFirstResponseSignal = make(chan struct{}) + info.channelFirstResultTracking = false + info.channelFirstResultRecorded.Store(false) + info.channelFirstResultSignal = nil + info.channelFirstResponseEvaluating = false + info.channelFirstResponseEvalAt = time.Time{} if model_setting.GetGlobalSettings().PassThroughRequestEnabled || channelMeta.ChannelSetting.PassThroughBodyEnabled { info.ReasoningEffort = "" } else { @@ -785,7 +794,9 @@ func (info *RelayInfo) setFirstResponseTimeLocked() { return } info.recordChannelHealthObservation("first_response") - info.recordChannelHealthSuccess() + if !info.channelFirstResultTracking { + info.recordChannelHealthSuccess() + } } const firstResponseEvaluationGrace = 100 * time.Millisecond @@ -799,7 +810,7 @@ func (info *RelayInfo) BeginFirstResponseEvaluation() func() { return func() {} } info.channelFirstResponseMu.Lock() - if !info.channelFirstResponseRecorded.Load() { + if !info.HasRecordedChannelFirstResult() { info.channelFirstResponseEvaluating = true info.channelFirstResponseEvalAt = time.Now() } @@ -828,9 +839,9 @@ func (info *RelayInfo) FirstResponseDeadlineExpired() bool { return true } -// FirstResponseSignal is closed once the current channel attempt has produced -// a first valid, deliverable result. It is used by the stream watchdog without -// exposing the internal deduplication state to relay callers. +// FirstResponseSignal is closed after the first accepted response frame for +// this attempt. Metadata may close it; strict adapters use FirstResultSignal +// for their output deadline instead. func (info *RelayInfo) FirstResponseSignal() <-chan struct{} { if info == nil { return nil @@ -991,7 +1002,7 @@ func (info *RelayInfo) RecordChannelTimeout(event string) { if start.IsZero() || threshold <= 0 || observedAt.Sub(start) < time.Duration(threshold)*time.Second { return } - if info.IsStream && !info.channelFirstResponseRecorded.Load() { + if info.IsStream && !info.HasRecordedChannelFirstResult() { timeoutKind = channelhealth.TimeoutKindStreamingFirstResult } else if !info.IsStream { timeoutKind = channelhealth.TimeoutKindNonStreamingResponse diff --git a/relay/common/stream_first_result.go b/relay/common/stream_first_result.go new file mode 100644 index 00000000..58b50ec6 --- /dev/null +++ b/relay/common/stream_first_result.go @@ -0,0 +1,66 @@ +package common + +import "time" + +// EnableFirstResultTracking is called by adapters before starting workers. +// Unclassified adapters retain their existing first-response semantics. +func (info *RelayInfo) EnableFirstResultTracking() { + if info == nil || info.channelFirstResultTracking { + return + } + info.channelFirstResultTracking = true + info.channelFirstResultSignal = make(chan struct{}) +} + +// SetFirstResultTime records a successfully delivered generation payload, +// independently of lifecycle/usage frames and their existing latency metric. +func (info *RelayInfo) SetFirstResultTime() { + if info == nil { + return + } + info.channelFirstResponseMu.Lock() + defer info.channelFirstResponseMu.Unlock() + info.setFirstResponseTimeLocked() + if !info.channelFirstResultTracking || !info.channelFirstResultRecorded.CompareAndSwap(false, true) { + return + } + if info.FirstResultTime.IsZero() { + info.FirstResultTime = time.Now() + } + close(info.channelFirstResultSignal) + if !info.IsChannelTest { + info.recordChannelHealthSuccess() + } +} + +func (info *RelayInfo) HasRecordedChannelFirstResult() bool { + if info == nil { + return false + } + if !info.channelFirstResultTracking { + return info.HasRecordedChannelFirstResponse() + } + return info.channelFirstResultRecorded.Load() +} + +func (info *RelayInfo) FirstResultSignal() <-chan struct{} { + if info == nil { + return nil + } + if !info.channelFirstResultTracking { + return info.FirstResponseSignal() + } + return info.channelFirstResultSignal +} + +func (info *RelayInfo) FirstResultDeadlineExpired() bool { + if info == nil { + return true + } + info.channelFirstResponseMu.Lock() + defer info.channelFirstResponseMu.Unlock() + if info.HasRecordedChannelFirstResult() { + return false + } + return !info.channelFirstResponseEvaluating || time.Since(info.channelFirstResponseEvalAt) >= firstResponseEvaluationGrace +} diff --git a/relay/common/stream_first_result_test.go b/relay/common/stream_first_result_test.go new file mode 100644 index 00000000..14889598 --- /dev/null +++ b/relay/common/stream_first_result_test.go @@ -0,0 +1,153 @@ +package common + +import ( + "context" + "net/http/httptest" + "testing" + "time" + + maxcommon "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/pkg/channelhealth" + "github.com/MAX-API-Next/MAX-API/setting/operation_setting" + "github.com/alicebob/miniredis/v2" + "github.com/gin-gonic/gin" + "github.com/go-redis/redis/v8" + "github.com/stretchr/testify/require" +) + +func TestFirstResultKeepsFirstFrameMetricAndAttemptSignalsSeparate(t *testing.T) { + start := time.Now().Add(-time.Second) + info := &RelayInfo{StartTime: start, isFirstResponse: true, channelFirstResponseSignal: make(chan struct{})} + info.EnableFirstResultTracking() + info.SetFirstResponseTime() + frameTime := info.FirstResponseTime + require.True(t, info.HasRecordedChannelFirstResponse()) + require.False(t, info.HasRecordedChannelFirstResult()) + require.True(t, info.FirstResultDeadlineExpired()) + select { + case <-info.FirstResultSignal(): + t.Fatal("metadata closed the output signal") + default: + } + finish := info.BeginFirstResponseEvaluation() + require.False(t, info.FirstResultDeadlineExpired()) + finish() + require.True(t, info.FirstResultDeadlineExpired()) + info.SetFirstResultTime() + resultTime := info.FirstResultTime + info.SetFirstResultTime() + require.Equal(t, frameTime, info.FirstResponseTime) + require.Equal(t, resultTime, info.FirstResultTime) + require.False(t, resultTime.Before(frameTime)) + require.False(t, info.FirstResultDeadlineExpired()) + select { + case <-info.FirstResultSignal(): + default: + t.Fatal("delivered output did not close the result signal") + } + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + info.InitChannelMeta(c) + info.EnableFirstResultTracking() + require.False(t, info.HasRecordedChannelFirstResponse()) + require.False(t, info.HasRecordedChannelFirstResult()) + require.True(t, info.FirstResultDeadlineExpired(), "retry gets a fresh attempt budget") + require.Equal(t, frameTime, info.FirstResponseTime, "request latency stays compatible") + select { + case <-info.FirstResultSignal(): + t.Fatal("retry reused a closed signal") + default: + } +} + +func TestStrictFirstResultHealthSamplesExcludeMetadata(t *testing.T) { + server := miniredis.RunT(t) + oldRDB, oldEnabled := maxcommon.RDB, maxcommon.RedisEnabled + maxcommon.RDB = redis.NewClient(&redis.Options{Addr: server.Addr()}) + maxcommon.RedisEnabled = true + t.Cleanup(func() { _ = maxcommon.RDB.Close(); maxcommon.RDB, maxcommon.RedisEnabled = oldRDB, oldEnabled }) + oldSetting := *operation_setting.GetMonitorSetting() + *operation_setting.GetMonitorSetting() = operation_setting.MonitorSetting{ + TimeoutAutoDisableEnabled: true, TimeoutAutoDisableMinimumSamples: 2, + StreamingFirstResultTimeoutSeconds: 1, + } + t.Cleanup(func() { *operation_setting.GetMonitorSetting() = oldSetting }) + info := &RelayInfo{IsStream: true, RequestId: "first-result-health", channelAttemptStartTime: time.Now(), + ChannelMeta: &ChannelMeta{ChannelId: 936, ChannelAutoBan: true}} + info.EnableFirstResultTracking() + info.SetFirstResponseTime() + assertFirstResultHealthCondition(t, false, func() bool { + state, err := channelhealth.LoadRuntimeState(context.Background(), 936) + require.NoError(t, err) + return state.SampleCount[channelhealth.RequestModeStreaming] != 0 + }, 150*time.Millisecond) + info.SetFirstResultTime() + info.SetFirstResultTime() + assertFirstResultHealthCondition(t, true, func() bool { + state, err := channelhealth.LoadRuntimeState(context.Background(), 936) + require.NoError(t, err) + return state.SampleCount[channelhealth.RequestModeStreaming] == 1 + }, time.Second) +} + +func TestStrictFirstResultTimeoutClassification(t *testing.T) { + server := miniredis.RunT(t) + oldRDB, oldEnabled := maxcommon.RDB, maxcommon.RedisEnabled + maxcommon.RDB = redis.NewClient(&redis.Options{Addr: server.Addr()}) + maxcommon.RedisEnabled = true + t.Cleanup(func() { _ = maxcommon.RDB.Close(); maxcommon.RDB, maxcommon.RedisEnabled = oldRDB, oldEnabled }) + oldSetting := *operation_setting.GetMonitorSetting() + *operation_setting.GetMonitorSetting() = operation_setting.MonitorSetting{ + AutoPriorityDemotionEnabled: true, StreamingFirstResultTimeoutSeconds: 1, + PriorityDeduction: 10, PenaltyCooldownSeconds: 60, + } + t.Cleanup(func() { *operation_setting.GetMonitorSetting() = oldSetting }) + for idx, test := range []struct { + event string + output, penalize bool + }{ + {"transport_timeout", false, true}, + {"transport_timeout", true, false}, + {"stream_idle_timeout", false, false}, + {"stream_idle_timeout", true, false}, + } { + t.Run(test.event+map[bool]string{false: "/metadata", true: "/output"}[test.output], func(t *testing.T) { + channelID := 1936 + idx + info := &RelayInfo{IsStream: true, RequestId: "classified-first-result", channelAttemptStartTime: time.Now().Add(-2 * time.Second), + ChannelMeta: &ChannelMeta{ChannelId: channelID, ChannelAutoBan: true}} + info.EnableFirstResultTracking() + info.SetFirstResponseTime() + if test.output { + info.SetFirstResultTime() + } + info.RecordChannelTimeout(test.event) + penaltyApplied := func() bool { + state, err := channelhealth.LoadRuntimeState(context.Background(), channelID) + require.NoError(t, err) + return state.Penalty == 10 + } + if test.penalize { + assertFirstResultHealthCondition(t, true, penaltyApplied, time.Second) + } else { + assertFirstResultHealthCondition(t, false, penaltyApplied, 150*time.Millisecond) + } + }) + } +} + +// Run polling on the test goroutine: testify's asynchronous Never/Eventually +// callbacks can outlive the deadline and access t or Redis after cleanup. +func assertFirstResultHealthCondition(t *testing.T, expected bool, condition func() bool, waitFor time.Duration) { + t.Helper() + deadline := time.Now().Add(waitFor) + for { + if condition() { + require.True(t, expected, "unexpected channel health state within the observation window") + return + } + if !time.Now().Before(deadline) { + break + } + time.Sleep(10 * time.Millisecond) + } + require.False(t, expected, "expected channel health state was not observed") +} diff --git a/relay/common/stream_status.go b/relay/common/stream_status.go index 4734b264..b265a407 100644 --- a/relay/common/stream_status.go +++ b/relay/common/stream_status.go @@ -36,6 +36,9 @@ type StreamStatus struct { mu sync.Mutex Errors []StreamErrorEntry ErrorCount int + + fatalMu sync.Mutex + fatalError error } func NewStreamStatus() *StreamStatus { @@ -67,6 +70,30 @@ func (s *StreamStatus) RecordError(msg string) { } } +// RecordFatalError preserves a handler error even when the scanner has +// already recorded EOF. The scanner and handler run concurrently, so the +// first transport end reason is not sufficient to classify a fatal callback +// failure. +func (s *StreamStatus) RecordFatalError(err error) { + if s == nil || err == nil { + return + } + s.fatalMu.Lock() + defer s.fatalMu.Unlock() + if s.fatalError == nil { + s.fatalError = err + } +} + +func (s *StreamStatus) FatalError() error { + if s == nil { + return nil + } + s.fatalMu.Lock() + defer s.fatalMu.Unlock() + return s.fatalError +} + func (s *StreamStatus) HasErrors() bool { if s == nil { return false @@ -89,9 +116,32 @@ func (s *StreamStatus) IsNormalEnd() bool { if s == nil { return true } - return s.EndReason == StreamEndReasonDone || + return !s.IsAbnormalEnd() && (s.EndReason == StreamEndReasonDone || s.EndReason == StreamEndReasonEOF || - s.EndReason == StreamEndReasonHandlerStop + s.EndReason == StreamEndReasonHandlerStop) +} + +// IsAbnormalEnd reports transport or handler endings that must not be treated +// as a successful provider completion. Plain EOF is intentionally excluded: +// callers still need to verify a protocol terminal event when a provider +// closes the stream without sending [DONE] or an equivalent event. A fatal +// handler error remains abnormal even if the scanner won the EOF race. +func (s *StreamStatus) IsAbnormalEnd() bool { + if s == nil { + return false + } + if s.FatalError() != nil { + return true + } + switch s.EndReason { + case StreamEndReasonTimeout, StreamEndReasonClientGone, + StreamEndReasonScannerErr, StreamEndReasonPanic, StreamEndReasonPingFail: + return true + case StreamEndReasonHandlerStop: + return s.EndError != nil + default: + return false + } } func (s *StreamStatus) Summary() string { @@ -100,8 +150,12 @@ func (s *StreamStatus) Summary() string { } b := &strings.Builder{} fmt.Fprintf(b, "reason=%s", s.EndReason) - if s.EndError != nil { - fmt.Fprintf(b, " end_error=%q", s.EndError.Error()) + endError := s.EndError + if endError == nil { + endError = s.FatalError() + } + if endError != nil { + fmt.Fprintf(b, " end_error=%q", endError.Error()) } s.mu.Lock() if s.ErrorCount > 0 { diff --git a/relay/common/stream_status_test.go b/relay/common/stream_status_test.go index 4a31cb79..5b8a2713 100644 --- a/relay/common/stream_status_test.go +++ b/relay/common/stream_status_test.go @@ -157,6 +157,38 @@ func TestStreamStatus_IsNormalEnd_NilSafe(t *testing.T) { assert.True(t, s.IsNormalEnd()) } +func TestStreamStatus_HandlerStopWithErrorIsAbnormal(t *testing.T) { + t.Parallel() + s := NewStreamStatus() + err := fmt.Errorf("invalid provider frame") + s.RecordFatalError(err) + s.SetEndReason(StreamEndReasonHandlerStop, err) + + assert.False(t, s.IsNormalEnd()) + assert.True(t, s.IsAbnormalEnd()) + assert.Equal(t, err, s.FatalError()) +} + +func TestStreamStatus_HandlerErrorSurvivesEOFRace(t *testing.T) { + t.Parallel() + s := NewStreamStatus() + s.SetEndReason(StreamEndReasonEOF, nil) + s.RecordFatalError(fmt.Errorf("handler failed after scanner EOF")) + + assert.False(t, s.IsNormalEnd()) + assert.True(t, s.IsAbnormalEnd()) + assert.Contains(t, s.Summary(), "handler failed after scanner EOF") +} + +func TestStreamStatus_HandlerStopWithoutErrorRemainsNormal(t *testing.T) { + t.Parallel() + s := NewStreamStatus() + s.SetEndReason(StreamEndReasonHandlerStop, nil) + + assert.True(t, s.IsNormalEnd()) + assert.False(t, s.IsAbnormalEnd()) +} + func TestStreamStatus_Summary(t *testing.T) { t.Parallel() diff --git a/relay/compatible_handler.go b/relay/compatible_handler.go index 9606d001..7f3d6f0b 100644 --- a/relay/compatible_handler.go +++ b/relay/compatible_handler.go @@ -78,6 +78,9 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (maxAPIError *types applySystemPromptIfNeeded(c, info, request) usage, maxApiErr := chatCompletionsViaResponses(c, info, adaptor, request) if maxApiErr != nil { + if usage != nil && info.IsStream && info.RelayFormat == types.RelayFormatOpenAI && types.IsSkipRetryError(maxApiErr) { + service.PostPartialConsumeQuota(c, info, usage) + } return maxApiErr } diff --git a/relay/file_search_stream_settlement_test.go b/relay/file_search_stream_settlement_test.go new file mode 100644 index 00000000..21948dc6 --- /dev/null +++ b/relay/file_search_stream_settlement_test.go @@ -0,0 +1,282 @@ +package relay + +import ( + "context" + "crypto/sha256" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/constant" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/model" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/service" + "github.com/MAX-API-Next/MAX-API/setting/config" + "github.com/MAX-API-Next/MAX-API/setting/operation_setting" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +// Cancellation follows a real tool-result flush, rather than the preceding +// created event or an upstream write that the adapter may not have processed. +type fileSearchCancelWriter struct { + gin.ResponseWriter + recorder *httptest.ResponseRecorder + cancel context.CancelFunc + triggered atomic.Bool +} + +func (w *fileSearchCancelWriter) Flush() { + w.ResponseWriter.Flush() + if strings.Contains(w.recorder.Body.String(), `"type":"file_search_call"`) && w.triggered.CompareAndSwap(false, true) { + w.cancel() + } +} + +func setFileSearchLedgerPriceForTest(t *testing.T, price float64) { + t.Helper() + prices := config.GlobalConfig.Get("tool_price_setting").(*operation_setting.ToolPriceSetting) + oldPrices := prices.Prices + prices.Prices = map[string]float64{dto.BuildInToolFileSearch: price} + operation_setting.RebuildToolPriceIndex() + t.Cleanup(func() { + prices.Prices = oldPrices + operation_setting.RebuildToolPriceIndex() + }) +} + +func TestFileSearchInterruptedStreamRetainsOriginalFunding(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitHttpClient() + service.InitTokenEncoders() + oldEmpty, oldRetries := common.EmptyCompletionRetryEnabled, common.RetryTimes + common.RetryTimes = 2 + setFileSearchLedgerPriceForTest(t, 0.02) + t.Cleanup(func() { + common.EmptyCompletionRetryEnabled, common.RetryTimes = oldEmpty, oldRetries + }) + for _, funding := range streamFundingCases() { + for _, usageKind := range []string{"known", "zero", "missing"} { + for _, retry := range []bool{false, true} { + for _, status := range []string{"completed", "legacy_missing"} { + for _, ending := range []string{"eof", "transport", "malformed", "cancel"} { + t.Run(fmt.Sprintf("%s/%s/retry_%t/%s/%s", funding.name, usageKind, retry, status, ending), func(t *testing.T) { + common.EmptyCompletionRetryEnabled = retry + db := setupStreamFundingLedger(t, funding) + _, frames := metadataFundingFrames(t, "responses_native", usageKind) + item := map[string]any{"type": "file_search_call", "id": "fs_completed", "queries": []string{"synthetic query"}, + "results": []any{map[string]any{"file_id": "file_synthetic", "text": "synthetic-search-result"}}} + if status == "completed" { + item["status"] = "completed" + } + done, err := common.Marshal(map[string]any{"type": "response.output_item.done", "output_index": 0, "item": item}) + require.NoError(t, err) + // The same completed invocation must not acquire a second fee. + frames = append(frames, string(done), string(done)) + if ending == "malformed" { + frames = append(frames, `{broken`) + } + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + if ending == "transport" { + w.Header().Set("Content-Length", "100000") + } + _, _ = fmt.Fprint(w, "data: "+strings.Join(frames, "\n\ndata: ")+"\n\n") + w.(http.Flusher).Flush() + if ending == "cancel" { + select { + case <-r.Context().Done(): + case <-time.After(3 * time.Second): + } + } + })) + t.Cleanup(upstream.Close) + fingerprint := sha256.Sum256([]byte(t.Name())) + requestID := fmt.Sprintf("file-search-%x", fingerprint[:16]) + c, info, recorder := streamFundingContext(t, "responses_native", requestID, upstream.URL, constant.ChannelTypeOpenAI, funding) + info.Request.(*dto.OpenAIResponsesRequest).Tools = []byte(`[{"type":"file_search","vector_store_ids":["vs_synthetic"]}]`) + info.ResponsesUsageInfo = &relaycommon.ResponsesUsageInfo{BuiltInTools: map[string]*relaycommon.BuildInToolInfo{ + dto.BuildInToolFileSearch: {ToolName: dto.BuildInToolFileSearch}, + }} + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + require.Equal(t, funding.source, info.BillingSource) + require.NoError(t, db.Create(&model.UserSubscription{Id: 956, UserId: 951, PlanId: 954, Status: "active", AmountTotal: 1000, + StartTime: time.Now().Unix() - 1800, EndTime: time.Now().Unix() + 1800}).Error) + info.UserSetting.BillingPreference = "wallet_only" + if funding.source == "wallet" { + info.UserSetting.BillingPreference = "subscription_only" + } + var cancelWriter *fileSearchCancelWriter + if ending == "cancel" { + ctx, cancel := context.WithCancel(c.Request.Context()) + t.Cleanup(cancel) + c.Request = c.Request.WithContext(ctx) + cancelWriter = &fileSearchCancelWriter{ResponseWriter: c.Writer, recorder: recorder, cancel: cancel} + c.Writer = cancelWriter + } + apiErr := ResponsesHelper(c, info) + require.NotNil(t, apiErr) + service.HandleFailedBilling(c, info, apiErr) + service.HandleFailedBilling(c, info, apiErr) + finalQuota, prompt, completion := 10, 0, 0 // 0.02 / 1000 * 500000; one actual call. + if usageKind == "known" { + finalQuota, prompt, completion = 16, 4, 2 + } + var settlement model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).First(&settlement).Error) + require.EqualValues(t, finalQuota-100, settlement.AppliedFundingDelta, "delivered File Search owes its fee, not a full refund") + require.EqualValues(t, finalQuota-100, settlement.AppliedTokenDelta) + require.Equal(t, funding.source, settlement.Source) + require.Equal(t, model.BillingSettlementStatusApplied, settlement.Status) + require.Equal(t, model.BillingSettlementEffectApplied, settlement.EffectStatus) + require.NotEmpty(t, settlement.EffectPayload) + require.True(t, types.IsSkipRetryError(apiErr)) + require.True(t, info.HasRecordedChannelFirstResult()) + require.False(t, info.FirstResultTime.IsZero()) + require.Equal(t, 1, info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolFileSearch].CallCount) + require.Contains(t, recorder.Body.String(), "synthetic-search-result") + require.NotContains(t, recorder.Body.String(), "response.completed") + require.EqualValues(t, 1, calls.Load()) + if cancelWriter != nil { + require.True(t, cancelWriter.triggered.Load()) + require.ErrorIs(t, c.Request.Context().Err(), context.Canceled) + } + assertStreamFundingBalances(t, db, funding, requestID, finalQuota, true, prompt, completion) + var log model.Log + require.NoError(t, db.Where("type = ?", model.LogTypeConsume).First(&log).Error) + var other map[string]any + require.NoError(t, common.UnmarshalJsonStr(log.Other, &other)) + require.EqualValues(t, 1, other["file_search_call_count"]) + require.EqualValues(t, 0.02, other["file_search_price"]) + var outbox int64 + require.NoError(t, db.Model(&model.CacheInvalidationTask{}).Count(&outbox).Error) + for range 2 { + applied, already, err := model.ApplyBillingSettlementOnce(model.BillingSettlementInput{ + OperationKey: settlement.OperationKey, Source: settlement.Source, UserID: settlement.UserID, + SubscriptionID: settlement.SubscriptionID, TokenID: settlement.TokenID, TokenKey: info.TokenKey, + FundingDelta: settlement.FundingDelta, TokenDelta: settlement.TokenDelta, + SubscriptionPreConsumeRequestID: settlement.SubscriptionPreConsumeRequestID, + FinalizeSubscriptionPreConsume: settlement.FinalizeSubscriptionPreConsume, + }) + require.NoError(t, err) + require.True(t, already) + require.EqualValues(t, finalQuota-100, applied) + model.ProcessPendingBillingSettlementsOnce() + } + var afterOutbox int64 + require.NoError(t, db.Model(&model.CacheInvalidationTask{}).Count(&afterOutbox).Error) + require.Equal(t, outbox, afterOutbox) + var otherSub model.UserSubscription + require.NoError(t, db.First(&otherSub, 956).Error) + require.Zero(t, otherSub.AmountUsed) + assertStreamFundingBalances(t, db, funding, requestID, finalQuota, true, prompt, completion) + }) + } + } + } + } + } +} + +func TestFileSearchLifecycleItemsRefundOriginalFunding(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitHttpClient() + service.InitTokenEncoders() + oldEmpty := common.EmptyCompletionRetryEnabled + common.EmptyCompletionRetryEnabled = false + t.Cleanup(func() { common.EmptyCompletionRetryEnabled = oldEmpty }) + setFileSearchLedgerPriceForTest(t, 0.02) + for _, funding := range streamFundingCases() { + for _, tc := range []struct{ name, event, status string }{ + {"added_running", dto.ResponsesOutputTypeItemAdded, "in_progress"}, + {"added_completed", dto.ResponsesOutputTypeItemAdded, "completed"}, + {"done_running", dto.ResponsesOutputTypeItemDone, "in_progress"}, + {"done_failed", dto.ResponsesOutputTypeItemDone, "failed"}, + {"done_cancelled", dto.ResponsesOutputTypeItemDone, "cancelled"}, + {"done_unknown", dto.ResponsesOutputTypeItemDone, "future_status"}, + {"nil_item", dto.ResponsesOutputTypeItemDone, ""}, + } { + t.Run(funding.name+"/"+tc.name, func(t *testing.T) { + db := setupStreamFundingLedger(t, funding) + _, frames := metadataFundingFrames(t, "responses_native", "known") + event := dto.ResponsesStreamResponse{Type: tc.event, Item: &dto.ResponsesOutput{Type: dto.BuildInCallFileSearchCall, ID: "fs_metadata", Status: tc.status}} + if tc.name == "nil_item" { + event.Item = nil + } + data, err := common.Marshal(event) + require.NoError(t, err) + frames = append(frames, string(data)) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, "data: "+strings.Join(frames, "\n\ndata: ")+"\n\n") + })) + t.Cleanup(upstream.Close) + requestID := "file-search-lifecycle-" + funding.name + "-" + tc.name + c, info, _ := streamFundingContext(t, "responses_native", requestID, upstream.URL, constant.ChannelTypeOpenAI, funding) + info.ResponsesUsageInfo = &relaycommon.ResponsesUsageInfo{BuiltInTools: map[string]*relaycommon.BuildInToolInfo{ + dto.BuildInToolFileSearch: {ToolName: dto.BuildInToolFileSearch}, + }} + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + apiErr := ResponsesHelper(c, info) + require.NotNil(t, apiErr) + service.HandleFailedBilling(c, info, apiErr) + service.HandleFailedBilling(c, info, apiErr) + require.False(t, info.HasRecordedChannelFirstResult()) + require.Zero(t, info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolFileSearch].CallCount) + assertStreamFundingBalances(t, db, funding, requestID, 0, false, 0, 0) + }) + } + } +} + +func TestFileSearchPartialSettlementAtAndAboveReservation(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitHttpClient() + service.InitTokenEncoders() + oldEmpty := common.EmptyCompletionRetryEnabled + common.EmptyCompletionRetryEnabled = false + t.Cleanup(func() { common.EmptyCompletionRetryEnabled = oldEmpty }) + for _, funding := range streamFundingCases() { + for _, quota := range []int{100, 110} { + t.Run(fmt.Sprintf("%s/quota_%d", funding.name, quota), func(t *testing.T) { + setFileSearchLedgerPriceForTest(t, float64(quota)/500) + db := setupStreamFundingLedger(t, funding) + _, frames := metadataFundingFrames(t, "responses_native", "zero") + frames = append(frames, `{"type":"response.output_item.done","output_index":0,"item":{"type":"file_search_call","id":"fs_boundary","status":"completed"}}`) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, "data: "+strings.Join(frames, "\n\ndata: ")+"\n\n") + })) + t.Cleanup(upstream.Close) + requestID := fmt.Sprintf("file-search-boundary-%s-%d", funding.name, quota) + c, info, _ := streamFundingContext(t, "responses_native", requestID, upstream.URL, constant.ChannelTypeOpenAI, funding) + info.ResponsesUsageInfo = &relaycommon.ResponsesUsageInfo{BuiltInTools: map[string]*relaycommon.BuildInToolInfo{ + dto.BuildInToolFileSearch: {ToolName: dto.BuildInToolFileSearch}, + }} + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + apiErr := ResponsesHelper(c, info) + require.NotNil(t, apiErr) + service.HandleFailedBilling(c, info, apiErr) + service.HandleFailedBilling(c, info, apiErr) + var settlement model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).First(&settlement).Error) + require.EqualValues(t, quota-100, settlement.AppliedFundingDelta) + require.EqualValues(t, quota-100, settlement.AppliedTokenDelta) + require.Equal(t, model.BillingSettlementStatusApplied, settlement.Status) + require.Equal(t, model.BillingSettlementEffectApplied, settlement.EffectStatus) + model.ProcessPendingBillingSettlementsOnce() + model.ProcessPendingBillingSettlementsOnce() + assertStreamFundingBalances(t, db, funding, requestID, quota, true, 0, 0) + }) + } + } +} diff --git a/relay/first_result_adapter_test.go b/relay/first_result_adapter_test.go new file mode 100644 index 00000000..bd8ccd00 --- /dev/null +++ b/relay/first_result_adapter_test.go @@ -0,0 +1,93 @@ +package relay + +import ( + "io" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/relay/channel/claude" + "github.com/MAX-API-Next/MAX-API/relay/channel/gemini" + "github.com/MAX-API-Next/MAX-API/relay/channel/openai" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/setting/operation_setting" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestResponsesAdaptersWaitForGenerationOutput(t *testing.T) { + for _, provider := range []struct { + name string + metadata, output string + forwardsMetadata bool + handler func(*gin.Context, *relaycommon.RelayInfo, *http.Response) (*dto.Usage, *types.MaxAPIError) + }{ + {"claude", `{"type":"message_start","message":{"id":"msg_test","model":"test-model","usage":{"input_tokens":10,"output_tokens":0}}}`, + `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}}`, true, + func(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.MaxAPIError) { + return claude.ClaudeStreamHandler(c, resp, info) + }}, + {"gemini", `{"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":0,"totalTokenCount":10}}`, + `{"candidates":[{"content":{"parts":[{"text":"hello"}]}}]}`, true, gemini.GeminiResponsesStreamHandler}, + {"chat_to_responses", `{"choices":[{"delta":{"role":"assistant"}}]}`, + `{"choices":[{"delta":{"content":"hello"}}]}`, true, openai.OaiChatToResponsesStreamHandler}, + {"responses", `{"type":"response.created","response":{"id":"resp_test","status":"in_progress","output":[]}}`, + `{"type":"response.output_text.delta","delta":"hello"}`, true, openai.OaiResponsesStreamHandler}, + {"responses_to_chat", `{"type":"response.created","response":{"id":"resp_test","status":"in_progress","output":[]}}`, + `{"type":"response.output_text.delta","delta":"hello"}`, false, openai.OaiResponsesToChatStreamHandler}, + } { + for _, hasOutput := range []bool{false, true} { + name := "metadata" + if hasOutput { + name = "output" + } + t.Run(provider.name+"/"+name, func(t *testing.T) { + oldSetting := *operation_setting.GetMonitorSetting() + *operation_setting.GetMonitorSetting() = operation_setting.MonitorSetting{StreamingFirstResultTimeoutSeconds: 1} + t.Cleanup(func() { *operation_setting.GetMonitorSetting() = oldSetting }) + oldRetry := common.EmptyCompletionRetryEnabled + common.EmptyCompletionRetryEnabled = false + t.Cleanup(func() { common.EmptyCompletionRetryEnabled = oldRetry }) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + info := &relaycommon.RelayInfo{IsStream: true, DisablePing: true, StartTime: time.Now(), RelayFormat: types.RelayFormatOpenAIResponses} + info.InitChannelMeta(c) + info.UpstreamModelName = "test-model" + info.SetEstimatePromptTokens(10) + if provider.name == "responses_to_chat" { + info.RelayFormat = types.RelayFormatOpenAI + } + r, w := io.Pipe() + finished := make(chan struct{}) + t.Cleanup(func() { _ = r.Close(); _ = w.Close(); <-finished }) + go func() { + defer close(finished) + defer w.Close() + if _, err := io.WriteString(w, "data: "+provider.metadata+"\n\n"); err != nil { + return + } + if hasOutput { + if _, err := io.WriteString(w, "data: "+provider.output+"\n\n"); err != nil { + return + } + } + time.Sleep(1500 * time.Millisecond) + }() + _, apiErr := provider.handler(c, info, &http.Response{StatusCode: 200, Body: r}) + require.Equal(t, hasOutput, info.HasRecordedChannelFirstResult()) + if hasOutput { + require.NotEqual(t, relaycommon.StreamEndReasonTimeout, info.StreamStatus.EndReason, "real output disables the first-result watchdog") + } else { + require.Equal(t, relaycommon.StreamEndReasonTimeout, info.StreamStatus.EndReason) + require.NotNil(t, apiErr) + require.Equal(t, types.ErrorCodeChannelResponseTimeExceeded, apiErr.GetErrorCode()) + require.Equal(t, provider.forwardsMetadata, types.IsSkipRetryError(apiErr)) + } + }) + } + } +} diff --git a/relay/helper/common.go b/relay/helper/common.go index 72fb8767..f968deaf 100644 --- a/relay/helper/common.go +++ b/relay/helper/common.go @@ -3,7 +3,9 @@ package helper import ( "errors" "fmt" + "io" "net/http" + "strings" "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/dto" @@ -12,6 +14,8 @@ import ( "github.com/gin-gonic/gin" "github.com/gorilla/websocket" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" ) func FlushWriter(c *gin.Context) (err error) { @@ -93,17 +97,101 @@ func ClaudeChunkData(c *gin.Context, resp dto.ClaudeResponse, data string) { } func ResponseChunkData(c *gin.Context, resp dto.ResponsesStreamResponse, data string) error { + _, err := ResponseChunkDataWithDelivery(c, resp, data) + return err +} + +// Delivery is true only after a complete SSE frame was written. Flush may +// still fail after delivery, and a short write must never count as a frame. +func ResponseChunkDataWithDelivery(c *gin.Context, resp dto.ResponsesStreamResponse, data string) (bool, error) { if c == nil || c.Writer == nil { - return errors.New("context or writer is nil") + return false, errors.New("context or writer is nil") + } + state := getResponsesStreamState(c) + state.mu.Lock() + defer state.mu.Unlock() + return writeResponsesStreamFrame(c, state, resp.Type, data) +} + +// The caller holds the request-local lock, including while flushing. A complete +// write still counts as delivery if flush fails, but no later frame may follow. +func writeResponsesStreamFrame(c *gin.Context, state *responsesStreamState, eventType, data string) (bool, error) { + if state.writeErr != nil { + return false, state.writeErr + } + if state.terminal { + return false, errResponsesStreamClosed } if requestContextDone(c) { - return fmt.Errorf("request context done: %w", c.Request.Context().Err()) + state.writeErr = fmt.Errorf("request context done: %w", c.Request.Context().Err()) + return false, state.writeErr + } + if (strings.HasPrefix(eventType, "response.") || eventType == "error") && !gjson.Get(data, "sequence_number").Exists() { + var err error + data, err = sjson.Set(data, "sequence_number", state.nextSequence) + if err != nil { + return false, err + } } - c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)}) - c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("data: %s", data)}) - return FlushWriter(c) + SetEventStreamHeaders(c) + frame := fmt.Sprintf("event: %s\ndata: %s\n\n", eventType, strings.ReplaceAll(data, "\r", "\\r")) + n, err := c.Writer.Write([]byte(frame)) + if err != nil { + state.writeErr = err + return false, err + } + if n != len(frame) { + state.writeErr = io.ErrShortWrite + return false, io.ErrShortWrite + } + state.recordDelivery(eventType, data) + err = flushResponsesWriter(c) + if err != nil { + state.writeErr = err + } + return true, err +} + +// ResponsesStreamEventHasOutput distinguishes delivered generation data from +// lifecycle metadata when deciding whether an interrupted stream has usage. +func ResponsesStreamEventHasOutput(event dto.ResponsesStreamResponse) bool { + switch event.Type { + case "response.output_text.delta", "response.reasoning_text.delta", "response.reasoning_summary_text.delta", + "response.refusal.delta", "response.function_call_arguments.delta", "response.custom_tool_call_input.delta": + return event.Delta != "" + case dto.ResponsesOutputTypeItemAdded, dto.ResponsesOutputTypeItemDone: + if event.Item != nil { + if event.Item.Type == dto.BuildInCallFileSearchCall { + // A search lifecycle/failed item is metadata. Only a finished + // invocation carries a delivered result and its existing fee. + status := strings.ToLower(strings.TrimSpace(event.Item.Status)) + return event.Type == dto.ResponsesOutputTypeItemDone && (status == "" || status == "completed") + } + return event.Item.Type == dto.BuildInCallFunctionCall || event.Item.Type == dto.BuildInCallCustomToolCall + } + } + return false +} + +// First-result timing excludes whitespace-only metadata/text while retaining +// the independent output-delivery predicate used for partial usage accounting. +func ResponsesStreamEventHasFirstResult(event dto.ResponsesStreamResponse) bool { + if !ResponsesStreamEventHasOutput(event) { + return false + } + if event.Type == dto.ResponsesOutputTypeItemAdded || event.Type == dto.ResponsesOutputTypeItemDone { + return true + } + return strings.TrimSpace(event.Delta) != "" +} + +// An empty usage object is not a reported zero. Require both numeric counters +// before treating an otherwise empty usage payload as authoritative. +func HasUsageTokenFields(data, path, inputField, outputField string) bool { + return gjson.Get(data, path+"."+inputField).Type == gjson.Number && + gjson.Get(data, path+"."+outputField).Type == gjson.Number } func StringData(c *gin.Context, str string) error { diff --git a/relay/helper/first_result_timeout.go b/relay/helper/first_result_timeout.go index d7d17a01..33da88c7 100644 --- a/relay/helper/first_result_timeout.go +++ b/relay/helper/first_result_timeout.go @@ -6,21 +6,25 @@ import ( relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" ) // FirstResultTimeoutError converts a stream scanner timeout that happened -// before any channel response was recorded into the retryable API error used -// by relay settlement and retry handling. A stream that already delivered a -// channel response keeps its existing terminal handling. -func FirstResultTimeoutError(info *relaycommon.RelayInfo) *types.MaxAPIError { +// before a generation payload into the API error used by relay settlement and +// retry handling. Even metadata or ping delivery forbids transparent replay. +func FirstResultTimeoutError(c *gin.Context, info *relaycommon.RelayInfo) *types.MaxAPIError { if info == nil || info.StreamStatus == nil || info.StreamStatus.EndReason != relaycommon.StreamEndReasonTimeout || - info.HasRecordedChannelFirstResponse() { + info.HasRecordedChannelFirstResult() { return nil } - return types.NewOpenAIError( - errors.New("upstream stream timed out before the first response"), + apiErr := types.NewOpenAIError( + errors.New("upstream stream timed out before the first result"), types.ErrorCodeChannelResponseTimeExceeded, http.StatusRequestTimeout, ) + if info.HasRecordedChannelFirstResponse() || (c != nil && c.Writer != nil && c.Writer.Written()) { + types.ErrOptionWithSkipRetry()(apiErr) + } + return apiErr } diff --git a/relay/helper/first_result_timeout_test.go b/relay/helper/first_result_timeout_test.go index ec91fe77..73d30772 100644 --- a/relay/helper/first_result_timeout_test.go +++ b/relay/helper/first_result_timeout_test.go @@ -1,10 +1,12 @@ package helper import ( + "net/http/httptest" "testing" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) @@ -12,19 +14,41 @@ func TestFirstResultTimeoutErrorRequiresMissingFirstResponse(t *testing.T) { info := &relaycommon.RelayInfo{StreamStatus: relaycommon.NewStreamStatus()} info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonTimeout, nil) - err := FirstResultTimeoutError(info) + err := FirstResultTimeoutError(nil, info) require.NotNil(t, err) require.Equal(t, types.ErrorCodeChannelResponseTimeExceeded, err.GetErrorCode()) info.SetFirstResponseTime() - require.Nil(t, FirstResultTimeoutError(info)) + require.Nil(t, FirstResultTimeoutError(nil, info)) info.StreamStatus = relaycommon.NewStreamStatus() info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonDone, nil) - require.Nil(t, FirstResultTimeoutError(info)) + require.Nil(t, FirstResultTimeoutError(nil, info)) +} + +func TestFirstResultTimeoutCannotReplayCommittedMetadataOrPing(t *testing.T) { + for _, committed := range []string{"none", "metadata", "ping"} { + t.Run(committed, func(t *testing.T) { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + info := &relaycommon.RelayInfo{StreamStatus: relaycommon.NewStreamStatus()} + info.EnableFirstResultTracking() + if committed == "metadata" { + info.SetFirstResponseTime() + } + if committed == "ping" { + _, _ = c.Writer.WriteString(": PING\n\n") + } + info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonTimeout, nil) + apiErr := FirstResultTimeoutError(c, info) + require.NotNil(t, apiErr) + require.Equal(t, committed != "none", types.IsSkipRetryError(apiErr)) + info.SetFirstResultTime() + require.Nil(t, FirstResultTimeoutError(c, info), "output-idle is not a first-result timeout") + }) + } } func TestFirstResultTimeoutErrorHandlesNilState(t *testing.T) { - require.Nil(t, FirstResultTimeoutError(nil)) - require.Nil(t, FirstResultTimeoutError(&relaycommon.RelayInfo{})) + require.Nil(t, FirstResultTimeoutError(nil, nil)) + require.Nil(t, FirstResultTimeoutError(nil, &relaycommon.RelayInfo{})) } diff --git a/relay/helper/responses_stream_state.go b/relay/helper/responses_stream_state.go new file mode 100644 index 00000000..c22ee044 --- /dev/null +++ b/relay/helper/responses_stream_state.go @@ -0,0 +1,156 @@ +package helper + +import ( + "errors" + "fmt" + "net/http" + "sync" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" +) + +const responsesStreamStateKey = "max_responses_stream_state" + +var errResponsesStreamClosed = errors.New("responses stream is already terminal") + +// Request-local downstream state, independent of provider status and billing. +// Keep only the latest lifecycle snapshot, not an unbounded history of deltas. +type responsesStreamState struct { + mu sync.Mutex + nextSequence int64 + response string + terminal bool + writeErr error +} + +func getResponsesStreamState(c *gin.Context) *responsesStreamState { + if value, ok := c.Get(responsesStreamStateKey); ok { + return value.(*responsesStreamState) + } + state := &responsesStreamState{} + c.Set(responsesStreamStateKey, state) + return state +} + +func isResponsesTerminalEvent(eventType string) bool { + switch eventType { + case "response.completed", "response.incomplete", "response.failed", "response.done", + "response.cancelled", "response.canceled", "error", "response.error": + return true + default: + return false + } +} + +// A failed ping can happen before any Responses event reaches the client. +func LatchResponsesStreamWriteError(c *gin.Context, err error) { + if c == nil || err == nil { + return + } + state := getResponsesStreamState(c) + state.mu.Lock() + defer state.mu.Unlock() + if state.writeErr == nil { + state.writeErr = err + } +} + +// Gin's Flush delegates to http.Flusher and discards FlushError. Prefer a +// supported error-returning flush through wrappers; retain legacy fallback. +func flushResponsesWriter(c *gin.Context) (err error) { + defer func() { + if r := recover(); r != nil { + err = fmt.Errorf("flush panic recovered: %v", r) + } + }() + if requestContextDone(c) { + return c.Request.Context().Err() + } + var writer http.ResponseWriter = c.Writer + for { + if flusher, ok := writer.(interface{ FlushError() error }); ok { + return flusher.FlushError() + } + unwrapper, ok := writer.(interface{ Unwrap() http.ResponseWriter }) + if !ok { + break + } + writer = unwrapper.Unwrap() + if writer == nil { + break + } + } + return FlushWriter(c) +} + +// WriteResponsesStreamFailure is the controller's fallback after commitment. +// An already delivered terminal or a broken downstream owns the outcome; no +// later error writer may append to it. Return errors for observability only. +func WriteResponsesStreamFailure(c *gin.Context, apiErr *types.MaxAPIError) error { + state := getResponsesStreamState(c) + state.mu.Lock() + defer state.mu.Unlock() + if state.writeErr != nil { + return state.writeErr + } + if state.terminal { + return nil + } + if requestContextDone(c) { + return c.Request.Context().Err() + } + oaiErr := apiErr.ToOpenAIError() + var errorCode any + if oaiErr.Code != nil { + errorCode = common.Interface2String(oaiErr.Code) + } + event := map[string]any{"type": "error", "code": errorCode, "message": oaiErr.Message, "param": nil} + if oaiErr.Param != "" { + event["param"] = oaiErr.Param + } + if state.response != "" { + var response map[string]any + if err := common.UnmarshalJsonStr(state.response, &response); err != nil { + return err + } + // No synthetic provider identity when the stream has no response object. + if id, _ := response["id"].(string); id != "" { + code := "server_error" + if common.Interface2String(oaiErr.Code) == "rate_limit_exceeded" { + code = "rate_limit_exceeded" + } + response["status"] = "failed" + response["error"] = map[string]any{"code": code, "message": oaiErr.Message} + response["incomplete_details"] = nil + // Usage in an earlier lifecycle snapshot is not final usage. + response["usage"] = nil + event = map[string]any{"type": "response.failed", "response": response} + } + } + event["sequence_number"] = state.nextSequence + data, err := common.Marshal(event) + if err != nil { + return err + } + _, err = writeResponsesStreamFrame(c, state, event["type"].(string), string(data)) + return err +} + +func (state *responsesStreamState) recordDelivery(eventType, data string) { + if sequence := gjson.Get(data, "sequence_number"); sequence.Exists() && sequence.Type == gjson.Number { + if next := sequence.Int() + 1; next > state.nextSequence { + state.nextSequence = next + } + } else { + state.nextSequence++ + } + if response := gjson.Get(data, "response"); response.IsObject() { + state.response = response.Raw + } + if isResponsesTerminalEvent(eventType) { + state.terminal = true + } +} diff --git a/relay/helper/responses_stream_state_test.go b/relay/helper/responses_stream_state_test.go new file mode 100644 index 00000000..c306e47b --- /dev/null +++ b/relay/helper/responses_stream_state_test.go @@ -0,0 +1,103 @@ +package helper + +import ( + "errors" + "io" + "net/http/httptest" + "strings" + "testing" + + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +type responsesFlushErrorWriter struct { + *httptest.ResponseRecorder + writes int + flushes int + err error +} + +func (w *responsesFlushErrorWriter) Write(p []byte) (int, error) { + w.writes++ + return w.ResponseRecorder.Write(p) +} + +func (w *responsesFlushErrorWriter) FlushError() error { + w.flushes++ + return w.err +} + +func TestResponsesFullWriteFlushErrorLatchesDelivery(t *testing.T) { + w := &responsesFlushErrorWriter{ResponseRecorder: httptest.NewRecorder(), err: io.ErrClosedPipe} + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("POST", "/v1/responses", nil) + delivered, err := ResponseChunkDataWithDelivery(c, dto.ResponsesStreamResponse{Type: "response.output_text.delta"}, `{"type":"response.output_text.delta","delta":"partial"}`) + require.True(t, delivered, "full write remains billable delivery when flush fails") + require.ErrorIs(t, err, io.ErrClosedPipe) + require.Equal(t, 1, w.writes) + require.Equal(t, 1, w.flushes) + before := w.Body.String() + apiErr := types.NewError(errors.New("upstream failure"), types.ErrorCodeBadResponse) + require.ErrorIs(t, WriteResponsesStreamFailure(c, apiErr), io.ErrClosedPipe) + delivered, err = ResponseChunkDataWithDelivery(c, dto.ResponsesStreamResponse{Type: "response.failed"}, `{"type":"response.failed"}`) + require.False(t, delivered) + require.ErrorIs(t, err, io.ErrClosedPipe) + require.Equal(t, 1, w.writes) + require.Equal(t, before, w.Body.String()) +} + +func TestResponsesMissingIdentityUsesErrorEventAndSequence(t *testing.T) { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("POST", "/v1/responses", nil) + require.NoError(t, ResponseChunkData(c, dto.ResponsesStreamResponse{Type: "response.output_text.delta"}, `{"type":"response.output_text.delta","delta":"partial"}`)) + apiErr := types.NewError(errors.New("local failure"), types.ErrorCodeBadResponse) + require.NoError(t, WriteResponsesStreamFailure(c, apiErr)) + frames := strings.Split(strings.TrimSpace(w.Body.String()), "\n\n") + require.Len(t, frames, 2) + for idx, frame := range frames { + data := strings.SplitN(frame, "data: ", 2)[1] + require.EqualValues(t, idx, gjson.Get(data, "sequence_number").Int()) + } + require.True(t, strings.HasPrefix(frames[1], "event: error\n")) + require.NotContains(t, frames[1], "resp_") + require.NotContains(t, frames[1], "[DONE]") + before := w.Body.String() + require.NoError(t, WriteResponsesStreamFailure(c, apiErr)) + require.Equal(t, before, w.Body.String()) + _, err := ResponseChunkDataWithDelivery(c, dto.ResponsesStreamResponse{Type: "response.completed"}, `{"type":"response.completed"}`) + require.ErrorIs(t, err, errResponsesStreamClosed) + require.Equal(t, before, w.Body.String()) +} + +func TestResponsesSequencePreservesProviderAndNonResponsesFrames(t *testing.T) { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("POST", "/v1/responses", nil) + raw := `{"type":"response.extension","sequence_number":41,"extension":{"keep":true}}` + require.NoError(t, ResponseChunkData(c, dto.ResponsesStreamResponse{Type: "response.extension"}, raw)) + require.Contains(t, w.Body.String(), "data: "+raw+"\n\n") + require.NoError(t, ResponseChunkData(c, dto.ResponsesStreamResponse{Type: "response.output_text.delta"}, `{"type":"response.output_text.delta","delta":"partial"}`)) + last := strings.Split(strings.TrimSpace(w.Body.String()), "\n\n")[1] + require.EqualValues(t, 42, gjson.Get(strings.SplitN(last, "data: ", 2)[1], "sequence_number").Int()) + image := `{"type":"image_generation.partial_image","b64_json":"synthetic"}` + require.NoError(t, ResponseChunkData(c, dto.ResponsesStreamResponse{Type: "image_generation.partial_image"}, image)) + require.Contains(t, w.Body.String(), "data: "+image+"\n\n") +} + +func TestResponsesPingFailureSuppressesFallbackWithoutEvent(t *testing.T) { + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest("POST", "/v1/responses", nil) + SetEventStreamHeaders(c) + require.NoError(t, PingData(c)) + LatchResponsesStreamWriteError(c, io.ErrClosedPipe) + before := w.Body.String() + apiErr := types.NewError(errors.New("timeout"), types.ErrorCodeBadResponse) + require.ErrorIs(t, WriteResponsesStreamFailure(c, apiErr), io.ErrClosedPipe) + require.Equal(t, before, w.Body.String()) +} diff --git a/relay/helper/responses_write_test.go b/relay/helper/responses_write_test.go new file mode 100644 index 00000000..6c15d6d0 --- /dev/null +++ b/relay/helper/responses_write_test.go @@ -0,0 +1,37 @@ +package helper + +import ( + "errors" + "io" + "net/http/httptest" + "testing" + + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type responsesFailureWriter struct { + *httptest.ResponseRecorder + err error +} + +func (w *responsesFailureWriter) Write(p []byte) (int, error) { + if w.err != nil { + return 0, w.err + } + return len(p) - 1, nil +} + +func TestResponseChunkDataPropagatesWriteFailure(t *testing.T) { + for _, writeErr := range []error{errors.New("synthetic write failure"), nil} { + w := &responsesFailureWriter{ResponseRecorder: httptest.NewRecorder(), err: writeErr} + c, _ := gin.CreateTestContext(w) + err := ResponseChunkData(c, dto.ResponsesStreamResponse{Type: "response.completed"}, `{}`) + if writeErr == nil { + require.ErrorIs(t, err, io.ErrShortWrite) + } else { + require.ErrorIs(t, err, writeErr) + } + } +} diff --git a/relay/helper/stream_result.go b/relay/helper/stream_result.go index 2ae8967d..7d2916a4 100644 --- a/relay/helper/stream_result.go +++ b/relay/helper/stream_result.go @@ -32,6 +32,7 @@ func (r *StreamResult) Error(err error) { func (r *StreamResult) Stop(err error) { if err != nil { r.status.RecordError(err.Error()) + r.status.RecordFatalError(err) } r.status.SetEndReason(relaycommon.StreamEndReasonHandlerStop, err) r.stopped = true diff --git a/relay/helper/stream_scanner.go b/relay/helper/stream_scanner.go index 8c432e22..8e3be8e8 100644 --- a/relay/helper/stream_scanner.go +++ b/relay/helper/stream_scanner.go @@ -15,6 +15,7 @@ import ( "github.com/MAX-API-Next/MAX-API/logger" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" "github.com/MAX-API-Next/MAX-API/setting/operation_setting" + "github.com/MAX-API-Next/MAX-API/types" "github.com/bytedance/gopkg/util/gopool" @@ -77,6 +78,10 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon if resp == nil || dataHandler == nil { return } + if info.RelayFormat == types.RelayFormatOpenAIResponses { + // Initialize before workers; ping failure and event delivery share state. + getResponsesStreamState(c) + } // Preserve an existing StreamStatus so callers can carry forward prior stream errors. if info.StreamStatus == nil { @@ -95,7 +100,7 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon } var firstResultTimer *time.Timer var firstResultTimeoutChan <-chan time.Time - if firstResultTimeoutSeconds > 0 && info.FirstResponseSignal() != nil { + if firstResultTimeoutSeconds > 0 && info.FirstResultSignal() != nil { firstResultTimeout := relaycommon.RemainingChannelFirstResponseTimeout( info.ChannelAttemptStartTime(), time.Duration(firstResultTimeoutSeconds)*time.Second, @@ -199,6 +204,9 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon err = pingDataWithTimeout(c) }() if err != nil { + if info.RelayFormat == types.RelayFormatOpenAIResponses { + LatchResponsesStreamWriteError(c, err) + } logger.LogError(c, "ping data error: "+err.Error()) info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonPingFail, err) stop() @@ -328,7 +336,7 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon // 主循环等待完成或超时。The first-result channel disables its timer as // soon as a handler confirms a deliverable payload; subsequent idle timeouts // retain their existing diagnostic-only semantics. - firstResultSignal := info.FirstResponseSignal() + firstResultSignal := info.FirstResultSignal() waitLoop: for { select { @@ -338,7 +346,7 @@ waitLoop: continue case <-firstResultTimeoutChan: firstResultTimeoutChan = nil - if info.FirstResponseDeadlineExpired() { + if info.FirstResultDeadlineExpired() { info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonTimeout, nil) info.RecordChannelTimeout("streaming_first_result_timeout") break waitLoop diff --git a/relay/metadata_funding_settlement_test.go b/relay/metadata_funding_settlement_test.go new file mode 100644 index 00000000..6b98c717 --- /dev/null +++ b/relay/metadata_funding_settlement_test.go @@ -0,0 +1,246 @@ +package relay + +import ( + "context" + "crypto/sha256" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/constant" + "github.com/MAX-API-Next/MAX-API/model" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/service" + "github.com/MAX-API-Next/MAX-API/setting/operation_setting" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func metadataFundingFrames(t *testing.T, route, usage string) (int, []string) { + t.Helper() + prompt, completion := 4, 2 + if usage == "zero" { + prompt, completion = 0, 0 + } + encode := func(v map[string]any) string { + data, err := common.Marshal(v) + require.NoError(t, err) + return string(data) + } + switch route { + case "chat_native", "chat_to_responses": + frame := map[string]any{"id": "chat_metadata", "choices": []any{map[string]any{"index": 0, "delta": map[string]any{"role": "assistant"}}}} + if usage != "missing" { + frame["usage"] = map[string]any{"prompt_tokens": prompt, "completion_tokens": completion, "total_tokens": prompt + completion} + } + channel := constant.ChannelTypeOpenAI + if route == "chat_to_responses" { + channel = constant.ChannelTypeAdvancedCustom + } + return channel, []string{encode(frame)} + case "claude_to_responses": + message := map[string]any{"id": "msg_metadata", "model": "gpt-test", "role": "assistant", "content": []any{}} + if usage != "missing" { + message["usage"] = map[string]any{"input_tokens": prompt, "output_tokens": completion} + } + return constant.ChannelTypeAnthropic, []string{encode(map[string]any{"type": "message_start", "message": message})} + case "gemini_to_responses": + frame := map[string]any{"candidates": []any{map[string]any{"content": map[string]any{"role": "model", "parts": []any{}}}}} + if usage != "missing" { + frame["usageMetadata"] = map[string]any{"promptTokenCount": prompt, "candidatesTokenCount": completion, "totalTokenCount": prompt + completion} + } + return constant.ChannelTypeGemini, []string{encode(frame)} + default: + response := map[string]any{"id": "resp_metadata", "model": "gpt-test", "status": "in_progress", "output": []any{}} + if usage != "missing" { + response["usage"] = map[string]any{"input_tokens": prompt, "output_tokens": completion, "total_tokens": prompt + completion} + } + channel := constant.ChannelTypeOpenAI + if route == "responses_to_chat_custom" { + channel = constant.ChannelTypeAdvancedCustom + } + return channel, []string{encode(map[string]any{"type": "response.created", "response": response})} + } +} + +// Cancel at the real downstream flush, after metadata was handled. A server +// write alone would not prove the adapter had processed the metadata. +type metadataCancelWriter struct { + gin.ResponseWriter + cancel context.CancelFunc + triggered atomic.Bool +} + +func (w *metadataCancelWriter) Flush() { + w.ResponseWriter.Flush() + if w.Size() > 0 && w.triggered.CompareAndSwap(false, true) { + w.cancel() + } +} + +func TestMetadataOnlyFailureRetainsOriginalFunding(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitHttpClient() + service.InitTokenEncoders() + oldEmptyRetry := common.EmptyCompletionRetryEnabled + common.EmptyCompletionRetryEnabled = false + t.Cleanup(func() { common.EmptyCompletionRetryEnabled = oldEmptyRetry }) + for _, route := range []string{"responses_native", "chat_native", "chat_to_responses", "responses_to_chat_custom", "responses_to_chat_global", "claude_to_responses", "gemini_to_responses"} { + for _, funding := range streamFundingCases() { + for _, usage := range []string{"known", "zero", "missing"} { + for _, ending := range []string{"eof", "transport", "malformed", "cancel", "timeout"} { + // Native Chat's normal EOF keeps its legacy empty-completion + // pricing policy. Its abnormal endings belong in this matrix. + if route == "chat_native" && ending == "eof" { + continue + } + bufferedMetadata := strings.HasPrefix(route, "responses_to_chat_") + // These converters defer Chat's role event until output. There + // is no downstream metadata flush at which to cancel a client. + if bufferedMetadata && ending == "cancel" { + continue + } + t.Run(strings.Join([]string{route, funding.name, usage, ending}, "/"), func(t *testing.T) { + db := setupStreamFundingLedger(t, funding) + oldMonitor := *operation_setting.GetMonitorSetting() + *operation_setting.GetMonitorSetting() = operation_setting.MonitorSetting{StreamingFirstResultTimeoutSeconds: 1} + t.Cleanup(func() { *operation_setting.GetMonitorSetting() = oldMonitor }) + channelType, frames := metadataFundingFrames(t, route, usage) + if ending == "malformed" { + frames = append(frames, `{broken`) + } + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + if ending == "transport" { + w.Header().Set("Content-Length", "100000") + } + _, _ = fmt.Fprint(w, "data: "+strings.Join(frames, "\n\ndata: ")+"\n\n") + w.(http.Flusher).Flush() + if ending == "timeout" || ending == "cancel" { + select { + case <-r.Context().Done(): + case <-time.After(3 * time.Second): + } + } + })) + t.Cleanup(upstream.Close) + fingerprint := sha256.Sum256([]byte(t.Name())) + requestID := fmt.Sprintf("metadata-%x", fingerprint[:16]) + c, info, recorder := streamFundingContext(t, route, requestID, upstream.URL, channelType, funding) + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + require.Equal(t, funding.source, info.BillingSource) + var selection model.BillingPreConsumeSelection + require.NoError(t, db.Where("request_id = ?", requestID).First(&selection).Error) + require.Equal(t, funding.source, selection.Source) + require.EqualValues(t, 100, selection.EffectiveQuota) + // A newly preferred subscription or changed preference must + // never receive this request's failure refund. + require.NoError(t, db.Create(&model.UserSubscription{Id: 956, UserId: 951, PlanId: 954, Status: "active", AmountTotal: 1000, + StartTime: time.Now().Unix() - 1800, EndTime: time.Now().Unix() + 1800}).Error) + info.UserSetting.BillingPreference = "wallet_only" + if funding.source == "wallet" { + info.UserSetting.BillingPreference = "subscription_only" + } + var cancelWriter *metadataCancelWriter + if ending == "cancel" { + ctx, cancel := context.WithCancel(c.Request.Context()) + t.Cleanup(cancel) + c.Request = c.Request.WithContext(ctx) + cancelWriter = &metadataCancelWriter{ResponseWriter: c.Writer, cancel: cancel} + c.Writer = cancelWriter + } + var apiErr *types.MaxAPIError + if info.RelayFormat == types.RelayFormatOpenAI { + apiErr = TextHelper(c, info) + } else { + apiErr = ResponsesHelper(c, info) + } + require.NotNil(t, apiErr) + require.False(t, info.HasRecordedChannelFirstResult()) + require.True(t, info.FirstResultTime.IsZero()) + if bufferedMetadata { + require.Empty(t, recorder.Body.String(), "created metadata remains buffered for Chat clients") + } else { + require.NotEmpty(t, recorder.Body.String(), "metadata was actually delivered") + } + require.NotContains(t, recorder.Body.String(), "response.completed") + if ending == "timeout" { + require.Equal(t, types.ErrorCodeChannelResponseTimeExceeded, apiErr.GetErrorCode()) + require.Equal(t, relaycommon.StreamEndReasonTimeout, info.StreamStatus.EndReason) + } + if cancelWriter != nil { + require.True(t, cancelWriter.triggered.Load()) + require.ErrorIs(t, c.Request.Context().Err(), context.Canceled) + } + if !bufferedMetadata { + require.True(t, types.IsSkipRetryError(apiErr), "delivered metadata closes transparent replay") + } + service.HandleFailedBilling(c, info, apiErr) + service.HandleFailedBilling(c, info, apiErr) + require.EqualValues(t, 1, calls.Load()) + var settlement model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).First(&settlement).Error) + require.Equal(t, funding.source, settlement.Source) + require.Equal(t, model.BillingSettlementStatusApplied, settlement.Status) + require.EqualValues(t, -100, settlement.AppliedFundingDelta) + require.EqualValues(t, -100, settlement.AppliedTokenDelta) + require.Empty(t, settlement.EffectPayload, "metadata-only failure has no consume projection") + assertStreamFundingBalances(t, db, funding, requestID, 0, false, 0, 0) + var outbox int64 + require.NoError(t, db.Model(&model.CacheInvalidationTask{}).Count(&outbox).Error) + for range 2 { + _, already, err := model.ApplyBillingSettlementOnce(model.BillingSettlementInput{ + OperationKey: settlement.OperationKey, Source: settlement.Source, UserID: settlement.UserID, + SubscriptionID: settlement.SubscriptionID, TokenID: settlement.TokenID, TokenKey: info.TokenKey, + FundingDelta: -100, TokenDelta: -100, SubscriptionPreConsumeRequestID: settlement.SubscriptionPreConsumeRequestID, + FinalizeSubscriptionPreConsume: settlement.FinalizeSubscriptionPreConsume, AllowMissingToken: true, + }) + require.NoError(t, err) + require.True(t, already) + model.ProcessPendingBillingSettlementsOnce() + } + var afterOutbox int64 + require.NoError(t, db.Model(&model.CacheInvalidationTask{}).Count(&afterOutbox).Error) + require.Equal(t, outbox, afterOutbox) + var other model.UserSubscription + require.NoError(t, db.First(&other, 956).Error) + require.Zero(t, other.AmountUsed) + assertStreamFundingBalances(t, db, funding, requestID, 0, false, 0, 0) + }) + } + } + } + } +} + +func TestMetadataOnlyCompletedResponsesPreservesReportedUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitHttpClient() + service.InitTokenEncoders() + oldEmptyRetry := common.EmptyCompletionRetryEnabled + common.EmptyCompletionRetryEnabled = false + t.Cleanup(func() { common.EmptyCompletionRetryEnabled = oldEmptyRetry }) + for _, funding := range streamFundingCases() { + t.Run(funding.name, func(t *testing.T) { + db := setupStreamFundingLedger(t, funding) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, "data: "+`{"type":"response.completed","response":{"id":"resp_empty_done","status":"completed","output":[],"usage":{"input_tokens":4,"output_tokens":2,"total_tokens":6}}}`+"\n\n") + })) + t.Cleanup(upstream.Close) + requestID := "metadata-complete-" + funding.name + c, info, _ := streamFundingContext(t, "responses_native", requestID, upstream.URL, constant.ChannelTypeOpenAI, funding) + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + require.Nil(t, ResponsesHelper(c, info), "a real terminal event keeps normal reported-usage billing") + assertStreamFundingBalances(t, db, funding, requestID, 6, true, 4, 2) + }) + } +} diff --git a/relay/request_settlement_replay_test.go b/relay/request_settlement_replay_test.go new file mode 100644 index 00000000..7ef72837 --- /dev/null +++ b/relay/request_settlement_replay_test.go @@ -0,0 +1,424 @@ +package relay + +import ( + "crypto/sha256" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/model" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/service" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +var requestReplayRoutes = []string{"responses_native", "chat_native", "chat_to_responses", "responses_to_chat_custom", "responses_to_chat_global", "claude_to_responses", "gemini_to_responses"} + +type requestReplayFixture struct { + db *gorm.DB + funding streamFundingCase + requestID string + c *gin.Context + info *relaycommon.RelayInfo + record model.BillingSettlement + effect model.BillingSettlementEffect + quota, prompt, complete int + rollback *atomic.Bool + sessions []*service.BillingSession +} + +func initRequestReplayTests(t *testing.T) { + t.Helper() + gin.SetMode(gin.TestMode) + service.InitHttpClient() + service.InitTokenEncoders() + oldEmpty, oldRetries := common.EmptyCompletionRetryEnabled, common.RetryTimes + common.EmptyCompletionRetryEnabled, common.RetryTimes = false, 2 + t.Cleanup(func() { common.EmptyCompletionRetryEnabled, common.RetryTimes = oldEmpty, oldRetries }) +} + +// Each fixture produces its finalize intent through real admission, HTTP and +// an adapter/Helper. The injected error is inside the funding transaction, +// immediately before its applied marker; all balance legs must roll back. +func newRequestReplayFixture(t *testing.T, route string, funding streamFundingCase, usage string, pending bool, sessionCount int) *requestReplayFixture { + t.Helper() + f := &requestReplayFixture{db: setupStreamFundingLedger(t, funding), funding: funding, prompt: 4, complete: 2, rollback: &atomic.Bool{}} + switch usage { + case "at_reservation": + f.prompt, f.complete = 80, 20 + case "above_reservation": + f.prompt, f.complete = 80, 50 + case "zero": + f.prompt, f.complete = 0, 0 + } + f.quota = f.prompt + f.complete + channelType, frames := streamFundingFrames(t, route, usage) + frames = append(frames, `{broken`) + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, "data: "+strings.Join(frames, "\n\ndata: ")+"\n\n") + })) + t.Cleanup(upstream.Close) + fingerprint := sha256.Sum256([]byte(t.Name())) + f.requestID = fmt.Sprintf("request-replay-%x", fingerprint[:16]) + f.c, f.info, _ = streamFundingContext(t, route, f.requestID, upstream.URL, channelType, funding) + require.Nil(t, service.PreConsumeBilling(f.c, 100, f.info)) + require.Equal(t, funding.source, f.info.BillingSource) + // Sessions are reconstructed before unresolved positive funding blocks new + // admission. Their pre-consume replay must not reserve a second time. + for range sessionCount { + info := &relaycommon.RelayInfo{RequestId: f.requestID, UserId: 951, TokenId: 952, TokenKey: f.info.TokenKey, + OriginModelName: f.info.OriginModelName, ForcePreConsume: true, UserSetting: dto.UserSetting{BillingPreference: "wallet_only"}} + if funding.source == "wallet" { + info.UserSetting.BillingPreference = "subscription_only" + } + session, apiErr := service.NewBillingSession(f.c, info, 100) + require.Nil(t, apiErr) + require.Equal(t, funding.source, info.BillingSource) + f.sessions = append(f.sessions, session) + } + f.rollback.Store(pending) + const callbackName = "test:request-finalize-rollback" + require.NoError(t, f.db.Callback().Update().Before("gorm:update").Register(callbackName, func(tx *gorm.DB) { + updates, ok := tx.Statement.Dest.(map[string]interface{}) + if f.rollback.Load() && tx.Statement.Table == "billing_settlements" && ok && updates["status"] == model.BillingSettlementStatusApplied { + tx.AddError(errors.New("synthetic request finalize transaction failure")) + } + })) + t.Cleanup(func() { _ = f.db.Callback().Update().Remove(callbackName) }) + var apiErr *types.MaxAPIError + if f.info.RelayFormat == types.RelayFormatOpenAI { + apiErr = TextHelper(f.c, f.info) + } else { + apiErr = ResponsesHelper(f.c, f.info) + } + require.NotNil(t, apiErr) + require.True(t, types.IsSkipRetryError(apiErr)) + service.HandleFailedBilling(f.c, f.info, apiErr) + service.HandleFailedBilling(f.c, f.info, apiErr) + require.EqualValues(t, 1, calls.Load()) + require.NoError(t, f.db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(f.requestID)).First(&f.record).Error) + require.EqualValues(t, f.quota-100, f.record.FundingDelta) + require.EqualValues(t, f.quota-100, f.record.TokenDelta) + require.NotEmpty(t, f.record.EffectPayload) + require.NoError(t, common.UnmarshalJsonStr(f.record.EffectPayload, &f.effect)) + require.True(t, f.effect.QuotaIsActual) + require.EqualValues(t, f.quota, f.effect.Quota) + if pending { + require.Equal(t, model.BillingSettlementStatusPending, f.record.Status) + require.Zero(t, f.record.AppliedFundingDelta) + require.Zero(t, f.record.AppliedTokenDelta) + assertRequestReplayReservation(t, f) + } else { + require.Equal(t, model.BillingSettlementStatusApplied, f.record.Status) + assertStreamFundingBalances(t, f.db, funding, f.requestID, f.quota, true, f.prompt, f.complete) + } + return f +} + +func assertRequestReplayReservation(t *testing.T, f *requestReplayFixture) { + t.Helper() + var user model.User + var token model.Token + var sub model.UserSubscription + var channel model.Channel + require.NoError(t, f.db.First(&user, 951).Error) + require.NoError(t, f.db.First(&token, 952).Error) + require.NoError(t, f.db.First(&sub, 955).Error) + require.NoError(t, f.db.First(&channel, 953).Error) + wallet, subUsed := f.funding.wallet, f.funding.subscriptionUsed + if f.funding.source == "wallet" { + wallet -= 100 + } else { + subUsed += 100 + } + require.EqualValues(t, wallet, user.Quota) + require.EqualValues(t, subUsed, sub.AmountUsed) + require.EqualValues(t, 900, token.RemainQuota) + require.EqualValues(t, 100, token.UsedQuota) + require.Zero(t, user.UsedQuota) + require.Zero(t, user.RequestCount) + require.Zero(t, channel.UsedQuota) + var logs, receipts int64 + require.NoError(t, f.db.Model(&model.Log{}).Where("type = ?", model.LogTypeConsume).Count(&logs).Error) + require.NoError(t, f.db.Model(&model.BillingLogReceipt{}).Count(&receipts).Error) + require.Zero(t, logs) + require.Zero(t, receipts) +} + +func requestReplayInput(record model.BillingSettlement, effect *model.BillingSettlementEffect) model.BillingSettlementInput { + return model.BillingSettlementInput{ + OperationKey: record.OperationKey, Source: record.Source, UserID: record.UserID, + SubscriptionID: record.SubscriptionID, TokenID: record.TokenID, + FundingDelta: record.FundingDelta, TokenDelta: record.TokenDelta, + TaskID: record.TaskID, TaskQuota: record.TaskQuota, TaskQuotaTarget: record.TaskQuotaTarget, + SubscriptionPreConsumeRequestID: record.SubscriptionPreConsumeRequestID, + FinalizeSubscriptionPreConsume: record.FinalizeSubscriptionPreConsume, + AllowMissingToken: record.AllowMissingToken, ManualOnFailure: record.ManualOnFailure, + PreConsumeRequestID: record.PreConsumeRequestID, PreConsumeModelName: record.PreConsumeModelName, + PreConsumeRequestedQuota: record.PreConsumeRequestedQuota, PreConsumeEffectiveQuota: record.PreConsumeEffectiveQuota, + Effect: effect, + } +} + +// Include durable records and all projections, not just the primary balance. +func requestReplaySnapshot(t *testing.T, db *gorm.DB) string { + t.Helper() + var snapshot struct { + Users []model.User + Tokens []model.Token + Channels []model.Channel + Subs []model.UserSubscription + Reserved []model.SubscriptionPreConsumeRecord + Selections []model.BillingPreConsumeSelection + Settlements []model.BillingSettlement + Logs []model.Log + Receipts []model.BillingLogReceipt + Outbox []model.CacheInvalidationTask + } + for _, target := range []any{&snapshot.Users, &snapshot.Tokens, &snapshot.Channels, &snapshot.Subs, &snapshot.Reserved, + &snapshot.Selections, &snapshot.Settlements, &snapshot.Logs, &snapshot.Receipts, &snapshot.Outbox} { + require.NoError(t, db.Order("id").Find(target).Error) + } + encoded, err := common.Marshal(snapshot) + require.NoError(t, err) + return fmt.Sprintf("%x", sha256.Sum256(encoded)) +} + +func requestReplayMutations() []struct { + name string + mutate func(*model.BillingSettlementInput) +} { + return []struct { + name string + mutate func(*model.BillingSettlementInput) + }{ + {"user", func(input *model.BillingSettlementInput) { input.UserID = 958 }}, + {"token", func(input *model.BillingSettlementInput) { input.TokenID = 959 }}, + {"source", func(input *model.BillingSettlementInput) { + input.Source = model.BillingSettlementSourceWallet + if input.Source == "wallet" && input.SubscriptionID == 0 { + input.Source = model.BillingSettlementSourceSubscription + } + }}, + {"subscription", func(input *model.BillingSettlementInput) { input.SubscriptionID = 956 }}, + {"funding_delta", func(input *model.BillingSettlementInput) { input.FundingDelta++ }}, + {"token_delta", func(input *model.BillingSettlementInput) { input.TokenDelta++ }}, + {"reservation_request", func(input *model.BillingSettlementInput) { + input.SubscriptionPreConsumeRequestID = "another-reservation" + }}, + {"finalize_reservation", func(input *model.BillingSettlementInput) { + input.FinalizeSubscriptionPreConsume = !input.FinalizeSubscriptionPreConsume + }}, + {"missing_token_policy", func(input *model.BillingSettlementInput) { input.AllowMissingToken = !input.AllowMissingToken }}, + {"manual_policy", func(input *model.BillingSettlementInput) { input.ManualOnFailure = !input.ManualOnFailure }}, + {"task_lifecycle", func(input *model.BillingSettlementInput) { input.TaskID = 987 }}, + {"task_quota", func(input *model.BillingSettlementInput) { input.TaskQuota++ }}, + {"task_target", func(input *model.BillingSettlementInput) { input.TaskQuotaTarget++ }}, + } +} + +func TestRequestFinalizeIdentityRejectsConflictingReplay(t *testing.T) { + initRequestReplayTests(t) + for _, route := range requestReplayRoutes { + for _, funding := range streamFundingCases() { + for _, usage := range []string{"known", "at_reservation", "above_reservation", "zero"} { + for _, pending := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/%s/%s/pending_%t", route, funding.name, usage, pending), func(t *testing.T) { + f := newRequestReplayFixture(t, route, funding, usage, pending, 0) + require.NoError(t, f.db.Create(&model.User{Id: 958, Username: "request-replay-decoy", AffCode: "replay-decoy", Quota: 1000}).Error) + require.NoError(t, f.db.Create(&model.Token{Id: 959, UserId: 958, Key: "synthetic-replay-decoy", RemainQuota: 1000}).Error) + require.NoError(t, f.db.Create(&model.UserSubscription{Id: 956, UserId: 951, PlanId: 954, Status: "active", AmountTotal: 1000, + StartTime: time.Now().Unix() - 1800, EndTime: time.Now().Unix() + 1800}).Error) + before := requestReplaySnapshot(t, f.db) + for _, mutation := range requestReplayMutations() { + t.Run(mutation.name, func(t *testing.T) { + input := requestReplayInput(f.record, &f.effect) + mutation.mutate(&input) + applied, already, err := model.ApplyBillingSettlementOnce(input) + require.ErrorIs(t, err, model.ErrBillingSettlementOperationConflict) + require.Zero(t, applied) + require.False(t, already) + require.Equal(t, before, requestReplaySnapshot(t, f.db), "conflicting replay must change no ledger or projection") + }) + } + f.rollback.Store(false) + for _, control := range []string{"effect_omitted", "description_drift"} { + t.Run(control, func(t *testing.T) { + input := requestReplayInput(f.record, nil) + if control == "description_drift" { + effect := f.effect + effect.Content, effect.ModelName = "different harmless description", "description-only-model" + input.Effect = &effect + } + for range 2 { + applied, _, err := model.ApplyBillingSettlementOnce(input) + require.NoError(t, err) + require.EqualValues(t, f.quota-100, applied) + require.NoError(t, model.ProcessBillingSettlementEffect(input.OperationKey)) + } + var record model.BillingSettlement + require.NoError(t, f.db.Where("operation_key = ?", input.OperationKey).First(&record).Error) + require.Equal(t, f.record.EffectPayload, record.EffectPayload) + require.Equal(t, model.BillingSettlementStatusApplied, record.Status) + assertStreamFundingBalances(t, f.db, funding, f.requestID, f.quota, true, f.prompt, f.complete) + }) + } + }) + } + } + } + } +} + +func TestRequestAdmissionRejectsIdentityAndReservationDrift(t *testing.T) { + initRequestReplayTests(t) + for _, funding := range streamFundingCases() { + t.Run(funding.name, func(t *testing.T) { + for _, mutation := range []struct { + name string + mutate func(*relaycommon.RelayInfo, *int) + }{ + {"user", func(info *relaycommon.RelayInfo, _ *int) { info.UserId = 958 }}, + {"token", func(info *relaycommon.RelayInfo, _ *int) { + info.TokenId, info.TokenKey = 959, "synthetic-admission-decoy" + }}, + {"model", func(info *relaycommon.RelayInfo, _ *int) { info.OriginModelName = "another-model" }}, + {"playground", func(info *relaycommon.RelayInfo, _ *int) { info.IsPlayground = true }}, + {"reservation_zero", func(_ *relaycommon.RelayInfo, quota *int) { *quota = 0 }}, + {"reservation_lower", func(_ *relaycommon.RelayInfo, quota *int) { *quota = 99 }}, + {"reservation_higher", func(_ *relaycommon.RelayInfo, quota *int) { *quota = 101 }}, + } { + t.Run(mutation.name, func(t *testing.T) { + f := newRequestReplayFixture(t, "responses_native", funding, "known", false, 0) + require.NoError(t, f.db.Create(&model.User{Id: 958, Username: "admission-replay-decoy", AffCode: "admission-decoy", Quota: 1000}).Error) + require.NoError(t, f.db.Create(&model.Token{Id: 959, UserId: 958, Key: "synthetic-admission-decoy", RemainQuota: 1000}).Error) + before := requestReplaySnapshot(t, f.db) + var beforeUser, afterUser model.User + require.NoError(t, f.db.First(&beforeUser, 951).Error) + info := &relaycommon.RelayInfo{RequestId: f.requestID, UserId: 951, TokenId: 952, TokenKey: f.info.TokenKey, + OriginModelName: f.info.OriginModelName, ForcePreConsume: true, UserSetting: dto.UserSetting{BillingPreference: funding.preference}} + quota := 100 + mutation.mutate(info, "a) + session, apiErr := service.NewBillingSession(f.c, info, quota) + require.NoError(t, f.db.First(&afterUser, 951).Error) + t.Logf("wallet quota before/after admission replay: %d/%d", beforeUser.Quota, afterUser.Quota) + require.NotNil(t, apiErr) + require.Nil(t, session) + require.Equal(t, before, requestReplaySnapshot(t, f.db), "conflicting admission must not reserve or project again") + }) + } + }) + } +} + +func TestRequestPlaygroundFreshIdentityKeepsFundingRules(t *testing.T) { + initRequestReplayTests(t) + for _, funding := range streamFundingCases() { + t.Run(funding.name, func(t *testing.T) { + db := setupStreamFundingLedger(t, funding) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + info := &relaycommon.RelayInfo{RequestId: "fresh-playground", UserId: 951, IsPlayground: true, ForcePreConsume: true, + OriginModelName: "gpt-test", UserSetting: dto.UserSetting{BillingPreference: funding.preference}} + session, apiErr := service.NewBillingSession(c, info, 100) + require.Nil(t, apiErr) + require.Equal(t, funding.source, info.BillingSource) + require.NoError(t, session.Settle(6)) + session.Refund(c) + var user model.User + var token model.Token + var sub model.UserSubscription + require.NoError(t, db.First(&user, 951).Error) + require.NoError(t, db.First(&token, 952).Error) + require.NoError(t, db.First(&sub, 955).Error) + wallet, subUsed := funding.wallet, funding.subscriptionUsed + if funding.source == "wallet" { + wallet -= 6 + } else { + subUsed += 6 + } + require.EqualValues(t, wallet, user.Quota) + require.EqualValues(t, subUsed, sub.AmountUsed) + require.EqualValues(t, 1000, token.RemainQuota) + require.Zero(t, token.UsedQuota) + var record model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(info.RequestId)).First(&record).Error) + require.Equal(t, model.BillingSettlementStatusApplied, record.Status) + require.EqualValues(t, -94, record.AppliedFundingDelta) + require.Zero(t, record.AppliedTokenDelta) + }) + } +} + +func TestRequestFinalizeConcurrentSettlementRefundAndRecovery(t *testing.T) { + initRequestReplayTests(t) + for _, route := range requestReplayRoutes { + for _, funding := range streamFundingCases() { + for _, usage := range []string{"known", "at_reservation", "above_reservation", "zero"} { + t.Run(strings.Join([]string{route, funding.name, usage}, "/"), func(t *testing.T) { + f := newRequestReplayFixture(t, route, funding, usage, true, 6) + f.rollback.Store(false) + require.NoError(t, f.db.Model(&model.BillingSettlement{}).Where("id = ?", f.record.ID).Update("next_attempt", 0).Error) + start := make(chan struct{}) + results := make(chan error, 7) + var workers sync.WaitGroup + launch := func(work func()) { + workers.Add(1) + go func() { defer workers.Done(); <-start; work() }() + } + for _, session := range f.sessions[:3] { + launch(func() { results <- session.SettleWithEffect(f.quota, &f.effect) }) + } + for _, session := range f.sessions[3:] { + launch(func() { session.Refund(f.c) }) + } + for range 4 { + launch(func() { + applied, _, err := model.ApplyBillingSettlementOnce(requestReplayInput(f.record, &f.effect)) + if err == nil && applied != int64(f.quota-100) { + err = fmt.Errorf("recorded delta differs: got %d", applied) + } + results <- err + }) + } + for range 3 { + launch(model.ProcessPendingBillingSettlementsOnce) + } + close(start) + workers.Wait() + close(results) + for err := range results { + require.NoError(t, err) + } + model.ProcessPendingBillingSettlementsOnce() + model.ProcessPendingBillingSettlementsOnce() + var record model.BillingSettlement + require.NoError(t, f.db.Where("id = ?", f.record.ID).First(&record).Error) + require.Equal(t, model.BillingSettlementStatusApplied, record.Status) + require.Equal(t, model.BillingSettlementEffectApplied, record.EffectStatus) + require.EqualValues(t, f.quota-100, record.AppliedFundingDelta) + require.EqualValues(t, f.quota-100, record.AppliedTokenDelta) + require.Equal(t, f.record.EffectPayload, record.EffectPayload) + assertStreamFundingBalances(t, f.db, funding, f.requestID, f.quota, true, f.prompt, f.complete) + var receipts int64 + require.NoError(t, f.db.Model(&model.BillingLogReceipt{}).Where("operation_key = ?", record.OperationKey).Count(&receipts).Error) + require.EqualValues(t, 1, receipts) + }) + } + } + } +} diff --git a/relay/responses_custom_tool_eof_settlement_test.go b/relay/responses_custom_tool_eof_settlement_test.go new file mode 100644 index 00000000..b6ee4410 --- /dev/null +++ b/relay/responses_custom_tool_eof_settlement_test.go @@ -0,0 +1,308 @@ +package relay + +import ( + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/constant" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/model" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + relayconstant "github.com/MAX-API-Next/MAX-API/relay/constant" + "github.com/MAX-API-Next/MAX-API/service" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" + "gorm.io/gorm" + "gorm.io/gorm/logger" +) + +type customToolEOFWriter struct { + *httptest.ResponseRecorder + ending string + secondInput string + flushFailure bool +} + +func (w *customToolEOFWriter) Write(p []byte) (int, error) { + data := strings.SplitN(string(p), "data: ", 2) + if len(data) != 2 { + return w.ResponseRecorder.Write(p) + } + eventType := gjson.Get(data[1], "type").String() + inputDelta := eventType == "response.custom_tool_call_input.delta" + fail := w.ending == "first_input_write_failure" && inputDelta || + (w.ending == "second_input_write_failure" || w.ending == "finish_second_input_write_failure") && inputDelta && gjson.Get(data[1], "delta").String() == w.secondInput || + w.ending == "terminal_write_failure" && (eventType == "response.completed" || eventType == "response.incomplete") || + w.ending == "no_delivery" && eventType == "response.output_item.added" + if fail { + return 0, io.ErrClosedPipe + } + if w.ending == "input_flush_failure" && inputDelta { + w.flushFailure = true + } + return w.ResponseRecorder.Write(p) +} + +func (w *customToolEOFWriter) FlushError() error { + if w.flushFailure { + return io.ErrClosedPipe + } + w.ResponseRecorder.Flush() + return nil +} + +// Exercise EOF finalization at the real ResponsesHelper and funding/effect seam. +// A late custom input must be billed once, and a failed input write must not be billed. +func TestResponsesCustomToolEOFFinalizationUsesDurableLedger(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitHttpClient() + inputs := []string{"first patch line\nsecond line", strings.Repeat("second patch body ", 50)} + for _, provider := range []string{"chat", "claude", "gemini"} { + for _, funding := range []string{"wallet_only", "subscription_only"} { + for _, usageKind := range []string{"missing", "known", "zero", "input_only", "zero_input_only"} { + if provider == "chat" && strings.HasSuffix(usageKind, "input_only") { + continue + } + for _, ending := range []string{"eof", "first_input_write_failure", "second_input_write_failure", "finish_second_input_write_failure", "input_flush_failure", "terminal_write_failure", "no_delivery"} { + t.Run(strings.Join([]string{provider, funding, usageKind, ending}, "/"), func(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}, &model.Channel{}, &model.Log{}, &model.BillingLogReceipt{}, + &model.BillingSettlement{}, &model.BillingPreConsumeSelection{}, &model.CacheInvalidationTask{}, + &model.UserSubscription{}, &model.SubscriptionPreConsumeRecord{}, &model.SubscriptionPlan{})) + oldDB, oldLogDB := model.DB, model.LOG_DB + oldSQLite, oldRedis, oldBatch, oldConsume := common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled + model.DB, model.LOG_DB = db, db + common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled = true, false, false, true + t.Cleanup(func() { + model.DB, model.LOG_DB = oldDB, oldLogDB + common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled = oldSQLite, oldRedis, oldBatch, oldConsume + _ = sqlDB.Close() + }) + require.NoError(t, db.Create(&model.User{Id: 901, Username: "custom-eof-ledger", Quota: 1000}).Error) + require.NoError(t, db.Create(&model.Token{Id: 902, UserId: 901, Key: "synthetic-custom-eof", RemainQuota: 1000, Status: common.TokenStatusEnabled}).Error) + require.NoError(t, db.Create(&model.Channel{Id: 903}).Error) + if funding == "subscription_only" { + require.NoError(t, db.Create(&model.SubscriptionPlan{Id: 904, Title: "custom-eof-plan", Enabled: true, TotalAmount: 1000, QuotaResetPeriod: model.SubscriptionResetNever}).Error) + require.NoError(t, db.Create(&model.UserSubscription{Id: 905, UserId: 901, PlanId: 904, AmountTotal: 1000, Status: "active", + StartTime: time.Now().Add(-time.Hour).Unix(), EndTime: time.Now().Add(time.Hour).Unix()}).Error) + } + prompt, completion := 10, 0 + deliveredInput := strings.Join(inputs, "") + if ending == "first_input_write_failure" || ending == "no_delivery" { + deliveredInput = "" + } else if ending == "second_input_write_failure" || ending == "finish_second_input_write_failure" || ending == "input_flush_failure" { + deliveredInput = inputs[0] + } + completion = service.EstimateTokenByModel("gpt-test", deliveredInput) + if usageKind == "known" { + prompt, completion = 4, 2 + } else if usageKind == "zero" { + prompt, completion = 0, 0 + } else if usageKind == "input_only" { + prompt = 4 + } else if usageKind == "zero_input_only" { + prompt = 0 + } + channelType, frames := customToolEOFFrames(t, provider, usageKind, inputs) + if ending == "finish_second_input_write_failure" && usageKind != "known" && usageKind != "zero" { + switch provider { + case "chat": + frames = append(frames, `{"choices":[{"delta":{},"finish_reason":"tool_calls"}]}`) + case "claude": + frames = append(frames, `{"type":"message_delta","delta":{"stop_reason":"tool_use"}}`) + case "gemini": + frames = append(frames, `{"candidates":[{"content":{"role":"model","parts":[]},"finishReason":"STOP"}]}`) + } + } + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, "data: "+strings.Join(frames, "\n\ndata: ")+"\n\n") + })) + t.Cleanup(upstream.Close) + writer := &customToolEOFWriter{ResponseRecorder: httptest.NewRecorder(), ending: ending, secondInput: inputs[1]} + c, _ := gin.CreateTestContext(writer) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + c.Set(string(constant.ContextKeyChannelType), channelType) + c.Set(string(constant.ContextKeyChannelBaseUrl), upstream.URL) + c.Set(string(constant.ContextKeyChannelKey), "synthetic-test-key") + c.Set(string(constant.ContextKeyChannelId), 903) + c.Set(string(constant.ContextKeyOriginalModel), "gpt-test") + if provider == "chat" { + c.Set(string(constant.ContextKeyChannelOtherSetting), dto.ChannelOtherSettings{AdvancedCustom: &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{ + {IncomingPath: "/v1/responses", UpstreamPath: "/v1/chat/completions", Converter: dto.AdvancedCustomConverterOpenAIResponsesToOpenAIChatCompletions}, + }}}) + } + requestID := "custom-eof-" + strings.Join([]string{provider, funding, usageKind, ending}, "-") + c.Set(common.RequestIdKey, requestID) + info := &relaycommon.RelayInfo{RequestId: requestID, UserId: 901, TokenId: 902, TokenKey: "synthetic-custom-eof", ForcePreConsume: true, + RelayMode: relayconstant.RelayModeResponses, RelayFormat: types.RelayFormatOpenAIResponses, IsStream: true, DisablePing: true, + OriginModelName: "gpt-test", StartTime: time.Now(), UsingGroup: "default", UserSetting: dto.UserSetting{BillingPreference: funding, QuotaWarningThreshold: 1}, + PriceData: types.PriceData{ModelRatio: 1, CompletionRatio: 1, GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 1}}, + Request: &dto.OpenAIResponsesRequest{Model: "gpt-test", Input: []byte(`"test"`), Stream: common.GetPointer(true), + Tools: []byte(`[{"type":"custom","name":"apply_patch"}]`)}} + info.SetEstimatePromptTokens(10) + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + apiErr := ResponsesHelper(c, info) + if ending == "eof" && (usageKind == "known" || usageKind == "zero") { + require.Nil(t, apiErr) + } else { + require.NotNil(t, apiErr) + require.True(t, types.IsSkipRetryError(apiErr)) + service.HandleFailedBilling(c, info, apiErr) + service.HandleFailedBilling(c, info, apiErr) + } + require.EqualValues(t, 1, calls.Load()) + quota := prompt + completion + if ending == "no_delivery" { + quota = 0 + } + var settlement model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).First(&settlement).Error) + require.Equal(t, model.BillingSettlementStatusApplied, settlement.Status) + require.EqualValues(t, quota-100, settlement.FundingDelta, "late delivered custom input must be reflected in actual funding") + require.EqualValues(t, quota-100, settlement.TokenDelta) + if ending != "no_delivery" { + var effect model.BillingSettlementEffect + require.NoError(t, common.UnmarshalJsonStr(settlement.EffectPayload, &effect)) + require.Equal(t, model.BillingSettlementEffectApplied, settlement.EffectStatus) + require.Equal(t, prompt, effect.PromptTokens) + require.Equal(t, completion, effect.CompletionTokens) + require.EqualValues(t, quota, effect.Quota) + require.True(t, effect.QuotaIsActual) + replayInfo := &relaycommon.RelayInfo{RequestId: requestID, UserId: info.UserId, TokenId: info.TokenId, TokenKey: info.TokenKey, + OriginModelName: info.OriginModelName, UserSetting: info.UserSetting, ForcePreConsume: true} + replay, replayErr := service.NewBillingSession(c, replayInfo, 100) + require.Nil(t, replayErr) + require.NoError(t, replay.SettleWithEffect(quota, &effect)) + require.NoError(t, replay.SettleWithEffect(quota, &effect)) + replay.Refund(c) + } + var user model.User + var token model.Token + var channel model.Channel + require.NoError(t, db.First(&user, 901).Error) + require.NoError(t, db.First(&token, 902).Error) + require.NoError(t, db.First(&channel, 903).Error) + walletQuota := 1000 - quota + if funding == "subscription_only" { + walletQuota = 1000 + var sub model.UserSubscription + require.NoError(t, db.First(&sub, 905).Error) + require.EqualValues(t, quota, sub.AmountUsed) + } + require.EqualValues(t, walletQuota, user.Quota) + require.EqualValues(t, 1000-quota, token.RemainQuota) + require.EqualValues(t, quota, user.UsedQuota) + require.EqualValues(t, quota, token.UsedQuota) + require.EqualValues(t, quota, channel.UsedQuota) + var count int64 + require.NoError(t, db.Model(&model.Log{}).Count(&count).Error) + wantLogs := 1 + if ending == "no_delivery" { + wantLogs = 0 + } + require.EqualValues(t, wantLogs, count) + wantRequests := wantLogs + if quota == 0 { + wantRequests = 0 + } + require.EqualValues(t, wantRequests, user.RequestCount) + terminalCount := 0 + var gotInput strings.Builder + for _, line := range strings.Split(writer.Body.String(), "\n") { + if !strings.HasPrefix(line, "data: ") { + continue + } + data := strings.TrimPrefix(line, "data: ") + if gjson.Get(data, "type").String() == "response.custom_tool_call_input.delta" { + gotInput.WriteString(gjson.Get(data, "delta").String()) + } + if gjson.Get(data, "type").String() == "response.completed" || gjson.Get(data, "type").String() == "response.incomplete" { + terminalCount++ + require.EqualValues(t, completion, gjson.Get(data, "response.usage.output_tokens").Int()) + require.EqualValues(t, prompt, gjson.Get(data, "response.usage.input_tokens").Int()) + } + } + require.Equal(t, deliveredInput, gotInput.String(), "only fully written input deltas may participate in settlement") + if ending == "eof" { + require.Equal(t, 1, terminalCount) + } else { + require.Zero(t, terminalCount, "no terminal may be appended after the downstream fails") + } + }) + } + } + } + } +} + +func customToolEOFFrames(t *testing.T, provider, usageKind string, inputs []string) (int, []string) { + t.Helper() + marshal := func(v any) string { b, err := common.Marshal(v); require.NoError(t, err); return string(b) } + prompt, completion := 4, 2 + if strings.HasPrefix(usageKind, "zero") { + prompt, completion = 0, 0 + } + complete := usageKind == "known" || usageKind == "zero" + var frames []string + switch provider { + case "chat": + if usageKind != "missing" { + frames = append(frames, marshal(map[string]any{"choices": []any{}, "usage": map[string]int{"prompt_tokens": prompt, "completion_tokens": completion, "total_tokens": prompt + completion}})) + } + for idx, input := range inputs { + frames = append(frames, marshal(map[string]any{"choices": []any{map[string]any{"delta": map[string]any{"tool_calls": []any{map[string]any{"index": idx, "id": fmt.Sprintf("call_%d", idx), "type": "function", "function": map[string]any{"name": "apply_patch", "arguments": marshal(map[string]string{"input": input})}}}}}}})) + } + if usageKind != "missing" { + frames = append(frames, `{"choices":[{"delta":{},"finish_reason":"tool_calls"}]}`) + } + return constant.ChannelTypeAdvancedCustom, frames + case "claude": + message := map[string]any{"id": "msg_custom_eof", "model": "gpt-test"} + if usageKind != "missing" { + message["usage"] = map[string]int{"input_tokens": prompt} + } + frames = append(frames, marshal(map[string]any{"type": "message_start", "message": message})) + for idx, input := range inputs { + frames = append(frames, marshal(map[string]any{"type": "content_block_start", "index": idx, "content_block": map[string]any{"type": "tool_use", "id": fmt.Sprintf("call_%d", idx), "name": "apply_patch", "input": map[string]any{}}}), + marshal(map[string]any{"type": "content_block_delta", "index": idx, "delta": map[string]any{"type": "input_json_delta", "partial_json": marshal(map[string]string{"input": input})}})) + } + if complete { + frames = append(frames, marshal(map[string]any{"type": "message_delta", "delta": map[string]string{"stop_reason": "tool_use"}, "usage": map[string]int{"output_tokens": completion}})) + } + return constant.ChannelTypeAnthropic, frames + default: + if usageKind != "missing" { + metadata := map[string]int{"promptTokenCount": prompt} + if complete { + metadata["candidatesTokenCount"], metadata["totalTokenCount"] = completion, prompt+completion + } + frames = append(frames, marshal(map[string]any{"usageMetadata": metadata})) + } + var parts []any + for idx, input := range inputs { + parts = append(parts, map[string]any{"functionCall": map[string]any{"id": fmt.Sprintf("call_%d", idx), "name": "apply_patch", "args": map[string]string{"input": input}}}) + } + frames = append(frames, marshal(map[string]any{"candidates": []any{map[string]any{"content": map[string]any{"role": "model", "parts": parts}}}})) + if complete { + frames = append(frames, `{"candidates":[{"content":{"role":"model","parts":[]},"finishReason":"STOP"}]}`) + } + return constant.ChannelTypeGemini, frames + } +} diff --git a/relay/responses_failure_boundary_test.go b/relay/responses_failure_boundary_test.go new file mode 100644 index 00000000..7549c55e --- /dev/null +++ b/relay/responses_failure_boundary_test.go @@ -0,0 +1,168 @@ +package relay + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/relay/channel/claude" + "github.com/MAX-API-Next/MAX-API/relay/channel/gemini" + "github.com/MAX-API-Next/MAX-API/relay/channel/openai" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/service" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/sjson" +) + +type responsesSelectiveFailureWriter struct { + *httptest.ResponseRecorder + failOn string +} + +func (w *responsesSelectiveFailureWriter) Write(p []byte) (int, error) { + if strings.Contains(string(p), w.failOn) { + return 0, io.ErrClosedPipe + } + return w.ResponseRecorder.Write(p) +} + +func TestResponsesStreamFailureRetainsOnlyDeliveredUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitTokenEncoders() + for _, provider := range []struct { + name string + frames []string + handle func(*gin.Context, *relaycommon.RelayInfo, *http.Response) (*dto.Usage, *types.MaxAPIError) + }{ + {name: "chat_bridge", frames: []string{ + `{"choices":[{"delta":{"content":"partial output"}}],"usage":{"prompt_tokens":10,"completion_tokens":2,"total_tokens":12}}`, + `{"choices":[{"delta":{"content":"undelivered"}}],"usage":{"prompt_tokens":1000,"completion_tokens":999,"total_tokens":1999}}`, + }, handle: openai.OaiChatToResponsesStreamHandler}, + {name: "gemini", frames: []string{ + `{"candidates":[{"content":{"parts":[{"text":"partial output"}]}}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":2,"totalTokenCount":12}}`, + `{"candidates":[{"content":{"parts":[{"text":"undelivered"}]}}],"usageMetadata":{"promptTokenCount":1000,"candidatesTokenCount":999,"totalTokenCount":1999}}`, + }, handle: gemini.GeminiResponsesStreamHandler}, + {name: "claude", frames: []string{ + `{"type":"message_start","message":{"model":"gpt-test","usage":{"input_tokens":10,"output_tokens":2}}}`, + `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial output"}}`, + `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"undelivered"}}`, + }, handle: func(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.MaxAPIError) { + return claude.ClaudeStreamHandler(c, resp, info) + }}, + {name: "responses_native", frames: []string{ + `{"type":"response.output_text.delta","delta":"partial output"}`, + `{"type":"response.output_text.delta","delta":"undelivered"}`, + }, handle: openai.OaiResponsesStreamHandler}, + } { + for _, failure := range []string{"first_write", "later_write", "malformed", "done_without_terminal"} { + t.Run(provider.name+"/"+failure, func(t *testing.T) { + recorder := httptest.NewRecorder() + var writer http.ResponseWriter = recorder + frames := append([]string(nil), provider.frames...) + if failure == "first_write" { + writer = &responsesSelectiveFailureWriter{recorder, "event:"} + } + if failure == "later_write" { + writer = &responsesSelectiveFailureWriter{recorder, "undelivered"} + } + if failure == "malformed" { + frames = append(frames[:len(frames)-1], `{broken`) + } + if failure == "done_without_terminal" { + frames = append(frames[:len(frames)-1], `[DONE]`) + } + c, _ := gin.CreateTestContext(writer) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + info := &relaycommon.RelayInfo{IsStream: true, RelayFormat: types.RelayFormatOpenAIResponses, DisablePing: true, + OriginModelName: "gpt-test", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-test"}} + info.SetEstimatePromptTokens(10) + usage, apiErr := provider.handle(c, info, &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader("data: " + strings.Join(frames, "\n\ndata: ") + "\n\n"))}) + require.NotNil(t, apiErr) + require.True(t, types.IsSkipRetryError(apiErr)) + if failure == "first_write" { + require.Nil(t, usage) + return + } + require.NotNil(t, usage) + require.Equal(t, 10, usage.PromptTokens) + require.Positive(t, usage.CompletionTokens) + require.Less(t, usage.CompletionTokens, 10, "failed chunk must not replace delivered usage") + require.NotContains(t, recorder.Body.String(), "undelivered") + require.NotContains(t, recorder.Body.String(), "event: response.completed") + }) + } + } +} + +func TestResponsesProviderExplicitZeroSurvivesFallback(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitTokenEncoders() + for _, provider := range []struct { + name string + frames []string + handle func(*gin.Context, *relaycommon.RelayInfo, *http.Response) (*dto.Usage, *types.MaxAPIError) + }{ + {name: "chat_bridge", frames: []string{`{"choices":[{"delta":{"content":"visible"}}]}`, `{"choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":0,"completion_tokens":0,"total_tokens":0}}`}, handle: openai.OaiChatToResponsesStreamHandler}, + {name: "gemini", frames: []string{`{"candidates":[{"content":{"parts":[{"text":"visible"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":0,"candidatesTokenCount":0,"totalTokenCount":0}}`}, handle: gemini.GeminiResponsesStreamHandler}, + {name: "claude", frames: []string{`{"type":"message_start","message":{"model":"gpt-test","usage":{"input_tokens":0}}}`, `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"visible"}}`, `{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":0}}`}, handle: func(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.MaxAPIError) { + return claude.ClaudeStreamHandler(c, resp, info) + }}, + {name: "responses_native", frames: []string{`{"type":"response.output_text.delta","delta":"visible"}`, `{"type":"response.completed","response":{"status":"completed","usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`}, handle: openai.OaiResponsesStreamHandler}, + } { + for _, emptyUsage := range []bool{false, true} { + t.Run(provider.name+map[bool]string{false: "/explicit_zero", true: "/empty_usage"}[emptyUsage], func(t *testing.T) { + frames := append([]string(nil), provider.frames...) + if emptyUsage { + for i, frame := range frames { + var path string + switch provider.name { + case "chat_bridge": + if strings.Contains(frame, `"usage"`) { + path = "usage" + } + case "gemini": + path = "usageMetadata" + case "claude": + if strings.Contains(frame, `"message_start"`) { + path = "message.usage" + } else if strings.Contains(frame, `"message_delta"`) { + path = "usage" + } + case "responses_native": + if strings.Contains(frame, `"response.completed"`) { + path = "response.usage" + } + } + if path != "" { + var err error + frames[i], err = sjson.SetRaw(frame, path, `{}`) + require.NoError(t, err) + } + } + } + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + info := &relaycommon.RelayInfo{IsStream: true, RelayFormat: types.RelayFormatOpenAIResponses, DisablePing: true, OriginModelName: "gpt-test", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-test"}} + info.SetEstimatePromptTokens(10) + usage, apiErr := provider.handle(c, info, &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader("data: " + strings.Join(frames, "\n\ndata: ") + "\n\n"))}) + require.Nil(t, apiErr) + require.NotNil(t, usage) + if emptyUsage { + require.Equal(t, 10, usage.PromptTokens) + require.Positive(t, usage.CompletionTokens) + } else { + require.Zero(t, usage.PromptTokens) + require.Zero(t, usage.CompletionTokens) + require.Zero(t, usage.TotalTokens) + require.NotNil(t, usage.BillingUsage) + require.True(t, usage.BillingUsage.TokenCountsReported) + } + }) + } + } +} diff --git a/relay/responses_first_result_settlement_test.go b/relay/responses_first_result_settlement_test.go new file mode 100644 index 00000000..393ce508 --- /dev/null +++ b/relay/responses_first_result_settlement_test.go @@ -0,0 +1,232 @@ +package relay + +import ( + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/constant" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/model" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + relayconstant "github.com/MAX-API-Next/MAX-API/relay/constant" + "github.com/MAX-API-Next/MAX-API/service" + "github.com/MAX-API-Next/MAX-API/setting/operation_setting" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "github.com/tidwall/sjson" + "gorm.io/gorm" + "gorm.io/gorm/logger" +) + +// Drive the production adapters, timeout watchdog and durable settlement seam. +// Whitespace is delivered generation data even though it is not a first result. +func TestResponsesFirstResultTimeoutSettlesDeliveredWhitespace(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitHttpClient() + service.InitTokenEncoders() + for _, provider := range []struct { + name string + channelType int + frames []string + textFrame int + textPath string + usageRoots map[int]string + counterPaths map[int][]string + }{ + {"native", constant.ChannelTypeOpenAI, []string{ + `{"type":"response.in_progress","response":{"id":"resp_timeout","status":"in_progress","usage":{"input_tokens":10,"output_tokens":2,"total_tokens":12}}}`, + `{"type":"response.output_text.delta","delta":"\n "}`, + }, 1, "delta", map[int]string{0: "response.usage"}, map[int][]string{0: {"response.usage.input_tokens", "response.usage.output_tokens", "response.usage.total_tokens"}}}, + {"claude", constant.ChannelTypeAnthropic, []string{ + `{"type":"message_start","message":{"id":"msg_timeout","model":"gpt-test","usage":{"input_tokens":10,"output_tokens":2}}}`, + `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"\n "}}`, + `{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":2}}`, + }, 1, "delta.text", map[int]string{0: "message.usage", 2: "usage"}, map[int][]string{0: {"message.usage.input_tokens", "message.usage.output_tokens"}, 2: {"usage.output_tokens"}}}, + {"gemini", constant.ChannelTypeGemini, []string{ + `{"candidates":[{"content":{"parts":[{"text":"\n "}]}}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":2,"totalTokenCount":12}}`, + }, 0, "candidates.0.content.parts.0.text", map[int]string{0: "usageMetadata"}, map[int][]string{0: {"usageMetadata.promptTokenCount", "usageMetadata.candidatesTokenCount", "usageMetadata.totalTokenCount"}}}, + {"chat_bridge", constant.ChannelTypeAdvancedCustom, []string{ + `{"choices":[{"delta":{"content":"\n "}}],"usage":{"prompt_tokens":10,"completion_tokens":2,"total_tokens":12}}`, + }, 0, "choices.0.delta.content", map[int]string{0: "usage"}, map[int][]string{0: {"usage.prompt_tokens", "usage.completion_tokens", "usage.total_tokens"}}}, + } { + for _, scenario := range []string{"paid", "free", "explicit_zero", "missing_usage", "metadata_only"} { + for _, funding := range []string{"wallet_only", "subscription_only"} { + t.Run(provider.name+"/"+scenario+"/"+funding, func(t *testing.T) { + oldMonitor := *operation_setting.GetMonitorSetting() + *operation_setting.GetMonitorSetting() = operation_setting.MonitorSetting{StreamingFirstResultTimeoutSeconds: 1} + t.Cleanup(func() { *operation_setting.GetMonitorSetting() = oldMonitor }) + frames := append([]string(nil), provider.frames...) + prompt, completion := 10, 2 + if scenario == "explicit_zero" { + prompt, completion = 0, 0 + for idx, paths := range provider.counterPaths { + for _, path := range paths { + var err error + frames[idx], err = sjson.Set(frames[idx], path, 0) + require.NoError(t, err) + } + } + } + if scenario == "missing_usage" { + for idx, root := range provider.usageRoots { + var err error + frames[idx], err = sjson.Delete(frames[idx], root) + require.NoError(t, err) + } + // The existing local fallback estimates this two-character + // whitespace string as one token; it does not use tiktoken. + completion = 1 + if provider.name == "native" { + // Native Responses retains its existing tiktoken fallback, + // which counts this whitespace as two tokens. + completion = 2 + } + } + if scenario == "metadata_only" { + var err error + frames[provider.textFrame], err = sjson.Delete(frames[provider.textFrame], provider.textPath) + require.NoError(t, err) + } + finalQuota := prompt + completion + if scenario == "free" || scenario == "metadata_only" { + finalQuota = 0 + } + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}, &model.Channel{}, &model.Log{}, &model.BillingLogReceipt{}, + &model.BillingSettlement{}, &model.BillingPreConsumeSelection{}, &model.CacheInvalidationTask{}, + &model.UserSubscription{}, &model.SubscriptionPreConsumeRecord{}, &model.SubscriptionPlan{})) + oldDB, oldLogDB := model.DB, model.LOG_DB + oldSQLite, oldRedis, oldBatch, oldConsume := common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled + model.DB, model.LOG_DB = db, db + common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled = true, false, false, true + t.Cleanup(func() { + model.DB, model.LOG_DB = oldDB, oldLogDB + common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled = oldSQLite, oldRedis, oldBatch, oldConsume + _ = sqlDB.Close() + }) + require.NoError(t, db.Create(&model.User{Id: 201, Username: "first-result-ledger", Quota: 1000}).Error) + require.NoError(t, db.Create(&model.Token{Id: 202, UserId: 201, Key: "synthetic-timeout-ledger-key", RemainQuota: 1000, Status: common.TokenStatusEnabled}).Error) + require.NoError(t, db.Create(&model.Channel{Id: 203}).Error) + if funding == "subscription_only" { + require.NoError(t, db.Create(&model.SubscriptionPlan{Id: 204, Title: "timeout-ledger-plan", Enabled: true, TotalAmount: 1000, QuotaResetPeriod: model.SubscriptionResetNever}).Error) + require.NoError(t, db.Create(&model.UserSubscription{Id: 205, UserId: 201, PlanId: 204, AmountTotal: 1000, Status: "active", + StartTime: time.Now().Add(-time.Hour).Unix(), EndTime: time.Now().Add(time.Hour).Unix()}).Error) + } + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, "data: "+strings.Join(frames, "\n\ndata: ")+"\n\n") + w.(http.Flusher).Flush() + select { + case <-r.Context().Done(): + case <-time.After(3 * time.Second): + } + })) + t.Cleanup(upstream.Close) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + c.Set(string(constant.ContextKeyChannelType), provider.channelType) + c.Set(string(constant.ContextKeyChannelBaseUrl), upstream.URL) + c.Set(string(constant.ContextKeyChannelKey), "synthetic-test-key") + c.Set(string(constant.ContextKeyChannelId), 203) + c.Set(string(constant.ContextKeyOriginalModel), "gpt-test") + if provider.name == "chat_bridge" { + c.Set(string(constant.ContextKeyChannelOtherSetting), dto.ChannelOtherSettings{AdvancedCustom: &dto.AdvancedCustomConfig{ + Routes: []dto.AdvancedCustomRoute{{IncomingPath: "/v1/responses", UpstreamPath: "/v1/chat/completions", Converter: dto.AdvancedCustomConverterOpenAIResponsesToOpenAIChatCompletions}}, + }}) + } + requestID := "timeout-ledger-" + provider.name + "-" + scenario + "-" + funding + c.Set(common.RequestIdKey, requestID) + info := &relaycommon.RelayInfo{RequestId: requestID, UserId: 201, TokenId: 202, TokenKey: "synthetic-timeout-ledger-key", ForcePreConsume: true, + RelayMode: relayconstant.RelayModeResponses, RelayFormat: types.RelayFormatOpenAIResponses, IsStream: true, DisablePing: true, + OriginModelName: "gpt-test", StartTime: time.Now(), UsingGroup: "default", UserSetting: dto.UserSetting{BillingPreference: funding, QuotaWarningThreshold: 1}, + PriceData: types.PriceData{ModelRatio: 1, CompletionRatio: 1, GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 1}}, + Request: &dto.OpenAIResponsesRequest{Model: "gpt-test", Input: []byte(`"test"`), Stream: common.GetPointer(true)}} + info.SetEstimatePromptTokens(10) + if scenario == "free" { + info.PriceData.ModelRatio = 0 + } + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + apiErr := ResponsesHelper(c, info) + require.NotNil(t, apiErr) + require.Equal(t, types.ErrorCodeChannelResponseTimeExceeded, apiErr.GetErrorCode()) + require.True(t, types.IsSkipRetryError(apiErr)) + require.Equal(t, relaycommon.StreamEndReasonTimeout, info.StreamStatus.EndReason) + require.False(t, info.HasRecordedChannelFirstResult()) + require.True(t, info.FirstResultTime.IsZero()) + require.EqualValues(t, 1, calls.Load()) + require.NotContains(t, recorder.Body.String(), "event: response.completed") + service.HandleFailedBilling(c, info, apiErr) + service.HandleFailedBilling(c, info, apiErr) + var settlement model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).First(&settlement).Error) + require.Equal(t, model.BillingSettlementStatusApplied, settlement.Status) + require.EqualValues(t, finalQuota-100, settlement.FundingDelta, "delivered whitespace must not receive a full failure refund") + require.EqualValues(t, finalQuota-100, settlement.TokenDelta) + if scenario != "metadata_only" { + require.Contains(t, recorder.Body.String(), "event: response.output_text.delta") + require.Equal(t, model.BillingSettlementEffectApplied, settlement.EffectStatus) + var effect model.BillingSettlementEffect + require.NoError(t, common.UnmarshalJsonStr(settlement.EffectPayload, &effect)) + require.Equal(t, prompt, effect.PromptTokens) + require.Equal(t, completion, effect.CompletionTokens) + require.True(t, effect.QuotaIsActual) + require.EqualValues(t, finalQuota, effect.Quota) + replayInfo := &relaycommon.RelayInfo{RequestId: requestID, UserId: info.UserId, TokenId: info.TokenId, TokenKey: info.TokenKey, + OriginModelName: info.OriginModelName, UserSetting: info.UserSetting, ForcePreConsume: true} + replay, replayErr := service.NewBillingSession(c, replayInfo, 100) + require.Nil(t, replayErr) + require.NoError(t, replay.SettleWithEffect(finalQuota, &effect)) + require.NoError(t, replay.SettleWithEffect(finalQuota, &effect)) + replay.Refund(c) + } + var user model.User + var token model.Token + var channel model.Channel + require.NoError(t, db.First(&user, 201).Error) + require.NoError(t, db.First(&token, 202).Error) + require.NoError(t, db.First(&channel, 203).Error) + walletQuota := 1000 - finalQuota + if funding == "subscription_only" { + walletQuota = 1000 + var sub model.UserSubscription + require.NoError(t, db.First(&sub, 205).Error) + require.EqualValues(t, finalQuota, sub.AmountUsed) + require.Equal(t, model.BillingSettlementSourceSubscription, settlement.Source) + } + require.EqualValues(t, walletQuota, user.Quota) + require.EqualValues(t, 1000-finalQuota, token.RemainQuota) + require.EqualValues(t, finalQuota, user.UsedQuota) + require.EqualValues(t, finalQuota, token.UsedQuota) + require.EqualValues(t, finalQuota, channel.UsedQuota) + var count int64 + require.NoError(t, db.Model(&model.Log{}).Count(&count).Error) + wantLogs := 1 + if scenario == "metadata_only" { + wantLogs = 0 + } + require.EqualValues(t, wantLogs, count) + wantRequests := wantLogs + if scenario == "explicit_zero" { + wantRequests = 0 + } + require.EqualValues(t, wantRequests, user.RequestCount) + }) + } + } + } +} diff --git a/relay/responses_handler.go b/relay/responses_handler.go index a1bb4eb5..4f05c677 100644 --- a/relay/responses_handler.go +++ b/relay/responses_handler.go @@ -141,6 +141,15 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (maxAPIError * usage, maxAPIError := adaptor.DoResponse(c, httpResp, info) if maxAPIError != nil { + // Responses streams can deliver billable output before an upstream + // cancellation, disconnect, or incomplete terminal event. Reuse the + // existing partial settlement path instead of refunding the whole + // reservation or claiming a successful request. + if partial, ok := usage.(*dto.Usage); ok && partial != nil && + types.IsSkipRetryError(maxAPIError) && info.IsStream && + info.RelayMode == relayconstant.RelayModeResponses { + service.PostPartialConsumeQuota(c, info, partial) + } // reset status code 重置状态码 service.ResetStatusCode(maxAPIError, statusCodeMappingStr) return maxAPIError diff --git a/relay/responses_partial_settlement_test.go b/relay/responses_partial_settlement_test.go new file mode 100644 index 00000000..493622f3 --- /dev/null +++ b/relay/responses_partial_settlement_test.go @@ -0,0 +1,245 @@ +package relay + +import ( + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/constant" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/model" + "github.com/MAX-API-Next/MAX-API/pkg/billingexpr" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + relayconstant "github.com/MAX-API-Next/MAX-API/relay/constant" + "github.com/MAX-API-Next/MAX-API/service" + "github.com/MAX-API-Next/MAX-API/setting/model_setting" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "github.com/tidwall/sjson" + "gorm.io/gorm" + "gorm.io/gorm/logger" +) + +// Unlike a settlement recorder, this drives the real Responses relay and the +// durable funding/effect seam against an isolated SQLite ledger. +func TestResponsesHelperPartialSettlementUsesDurableLedger(t *testing.T) { + for _, provider := range []struct { + name string + channelType int + frames []string + }{ + {name: "responses_to_chat_global", channelType: constant.ChannelTypeOpenAI, frames: []string{ + `{"type":"response.created","response":{"id":"resp_partial_chat","status":"in_progress","usage":{"input_tokens":10,"output_tokens":2,"total_tokens":12}}}`, + `{"type":"response.output_text.delta","delta":"partial output"}`, + `{broken`, + }}, + {name: "responses_to_chat", channelType: constant.ChannelTypeAdvancedCustom, frames: []string{ + `{"type":"response.created","response":{"id":"resp_partial_chat","status":"in_progress","usage":{"input_tokens":10,"output_tokens":2,"total_tokens":12}}}`, + `{"type":"response.output_text.delta","delta":"partial output"}`, + `{broken`, + }}, + {name: "claude", channelType: constant.ChannelTypeAnthropic, frames: []string{ + `{"type":"message_start","message":{"id":"msg_test","model":"gpt-test","usage":{"input_tokens":10,"output_tokens":2}}}`, + `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial output"}}`, + `{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":2}}`, + `{"type":"error","error":{"type":"overloaded_error","message":"synthetic provider failure"}}`, + }}, + {name: "gemini", channelType: constant.ChannelTypeGemini, frames: []string{ + `{"candidates":[{"content":{"role":"model","parts":[{"text":"partial output"}]}}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":2,"totalTokenCount":12}}`, + }}, + {name: "openai", channelType: constant.ChannelTypeOpenAI, frames: []string{ + `{"type":"response.output_text.delta","delta":"partial output"}`, + `{"type":"response.failed","response":{"status":"failed","usage":{"input_tokens":10,"output_tokens":2,"total_tokens":12},"error":{"type":"server_error","message":"synthetic provider failure"}}}`, + }}, + } { + for _, scenario := range []struct { + name string + ratio float64 + quota int + zeroUsage bool + expr string + }{ + {name: "paid", ratio: 1, quota: 12}, + {name: "free", ratio: 0, quota: 0}, + {name: "minimum_quota", ratio: 0.001, quota: 1}, + {name: "additional_debit", ratio: 10, quota: 120}, + {name: "explicit_zero", ratio: 1, quota: 0, zeroUsage: true}, + {name: "tiered_zero", ratio: 1, quota: 0, expr: `tier("free", 0)`}, + {name: "tiered_rounding_zero", ratio: 1, quota: 0, expr: `tier("small", 0.01)`}, + {name: "expression_failure", ratio: 1, quota: 100, expr: `unknown_function()`}, + } { + for _, disconnected := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/%s/disconnect_%t", provider.name, scenario.name, disconnected), func(t *testing.T) { + finalQuota := scenario.quota + if strings.HasPrefix(provider.name, "responses_to_chat") && disconnected && finalQuota == 0 { + // Chat retains its original abnormal-stream reservation + // fallback; Responses-only zero-price exemptions do not expand. + finalQuota = 100 + } + promptTokens, completionTokens := 10, 2 + frames := append([]string(nil), provider.frames...) + if strings.HasPrefix(provider.name, "responses_to_chat") && disconnected { + // Isolate a transport error without a competing malformed + // handler stop, keeping the existing EndReason deterministic. + frames = frames[:2] + } + if scenario.zeroUsage { + promptTokens, completionTokens = 0, 0 + paths := map[string][]string{ + "responses_to_chat_global": {"response.usage.input_tokens", "response.usage.output_tokens", "response.usage.total_tokens"}, + "responses_to_chat": {"response.usage.input_tokens", "response.usage.output_tokens", "response.usage.total_tokens"}, + "claude": {"message.usage.input_tokens", "message.usage.output_tokens", "usage.output_tokens"}, + "gemini": {"usageMetadata.promptTokenCount", "usageMetadata.candidatesTokenCount", "usageMetadata.totalTokenCount"}, + "openai": {"response.usage.input_tokens", "response.usage.output_tokens", "response.usage.total_tokens"}, + } + for idx, frame := range frames { + for _, path := range paths[provider.name] { + // Only replace existing counters, preserving absent usage on text frames. + if !strings.Contains(frame, "usage") { + continue + } + if provider.name == "claude" && ((idx == 0 && strings.HasPrefix(path, "usage.")) || (idx != 0 && strings.HasPrefix(path, "message."))) { + continue + } + var err error + frame, err = sjson.Set(frame, path, 0) + require.NoError(t, err) + } + frames[idx] = frame + } + } + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}, &model.Channel{}, &model.Log{}, &model.BillingLogReceipt{}, + &model.BillingSettlement{}, &model.BillingPreConsumeSelection{}, &model.CacheInvalidationTask{}, + &model.UserSubscription{}, &model.SubscriptionPreConsumeRecord{}, &model.SubscriptionPlan{})) + oldDB, oldLogDB := model.DB, model.LOG_DB + oldSQLite, oldRedis, oldBatch, oldLogConsume := common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled + model.DB, model.LOG_DB = db, db + common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled = true, false, false, true + t.Cleanup(func() { + model.DB, model.LOG_DB = oldDB, oldLogDB + common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled = oldSQLite, oldRedis, oldBatch, oldLogConsume + _ = sqlDB.Close() + }) + require.NoError(t, db.Create(&model.User{Id: 101, Username: "responses-ledger-test", Quota: 1000}).Error) + require.NoError(t, db.Create(&model.Token{Id: 102, UserId: 101, Key: "synthetic-responses-ledger-key", RemainQuota: 1000, Status: common.TokenStatusEnabled}).Error) + require.NoError(t, db.Create(&model.Channel{Id: 103}).Error) + gin.SetMode(gin.TestMode) + service.InitHttpClient() + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + if disconnected { + w.Header().Set("Content-Length", "100000") + } + _, _ = w.Write([]byte("data: " + strings.Join(frames, "\n\ndata: ") + "\n\n")) + })) + t.Cleanup(upstream.Close) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + c.Set(string(constant.ContextKeyChannelType), provider.channelType) + c.Set(string(constant.ContextKeyChannelBaseUrl), upstream.URL) + c.Set(string(constant.ContextKeyChannelKey), "synthetic-test-key") + c.Set(string(constant.ContextKeyChannelId), 103) + c.Set(string(constant.ContextKeyOriginalModel), "gpt-test") + requestID := "responses-durable-" + provider.name + c.Set(common.RequestIdKey, requestID) + info := &relaycommon.RelayInfo{RequestId: requestID, UserId: 101, TokenId: 102, TokenKey: "synthetic-responses-ledger-key", ForcePreConsume: true, + RelayMode: relayconstant.RelayModeResponses, RelayFormat: types.RelayFormatOpenAIResponses, IsStream: true, DisablePing: true, + OriginModelName: "gpt-test", StartTime: time.Now(), UsingGroup: "default", + UserSetting: dto.UserSetting{BillingPreference: "wallet_only", QuotaWarningThreshold: 1}, + PriceData: types.PriceData{ModelRatio: 1, CompletionRatio: 1, GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 1}}, + Request: &dto.OpenAIResponsesRequest{Model: "gpt-test", Input: []byte(`"test"`), Stream: common.GetPointer(true)}} + info.SetEstimatePromptTokens(10) + if strings.HasPrefix(provider.name, "responses_to_chat") { + oldGlobal := *model_setting.GetGlobalSettings() + *model_setting.GetGlobalSettings() = model_setting.GlobalSettings{} + t.Cleanup(func() { *model_setting.GetGlobalSettings() = oldGlobal }) + if provider.name == "responses_to_chat_global" { + model_setting.GetGlobalSettings().ChatCompletionsToResponsesPolicy = model_setting.ChatCompletionsToResponsesPolicy{Enabled: true, AllChannels: true, ModelPatterns: []string{"^gpt-test$"}} + } + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + c.Set(string(constant.ContextKeyChannelOtherSetting), dto.ChannelOtherSettings{AdvancedCustom: &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{ + {IncomingPath: "/v1/chat/completions", UpstreamPath: "/v1/responses", Converter: dto.AdvancedCustomConverterOpenAIChatCompletionsToOpenAIResponses}, + }}}) + info.RelayMode, info.RelayFormat = relayconstant.RelayModeChatCompletions, types.RelayFormatOpenAI + info.Request = &dto.GeneralOpenAIRequest{Model: "gpt-test", Stream: common.GetPointer(true)} + } + info.PriceData.ModelRatio = scenario.ratio + if scenario.expr != "" { + info.TieredBillingSnapshot = &billingexpr.BillingSnapshot{BillingMode: "tiered_expr", ExprString: scenario.expr, ExprHash: billingexpr.ExprHashString(scenario.expr), GroupRatio: 1, QuotaPerUnit: common.QuotaPerUnit, EstimatedQuotaAfterGroup: 100} + } + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + var apiErr *types.MaxAPIError + if strings.HasPrefix(provider.name, "responses_to_chat") { + apiErr = TextHelper(c, info) + } else { + apiErr = ResponsesHelper(c, info) + } + require.NotNil(t, apiErr) + require.True(t, types.IsSkipRetryError(apiErr)) + require.EqualValues(t, 1, calls.Load()) + service.HandleFailedBilling(c, info, apiErr) + service.HandleFailedBilling(c, info, apiErr) + var settlement model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).First(&settlement).Error) + require.Equal(t, model.BillingSettlementStatusApplied, settlement.Status) + require.Equal(t, model.BillingSettlementEffectApplied, settlement.EffectStatus) + require.EqualValues(t, finalQuota-100, settlement.FundingDelta) + require.EqualValues(t, finalQuota-100, settlement.TokenDelta) + var effect model.BillingSettlementEffect + require.NoError(t, common.UnmarshalJsonStr(settlement.EffectPayload, &effect)) + require.Equal(t, promptTokens, effect.PromptTokens) + require.Equal(t, completionTokens, effect.CompletionTokens) + require.True(t, effect.QuotaIsActual) + require.EqualValues(t, finalQuota, effect.Quota) + // Reconstruct a session to exercise durable duplicate handling rather + // than relying only on the original session's in-memory settled bit. + replayInfo := &relaycommon.RelayInfo{RequestId: info.RequestId, UserId: info.UserId, TokenId: info.TokenId, + TokenKey: info.TokenKey, OriginModelName: info.OriginModelName, UserSetting: info.UserSetting, ForcePreConsume: true} + replay, replayErr := service.NewBillingSession(c, replayInfo, 100) + require.Nil(t, replayErr) + require.NoError(t, replay.SettleWithEffect(finalQuota, &effect)) + require.NoError(t, replay.SettleWithEffect(finalQuota, &effect)) + replay.Refund(c) + var user model.User + require.NoError(t, db.First(&user, 101).Error) + var token model.Token + require.NoError(t, db.First(&token, 102).Error) + var channel model.Channel + require.NoError(t, db.First(&channel, 103).Error) + require.EqualValues(t, 1000-finalQuota, user.Quota) + require.EqualValues(t, finalQuota, user.UsedQuota) + requestCount := 1 + if scenario.zeroUsage && finalQuota == 0 { + requestCount = 0 + } + require.EqualValues(t, requestCount, user.RequestCount) + require.EqualValues(t, 1000-finalQuota, token.RemainQuota) + require.EqualValues(t, finalQuota, token.UsedQuota) + require.EqualValues(t, finalQuota, channel.UsedQuota) + var count int64 + require.NoError(t, db.Model(&model.Log{}).Count(&count).Error) + require.EqualValues(t, 1, count, fmt.Sprintf("duplicate consume log for %s", provider.name)) + var log model.Log + require.NoError(t, db.First(&log).Error) + require.Equal(t, finalQuota, log.Quota) + require.Equal(t, promptTokens, log.PromptTokens) + require.Equal(t, completionTokens, log.CompletionTokens) + }) + } + } + } +} diff --git a/relay/responses_to_chat_tool_settlement_test.go b/relay/responses_to_chat_tool_settlement_test.go new file mode 100644 index 00000000..781162f1 --- /dev/null +++ b/relay/responses_to_chat_tool_settlement_test.go @@ -0,0 +1,194 @@ +package relay + +import ( + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/constant" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/model" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + relayconstant "github.com/MAX-API-Next/MAX-API/relay/constant" + "github.com/MAX-API-Next/MAX-API/service" + "github.com/MAX-API-Next/MAX-API/setting/model_setting" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "gorm.io/gorm" + "gorm.io/gorm/logger" +) + +// Exercise both real TextHelper routes and the original funding/effect ledger. +// Unknown usage must count only delivered tool data, including after write failure. +func TestResponsesToChatToolPartialSettlementUsesDurableLedger(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitHttpClient() + const args = `{"query":"weather tomorrow","city":"Shanghai"}` + for _, route := range []string{"custom", "global"} { + for _, funding := range []string{"wallet_only", "subscription_only"} { + for _, usageKind := range []string{"missing", "known", "zero"} { + for _, ending := range []string{"eof", "malformed", "write_failure", "no_delivery"} { + t.Run(strings.Join([]string{route, funding, usageKind, ending}, "/"), func(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}, &model.Channel{}, &model.Log{}, &model.BillingLogReceipt{}, + &model.BillingSettlement{}, &model.BillingPreConsumeSelection{}, &model.CacheInvalidationTask{}, + &model.UserSubscription{}, &model.SubscriptionPreConsumeRecord{}, &model.SubscriptionPlan{})) + oldDB, oldLogDB := model.DB, model.LOG_DB + oldSQLite, oldRedis, oldBatch, oldConsume := common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled + oldGlobal := *model_setting.GetGlobalSettings() + model.DB, model.LOG_DB = db, db + common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled = true, false, false, true + *model_setting.GetGlobalSettings() = model_setting.GlobalSettings{} + t.Cleanup(func() { + model.DB, model.LOG_DB = oldDB, oldLogDB + common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled = oldSQLite, oldRedis, oldBatch, oldConsume + *model_setting.GetGlobalSettings() = oldGlobal + _ = sqlDB.Close() + }) + require.NoError(t, db.Create(&model.User{Id: 801, Username: "tool-ledger-test", Quota: 1000}).Error) + require.NoError(t, db.Create(&model.Token{Id: 802, UserId: 801, Key: "synthetic-tool-ledger", RemainQuota: 1000, Status: common.TokenStatusEnabled}).Error) + require.NoError(t, db.Create(&model.Channel{Id: 803}).Error) + if funding == "subscription_only" { + require.NoError(t, db.Create(&model.SubscriptionPlan{Id: 804, Title: "tool-ledger-plan", Enabled: true, TotalAmount: 1000, QuotaResetPeriod: model.SubscriptionResetNever}).Error) + require.NoError(t, db.Create(&model.UserSubscription{Id: 805, UserId: 801, PlanId: 804, AmountTotal: 1000, Status: "active", + StartTime: time.Now().Add(-time.Hour).Unix(), EndTime: time.Now().Add(time.Hour).Unix()}).Error) + } + response := map[string]any{"id": "resp_tool_partial", "status": "in_progress"} + prompt, completion := 10, service.EstimateTokenByModel("gpt-test", "lookup"+args) + require.Positive(t, completion) + if usageKind != "missing" { + completion = 2 + if usageKind == "zero" { + prompt, completion = 0, 0 + } + response["usage"] = map[string]any{"input_tokens": prompt, "output_tokens": completion, "total_tokens": prompt + completion} + } + created, err := common.Marshal(map[string]any{"type": "response.created", "response": response}) + require.NoError(t, err) + argumentDelta, err := common.Marshal(map[string]any{"type": "response.function_call_arguments.delta", "item_id": "fc_partial", "delta": args}) + require.NoError(t, err) + frames := []string{string(created), `{"type":"response.output_item.added","item":{"type":"function_call","id":"fc_partial","call_id":"call_partial","name":"lookup","arguments":""}}`, string(argumentDelta)} + if ending == "malformed" { + frames = append(frames, `{broken`) + } + if ending == "write_failure" { + frames = append(frames, `{"type":"response.function_call_arguments.delta","item_id":"fc_partial","delta":"undelivered"}`) + } + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, "data: "+strings.Join(frames, "\n\ndata: ")+"\n\n") + })) + t.Cleanup(upstream.Close) + recorder := httptest.NewRecorder() + var writer http.ResponseWriter = recorder + if ending == "write_failure" { + writer = &responsesSelectiveFailureWriter{recorder, "undelivered"} + } + if ending == "no_delivery" { + writer = &responsesSelectiveFailureWriter{recorder, `"name":"lookup"`} + } + c, _ := gin.CreateTestContext(writer) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + channelType := constant.ChannelTypeAdvancedCustom + if route == "global" { + channelType = constant.ChannelTypeOpenAI + model_setting.GetGlobalSettings().ChatCompletionsToResponsesPolicy = model_setting.ChatCompletionsToResponsesPolicy{Enabled: true, AllChannels: true, ModelPatterns: []string{"^gpt-test$"}} + } + c.Set(string(constant.ContextKeyChannelType), channelType) + c.Set(string(constant.ContextKeyChannelBaseUrl), upstream.URL) + c.Set(string(constant.ContextKeyChannelKey), "synthetic-test-key") + c.Set(string(constant.ContextKeyChannelId), 803) + c.Set(string(constant.ContextKeyOriginalModel), "gpt-test") + c.Set(string(constant.ContextKeyChannelOtherSetting), dto.ChannelOtherSettings{AdvancedCustom: &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{ + {IncomingPath: "/v1/chat/completions", UpstreamPath: "/v1/responses", Converter: dto.AdvancedCustomConverterOpenAIChatCompletionsToOpenAIResponses}, + }}}) + requestID := "tool-ledger-" + strings.Join([]string{route, funding, usageKind, ending}, "-") + c.Set(common.RequestIdKey, requestID) + info := &relaycommon.RelayInfo{RequestId: requestID, UserId: 801, TokenId: 802, TokenKey: "synthetic-tool-ledger", ForcePreConsume: true, + RelayMode: relayconstant.RelayModeChatCompletions, RelayFormat: types.RelayFormatOpenAI, IsStream: true, DisablePing: true, + OriginModelName: "gpt-test", StartTime: time.Now(), UsingGroup: "default", UserSetting: dto.UserSetting{BillingPreference: funding, QuotaWarningThreshold: 1}, + PriceData: types.PriceData{ModelRatio: 1, CompletionRatio: 1, GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 1}}, + Request: &dto.GeneralOpenAIRequest{Model: "gpt-test", Stream: common.GetPointer(true)}} + info.SetEstimatePromptTokens(10) + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + apiErr := TextHelper(c, info) + require.NotNil(t, apiErr) + require.True(t, types.IsSkipRetryError(apiErr)) + require.EqualValues(t, 1, calls.Load()) + require.NotContains(t, recorder.Body.String(), "undelivered") + require.NotContains(t, recorder.Body.String(), `"finish_reason":"stop"`) + service.HandleFailedBilling(c, info, apiErr) + service.HandleFailedBilling(c, info, apiErr) + quota := prompt + completion + if ending == "no_delivery" { + quota = 0 + } + var settlement model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).First(&settlement).Error) + require.Equal(t, model.BillingSettlementStatusApplied, settlement.Status) + require.EqualValues(t, quota-100, settlement.FundingDelta, "delivered tool data must participate in actual settlement") + require.EqualValues(t, quota-100, settlement.TokenDelta) + if ending != "no_delivery" { + require.Equal(t, model.BillingSettlementEffectApplied, settlement.EffectStatus) + var effect model.BillingSettlementEffect + require.NoError(t, common.UnmarshalJsonStr(settlement.EffectPayload, &effect)) + require.Equal(t, prompt, effect.PromptTokens) + require.Equal(t, completion, effect.CompletionTokens) + require.EqualValues(t, quota, effect.Quota) + require.True(t, effect.QuotaIsActual) + replayInfo := &relaycommon.RelayInfo{RequestId: requestID, UserId: info.UserId, TokenId: info.TokenId, TokenKey: info.TokenKey, + OriginModelName: info.OriginModelName, UserSetting: info.UserSetting, ForcePreConsume: true} + replay, replayErr := service.NewBillingSession(c, replayInfo, 100) + require.Nil(t, replayErr) + require.NoError(t, replay.SettleWithEffect(quota, &effect)) + require.NoError(t, replay.SettleWithEffect(quota, &effect)) + replay.Refund(c) + } + var user model.User + var token model.Token + var channel model.Channel + require.NoError(t, db.First(&user, 801).Error) + require.NoError(t, db.First(&token, 802).Error) + require.NoError(t, db.First(&channel, 803).Error) + walletQuota := 1000 - quota + if funding == "subscription_only" { + walletQuota = 1000 + var subscription model.UserSubscription + require.NoError(t, db.First(&subscription, 805).Error) + require.EqualValues(t, quota, subscription.AmountUsed) + require.Equal(t, model.BillingSettlementSourceSubscription, settlement.Source) + } + require.EqualValues(t, walletQuota, user.Quota) + require.EqualValues(t, quota, user.UsedQuota) + require.EqualValues(t, 1000-quota, token.RemainQuota) + require.EqualValues(t, quota, token.UsedQuota) + require.EqualValues(t, quota, channel.UsedQuota) + var count int64 + require.NoError(t, db.Model(&model.Log{}).Count(&count).Error) + wantLogs, wantRequests := 1, 1 + if ending == "no_delivery" { + wantLogs, wantRequests = 0, 0 + } else if usageKind == "zero" { + wantRequests = 0 + } + require.EqualValues(t, wantLogs, count) + require.EqualValues(t, wantRequests, user.RequestCount) + }) + } + } + } + } +} diff --git a/relay/responses_tools_roundtrip_test.go b/relay/responses_tools_roundtrip_test.go new file mode 100644 index 00000000..adb7749f --- /dev/null +++ b/relay/responses_tools_roundtrip_test.go @@ -0,0 +1,97 @@ +package relay + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/relay/channel/claude" + "github.com/MAX-API-Next/MAX-API/relay/channel/gemini" + "github.com/MAX-API-Next/MAX-API/relay/channel/openai" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestResponsesCustomToolRoundTripAcrossProviders(t *testing.T) { + tests := []struct { + name string + body string + frames []string + handle func(*gin.Context, *relaycommon.RelayInfo, *http.Response) (*dto.Usage, *types.MaxAPIError) + stream func(*gin.Context, *relaycommon.RelayInfo, *http.Response) (*dto.Usage, *types.MaxAPIError) + }{ + {name: "openai", body: `{"model":"test","choices":[{"message":{"role":"assistant","tool_calls":[{"id":"call_test","type":"function","function":{"name":"apply_patch","arguments":"{\"input\":\"patch body\"}"}}]},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":4,"completion_tokens":2,"total_tokens":6}}`, + frames: []string{ + `{"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_test","type":"function","function":{"arguments":"{\"input\":"}}]}}]}`, + `{"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"name":"apply_patch","arguments":"\"patch body\"}"}}]}}]}`, + `{"choices":[{"delta":{},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":4,"completion_tokens":2,"total_tokens":6}}`, + }, handle: openai.OaiChatToResponsesHandler, stream: openai.OaiChatToResponsesStreamHandler}, + {name: "claude", body: `{"type":"message","model":"test","stop_reason":"tool_use","content":[{"type":"tool_use","id":"call_test","name":"apply_patch","input":{"input":"patch body"}}],"usage":{"input_tokens":4,"output_tokens":2}}`, + frames: []string{ + `{"type":"message_start","message":{"model":"test","usage":{"input_tokens":4}}}`, + `{"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"call_test","name":"apply_patch","input":{}}}`, + `{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"input\":\"patch body\"}"}}`, + `{"type":"message_delta","delta":{"stop_reason":"tool_use"},"usage":{"output_tokens":2}}`, + }, handle: func(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.MaxAPIError) { + return claude.ClaudeHandler(c, resp, info) + }, + stream: func(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.MaxAPIError) { + return claude.ClaudeStreamHandler(c, resp, info) + }}, + {name: "gemini", body: `{"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"id":"call_test","name":"apply_patch","args":{"input":"patch body"}}}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":4,"candidatesTokenCount":2,"totalTokenCount":6}}`, + frames: []string{`{"candidates":[{"content":{"role":"model","parts":[{"functionCall":{"id":"call_test","name":"apply_patch","args":{"input":"patch body"}}}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":4,"candidatesTokenCount":2,"totalTokenCount":6}}`}, + handle: gemini.GeminiResponsesHandler, stream: gemini.GeminiResponsesStreamHandler}, + } + for _, tt := range tests { + for _, stream := range []bool{false, true} { + t.Run(tt.name+map[bool]string{false: "/nonstream", true: "/stream"}[stream], func(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + info := &relaycommon.RelayInfo{IsStream: stream, RelayFormat: types.RelayFormatOpenAIResponses, DisablePing: true, + OriginModelName: "test", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "test"}, + Request: &dto.OpenAIResponsesRequest{Tools: []byte(`[{"type":"custom","name":"apply_patch"}]`)}} + body := tt.body + handler := tt.handle + if stream { + body = "data: " + strings.Join(tt.frames, "\n\ndata: ") + "\n\n" + handler = tt.stream + } + usage, apiErr := handler(c, info, &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(body))}) + require.Nil(t, apiErr) + require.NotNil(t, usage) + var converted dto.OpenAIResponsesResponse + if stream { + require.Contains(t, recorder.Body.String(), "response.custom_tool_call_input.delta") + require.Contains(t, recorder.Body.String(), "response.custom_tool_call_input.done") + require.NotContains(t, recorder.Body.String(), "response.function_call_arguments") + for _, line := range strings.Split(recorder.Body.String(), "\n") { + if !strings.HasPrefix(line, "data: ") { + continue + } + var event dto.ResponsesStreamResponse + require.NoError(t, common.UnmarshalJsonStr(strings.TrimPrefix(line, "data: "), &event)) + if event.Type == "response.completed" { + converted = *event.Response + } + } + } else { + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &converted)) + } + require.Len(t, converted.Output, 1) + require.Equal(t, "custom_tool_call", converted.Output[0].Type) + require.Equal(t, "apply_patch", converted.Output[0].Name) + require.Equal(t, "patch body", converted.Output[0].Input) + require.Empty(t, converted.Output[0].Arguments) + require.Equal(t, "call_test", converted.Output[0].CallId) + }) + } + } +} diff --git a/relay/subscription_period_settlement_test.go b/relay/subscription_period_settlement_test.go new file mode 100644 index 00000000..e84c0fab --- /dev/null +++ b/relay/subscription_period_settlement_test.go @@ -0,0 +1,240 @@ +package relay + +import ( + "crypto/sha256" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/model" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/service" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestSubscriptionResetPreservesUnsettledStreamReservation(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitHttpClient() + service.InitTokenEncoders() + for _, advance := range []bool{false, true} { + t.Run(fmt.Sprintf("advance_reset_time_%t", advance), func(t *testing.T) { + funding := streamFundingCases()[0] + db := setupStreamFundingLedger(t, funding) + // Never-reset plans keep a zero period anchor even when the admin + // requests a new reset time. This makes the boundary deterministic. + require.NoError(t, db.Model(&model.UserSubscription{}).Where("id = ?", 955). + Updates(map[string]any{"last_reset_time": 0, "next_reset_time": 0}).Error) + channelType, frames := streamFundingFrames(t, "responses_native", "known") + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, "data: "+strings.Join(frames, "\n\ndata: ")+"\n\n") + })) + t.Cleanup(upstream.Close) + requestID := fmt.Sprintf("period-pending-%t", advance) + c, info, _ := streamFundingContext(t, "responses_native", requestID, upstream.URL, channelType, funding) + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + var beforeOutbox int64 + require.NoError(t, db.Model(&model.CacheInvalidationTask{}).Count(&beforeOutbox).Error) + + _, resetErr := model.AdminResetUserSubscriptionsByPlan(951, 954, advance) + require.Error(t, resetErr, "an unchanged period anchor must not erase an unsettled reservation") + var sub model.UserSubscription + require.NoError(t, db.First(&sub, 955).Error) + require.EqualValues(t, 100, sub.AmountUsed) + require.Zero(t, sub.LastResetTime) + var afterOutbox int64 + require.NoError(t, db.Model(&model.CacheInvalidationTask{}).Count(&afterOutbox).Error) + require.Equal(t, beforeOutbox, afterOutbox) + + apiErr := ResponsesHelper(c, info) + require.NotNil(t, apiErr) // The synthetic stream has no terminal event. + require.True(t, types.IsSkipRetryError(apiErr)) + service.HandleFailedBilling(c, info, apiErr) + service.HandleFailedBilling(c, info, apiErr) + assertStreamFundingBalances(t, db, funding, requestID, 6, true, 4, 2) + _, resetErr = model.AdminResetUserSubscriptionsByPlan(951, 954, advance) + require.NoError(t, resetErr, "a completed funding settlement no longer blocks the reset") + require.NoError(t, db.First(&sub, 955).Error) + require.Zero(t, sub.AmountUsed) + model.ProcessPendingBillingSettlementsOnce() + require.NoError(t, db.First(&sub, 955).Error) + require.Zero(t, sub.AmountUsed, "completed settlement replay must not debit the reset quota") + }) + } +} + +func TestSubscriptionPeriodBoundaryDelayedStreamSettlement(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitHttpClient() + service.InitTokenEncoders() + for _, route := range []string{"responses_native", "chat_native", "chat_to_responses", "responses_to_chat_custom", "responses_to_chat_global", "claude_to_responses", "gemini_to_responses"} { + for _, usage := range []string{"known", "zero", "missing", "above_reservation", "at_reservation"} { + for _, boundary := range []string{"keep_time_rejected", "admin_advanced", "due_reset", "expired"} { + t.Run(strings.Join([]string{route, usage, boundary}, "/"), func(t *testing.T) { + funding := streamFundingCases()[0] + db := setupStreamFundingLedger(t, funding) + now := model.GetDBTimestamp() + require.NoError(t, db.Model(&model.SubscriptionPlan{}).Where("id = ?", 954). + Update("quota_reset_period", model.SubscriptionResetDaily).Error) + model.InvalidateSubscriptionPlanCache(954) + t.Cleanup(func() { model.InvalidateSubscriptionPlanCache(954) }) + anchor := now - 3600 + if boundary == "due_reset" { + // Advance persisted scheduler timestamps instead of sleeping + // a day; the reservation captures the actual old-period anchor. + anchor = now - 86400 - 60 + } + require.NoError(t, db.Model(&model.UserSubscription{}).Where("id = ?", 955). + Updates(map[string]any{"last_reset_time": anchor, "next_reset_time": now + 86400, "end_time": now + 172800}).Error) + channelType, frames := streamFundingFrames(t, route, usage) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, "data: "+strings.Join(frames, "\n\ndata: ")+"\n\n") + })) + t.Cleanup(upstream.Close) + fingerprint := sha256.Sum256([]byte(t.Name())) + requestID := fmt.Sprintf("period-%x", fingerprint[:16]) + c, info, recorder := streamFundingContext(t, route, requestID, upstream.URL, channelType, funding) + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + switch boundary { + case "keep_time_rejected": + _, err := model.AdminResetUserSubscriptionsByPlan(951, 954, false) + require.ErrorIs(t, err, model.ErrSubscriptionResetPendingBilling) + case "admin_advanced": + _, err := model.AdminResetUserSubscriptionsByPlan(951, 954, true) + require.NoError(t, err) + case "due_reset": + require.NoError(t, db.Model(&model.UserSubscription{}).Where("id = ?", 955).Update("next_reset_time", now-1).Error) + count, err := model.ResetDueSubscriptions(10) + require.NoError(t, err) + require.Equal(t, 1, count) + var reset model.UserSubscription + require.NoError(t, db.First(&reset, 955).Error) + require.NotEqual(t, anchor, reset.LastResetTime) + case "expired": + require.NoError(t, db.Model(&model.UserSubscription{}).Where("id = ?", 955).Update("end_time", now-1).Error) + count, err := model.ExpireDueSubscriptions(10) + require.NoError(t, err) + require.Equal(t, 1, count) + require.NoError(t, db.Create(&model.UserSubscription{Id: 956, UserId: 951, PlanId: 954, Status: "active", AmountTotal: 1000, + StartTime: now - 3600, EndTime: now + 3600, LastResetTime: now}).Error) + } + newInfo := &relaycommon.RelayInfo{RequestId: requestID + "-new", UserId: 951, TokenId: 952, TokenKey: info.TokenKey, + OriginModelName: "gpt-test", ForcePreConsume: true, UserSetting: dto.UserSetting{BillingPreference: "subscription_only"}} + _, admissionErr := service.NewBillingSession(c, newInfo, 200) + require.Nil(t, admissionErr) + require.NoError(t, db.Create(&model.UserSubscription{Id: 957, UserId: 951, PlanId: 954, Status: "active", AmountTotal: 1000, + StartTime: now - 1800, EndTime: now + 1800, LastResetTime: now}).Error) + info.UserSetting.BillingPreference = "wallet_only" + var apiErr *types.MaxAPIError + if info.RelayFormat == types.RelayFormatOpenAI { + apiErr = TextHelper(c, info) + } else { + apiErr = ResponsesHelper(c, info) + } + if apiErr != nil { + require.True(t, types.IsSkipRetryError(apiErr)) + service.HandleFailedBilling(c, info, apiErr) + service.HandleFailedBilling(c, info, apiErr) + } + require.Contains(t, recorder.Body.String(), "partial output") + quota := 6 + switch usage { + case "zero": + quota = 0 + case "missing": + quota = 10 + service.EstimateTokenByModel("gpt-test", "partial output") + if route == "responses_native" { + quota = 10 + service.CountTextToken("partial output", "gpt-test") + } + case "above_reservation": + quota = 130 + case "at_reservation": + quota = 100 + } + var settlement model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).First(&settlement).Error) + require.Equal(t, "subscription", settlement.Source) + require.Equal(t, 955, settlement.SubscriptionID) + require.EqualValues(t, quota-100, settlement.FundingDelta) + crossed := boundary == "admin_advanced" || boundary == "due_reset" + if crossed { + require.Equal(t, model.BillingSettlementStatusManual, settlement.Status) + require.Contains(t, settlement.LastError, model.ErrSubscriptionSettlementPeriodChanged.Error()) + require.Zero(t, settlement.AppliedFundingDelta) + require.Zero(t, settlement.AppliedTokenDelta) + } else { + require.Equal(t, model.BillingSettlementStatusApplied, settlement.Status) + require.EqualValues(t, quota-100, settlement.AppliedFundingDelta) + require.EqualValues(t, quota-100, settlement.AppliedTokenDelta) + } + var outboxBefore int64 + require.NoError(t, db.Model(&model.CacheInvalidationTask{}).Count(&outboxBefore).Error) + for range 2 { + _, already, err := model.ApplyBillingSettlementOnce(model.BillingSettlementInput{ + OperationKey: settlement.OperationKey, Source: settlement.Source, UserID: settlement.UserID, + SubscriptionID: settlement.SubscriptionID, TokenID: settlement.TokenID, TokenKey: info.TokenKey, + FundingDelta: settlement.FundingDelta, TokenDelta: settlement.TokenDelta, SubscriptionPreConsumeRequestID: requestID, + }) + if crossed { + require.ErrorIs(t, err, model.ErrBillingSettlementManualReview) + require.False(t, already) + } else { + require.NoError(t, err) + require.True(t, already) + } + model.ProcessPendingBillingSettlementsOnce() + } + var outboxAfter int64 + require.NoError(t, db.Model(&model.CacheInvalidationTask{}).Count(&outboxAfter).Error) + require.Equal(t, outboxBefore, outboxAfter) + var user model.User + var token model.Token + var channel model.Channel + var sub, other model.UserSubscription + require.NoError(t, db.First(&user, 951).Error) + require.NoError(t, db.First(&token, 952).Error) + require.NoError(t, db.First(&channel, 953).Error) + require.NoError(t, db.First(&sub, 955).Error) + require.NoError(t, db.First(&other, 957).Error) + require.EqualValues(t, 1000, user.Quota) + require.Zero(t, other.AmountUsed) + wantQuota, wantUsed, wantLogs := quota, quota, 1 + if boundary == "keep_time_rejected" { + wantUsed += 200 + } else if crossed { + wantQuota, wantUsed, wantLogs = 100, 200, 0 + } + require.EqualValues(t, wantUsed, sub.AmountUsed) + require.EqualValues(t, 800-wantQuota, token.RemainQuota) + require.EqualValues(t, 200+wantQuota, token.UsedQuota) + projectedQuota, requests := quota, 1 + if crossed { + projectedQuota, requests = 0, 0 + } else if quota == 0 { + requests = 0 + } + require.EqualValues(t, projectedQuota, user.UsedQuota) + require.EqualValues(t, requests, user.RequestCount) + require.EqualValues(t, projectedQuota, channel.UsedQuota) + var logs, finalizes int64 + require.NoError(t, db.Model(&model.Log{}).Where("type = ?", model.LogTypeConsume).Count(&logs).Error) + require.EqualValues(t, wantLogs, logs) + require.NoError(t, db.Model(&model.BillingSettlement{}).Where("operation_key = ?", settlement.OperationKey).Count(&finalizes).Error) + require.EqualValues(t, 1, finalizes) + if boundary == "expired" { + require.Equal(t, "expired", sub.Status) + other = model.UserSubscription{} + require.NoError(t, db.First(&other, 956).Error) + require.EqualValues(t, 200, other.AmountUsed) + } + }) + } + } + } +} diff --git a/relay/subscription_stream_settlement_test.go b/relay/subscription_stream_settlement_test.go new file mode 100644 index 00000000..9565245f --- /dev/null +++ b/relay/subscription_stream_settlement_test.go @@ -0,0 +1,394 @@ +package relay + +import ( + "crypto/sha256" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/constant" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/model" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + relayconstant "github.com/MAX-API-Next/MAX-API/relay/constant" + "github.com/MAX-API-Next/MAX-API/service" + "github.com/MAX-API-Next/MAX-API/setting/model_setting" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "gorm.io/gorm" + "gorm.io/gorm/logger" +) + +type streamFundingCase struct { + name, preference, source string + wallet, subscriptionUsed int64 + allowOverflow bool +} + +func streamFundingCases() []streamFundingCase { + return []streamFundingCase{ + {name: "subscription_only", preference: "subscription_only", source: "subscription", wallet: 1000}, + {name: "subscription_first", preference: "subscription_first", source: "subscription", wallet: 1000}, + {name: "wallet_first_to_subscription", preference: "wallet_first", source: "subscription"}, + {name: "wallet_only", preference: "wallet_only", source: "wallet", wallet: 1000}, + {name: "wallet_first", preference: "wallet_first", source: "wallet", wallet: 1000}, + {name: "subscription_first_to_wallet", preference: "subscription_first", source: "wallet", wallet: 1000, subscriptionUsed: 1000, allowOverflow: true}, + } +} + +func setupStreamFundingLedger(t *testing.T, funding streamFundingCase) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + t.Cleanup(func() { _ = sqlDB.Close() }) + require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}, &model.Channel{}, &model.Log{}, &model.BillingLogReceipt{}, + &model.BillingSettlement{}, &model.BillingPreConsumeSelection{}, &model.CacheInvalidationTask{}, + &model.UserSubscription{}, &model.SubscriptionPreConsumeRecord{}, &model.SubscriptionPlan{})) + oldDB, oldLogDB := model.DB, model.LOG_DB + oldSQLite, oldRedis, oldBatch, oldConsume := common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled + model.DB, model.LOG_DB = db, db + common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled = true, false, false, true + t.Cleanup(func() { + model.DB, model.LOG_DB = oldDB, oldLogDB + common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled = oldSQLite, oldRedis, oldBatch, oldConsume + }) + require.NoError(t, db.Create(&model.User{Id: 951, Username: "stream-funding", Quota: funding.wallet}).Error) + require.NoError(t, db.Create(&model.Token{Id: 952, UserId: 951, Key: "synthetic-stream-funding", RemainQuota: 1000, Status: common.TokenStatusEnabled}).Error) + require.NoError(t, db.Create(&model.Channel{Id: 953}).Error) + require.NoError(t, db.Create(&model.SubscriptionPlan{Id: 954, Title: "stream-funding-plan", Enabled: true, + TotalAmount: 1000, AllowWalletOverflow: &funding.allowOverflow, QuotaResetPeriod: model.SubscriptionResetNever}).Error) + now := time.Now().Unix() + require.NoError(t, db.Create(&model.UserSubscription{Id: 955, UserId: 951, PlanId: 954, Status: "active", + AmountTotal: 1000, AmountUsed: funding.subscriptionUsed, AllowWalletOverflow: funding.allowOverflow, + StartTime: now - 3600, EndTime: now + 3600, LastResetTime: now - 3600}).Error) + return db +} + +func streamFundingFrames(t *testing.T, route, usage string) (int, []string) { + t.Helper() + prompt, completion := 4, 2 + if usage == "zero" { + prompt, completion = 0, 0 + } else if usage == "above_reservation" { + prompt, completion = 80, 50 + } else if usage == "at_reservation" { + prompt, completion = 80, 20 + } + encode := func(v map[string]any) string { + b, err := common.Marshal(v) + require.NoError(t, err) + return string(b) + } + switch route { + case "chat_native", "chat_to_responses": + frame := map[string]any{"choices": []any{map[string]any{"delta": map[string]any{"content": "partial output"}}}} + if usage != "missing" { + frame["usage"] = map[string]any{"prompt_tokens": prompt, "completion_tokens": completion, "total_tokens": prompt + completion} + } + channel := constant.ChannelTypeOpenAI + if route == "chat_to_responses" { + channel = constant.ChannelTypeAdvancedCustom + } + return channel, []string{encode(frame)} + case "claude_to_responses": + message := map[string]any{"id": "msg_funding", "model": "gpt-test"} + if usage != "missing" { + message["usage"] = map[string]any{"input_tokens": prompt, "output_tokens": 0} + } + frames := []string{encode(map[string]any{"type": "message_start", "message": message}), + `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial output"}}`} + if usage != "missing" { + frames = append(frames, encode(map[string]any{"type": "message_delta", "delta": map[string]any{"stop_reason": "end_turn"}, "usage": map[string]any{"output_tokens": completion}})) + } + return constant.ChannelTypeAnthropic, frames + case "gemini_to_responses": + frame := map[string]any{"candidates": []any{map[string]any{"content": map[string]any{"role": "model", "parts": []any{map[string]any{"text": "partial output"}}}}}} + if usage != "missing" { + frame["usageMetadata"] = map[string]any{"promptTokenCount": prompt, "candidatesTokenCount": completion, "totalTokenCount": prompt + completion} + } + return constant.ChannelTypeGemini, []string{encode(frame)} + default: + response := map[string]any{"id": "resp_funding", "model": "gpt-test", "status": "in_progress", "output": []any{}} + if usage != "missing" { + response["usage"] = map[string]any{"input_tokens": prompt, "output_tokens": completion, "total_tokens": prompt + completion} + } + channel := constant.ChannelTypeOpenAI + if route == "responses_to_chat_custom" { + channel = constant.ChannelTypeAdvancedCustom + } + return channel, []string{encode(map[string]any{"type": "response.created", "response": response}), + `{"type":"response.output_text.delta","delta":"partial output"}`} + } +} + +func streamFundingContext(t *testing.T, route, requestID, upstream string, channelType int, funding streamFundingCase) (*gin.Context, *relaycommon.RelayInfo, *httptest.ResponseRecorder) { + t.Helper() + oldGlobal := *model_setting.GetGlobalSettings() + *model_setting.GetGlobalSettings() = model_setting.GlobalSettings{} + t.Cleanup(func() { *model_setting.GetGlobalSettings() = oldGlobal }) + if route == "responses_to_chat_global" { + model_setting.GetGlobalSettings().ChatCompletionsToResponsesPolicy = model_setting.ChatCompletionsToResponsesPolicy{Enabled: true, AllChannels: true, ModelPatterns: []string{"^gpt-test$"}} + } + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + chat := route == "chat_native" || strings.HasPrefix(route, "responses_to_chat_") + path := "/v1/responses" + if chat { + path = "/v1/chat/completions" + } + c.Request = httptest.NewRequest(http.MethodPost, path, nil) + c.Set(common.RequestIdKey, requestID) + c.Set(string(constant.ContextKeyChannelType), channelType) + c.Set(string(constant.ContextKeyChannelBaseUrl), upstream) + c.Set(string(constant.ContextKeyChannelKey), "synthetic-stream-key") + c.Set(string(constant.ContextKeyChannelId), 953) + c.Set(string(constant.ContextKeyOriginalModel), "gpt-test") + if channelType == constant.ChannelTypeAdvancedCustom { + upstreamPath, converter := "/v1/chat/completions", dto.AdvancedCustomConverterOpenAIResponsesToOpenAIChatCompletions + if chat { + upstreamPath, converter = "/v1/responses", dto.AdvancedCustomConverterOpenAIChatCompletionsToOpenAIResponses + } + c.Set(string(constant.ContextKeyChannelOtherSetting), dto.ChannelOtherSettings{AdvancedCustom: &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{ + {IncomingPath: path, UpstreamPath: upstreamPath, Converter: converter}, + }}}) + } + info := &relaycommon.RelayInfo{RequestId: requestID, UserId: 951, TokenId: 952, TokenKey: "synthetic-stream-funding", ForcePreConsume: true, + RelayMode: relayconstant.RelayModeResponses, RelayFormat: types.RelayFormatOpenAIResponses, IsStream: true, DisablePing: true, + OriginModelName: "gpt-test", StartTime: time.Now(), UsingGroup: "default", UserSetting: dto.UserSetting{BillingPreference: funding.preference, QuotaWarningThreshold: 1}, + PriceData: types.PriceData{ModelRatio: 1, CompletionRatio: 1, GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 1}}, + Request: &dto.OpenAIResponsesRequest{Model: "gpt-test", Input: []byte(`"test"`), Stream: common.GetPointer(true)}} + if chat { + info.RelayMode, info.RelayFormat = relayconstant.RelayModeChatCompletions, types.RelayFormatOpenAI + info.Request = &dto.GeneralOpenAIRequest{Model: "gpt-test", Messages: []dto.Message{{Role: "user", Content: "test"}}, Stream: common.GetPointer(true)} + } + info.SetEstimatePromptTokens(10) + return c, info, recorder +} + +// Test the complete relay-to-ledger seam rather than a mock BillingSettler. +// Reconstructed sessions must retain the source even after preferences change +// and a more attractive subscription becomes available. +func TestStreamSettlementRetainsOriginalFunding(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitHttpClient() + service.InitTokenEncoders() + for _, route := range []string{"responses_native", "chat_native", "chat_to_responses", "responses_to_chat_custom", "responses_to_chat_global", "claude_to_responses", "gemini_to_responses"} { + for _, funding := range streamFundingCases() { + for _, usageKind := range []string{"known", "zero", "missing", "above_reservation", "at_reservation"} { + for _, ending := range []string{"malformed", "transport", "eof", "no_delivery"} { + t.Run(strings.Join([]string{route, funding.name, usageKind, ending}, "/"), func(t *testing.T) { + db := setupStreamFundingLedger(t, funding) + channelType, frames := streamFundingFrames(t, route, usageKind) + if ending == "malformed" { + frames = append(frames, `{broken`) + } else if ending == "no_delivery" { + frames = []string{`{broken`} + } + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + if ending == "transport" { + w.Header().Set("Content-Length", "100000") + } + _, _ = fmt.Fprint(w, "data: "+strings.Join(frames, "\n\ndata: ")+"\n\n") + })) + t.Cleanup(upstream.Close) + fingerprint := sha256.Sum256([]byte(t.Name())) + requestID := fmt.Sprintf("funding-%x", fingerprint[:16]) + c, info, recorder := streamFundingContext(t, route, requestID, upstream.URL, channelType, funding) + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + require.Equal(t, funding.source, info.BillingSource) + require.Equal(t, 100, info.FinalPreConsumedQuota) + var selection model.BillingPreConsumeSelection + require.NoError(t, db.Where("request_id = ?", requestID).First(&selection).Error) + require.Equal(t, funding.source, selection.Source) + require.EqualValues(t, 100, selection.RequestedQuota) + require.EqualValues(t, 100, selection.EffectiveQuota) + var apiErr *types.MaxAPIError + if info.RelayFormat == types.RelayFormatOpenAI { + apiErr = TextHelper(c, info) + } else { + apiErr = ResponsesHelper(c, info) + } + if apiErr != nil { + if ending != "no_delivery" { + require.True(t, types.IsSkipRetryError(apiErr)) + } + service.HandleFailedBilling(c, info, apiErr) + service.HandleFailedBilling(c, info, apiErr) + } else { + require.Equal(t, "eof", ending) + } + require.EqualValues(t, 1, calls.Load()) + prompt, completion := 4, 2 + if usageKind == "zero" { + prompt, completion = 0, 0 + } else if usageKind == "above_reservation" { + prompt, completion = 80, 50 + } else if usageKind == "at_reservation" { + prompt, completion = 80, 20 + } else if usageKind == "missing" { + prompt, completion = 10, service.EstimateTokenByModel("gpt-test", "partial output") + if route == "responses_native" { + completion = service.CountTextToken("partial output", "gpt-test") + } + } + quota := prompt + completion + if usageKind == "zero" && ending == "transport" && info.RelayFormat == types.RelayFormatOpenAI { + quota = 100 // Preserve Chat's abnormal zero-usage reservation policy. + } + if ending == "no_delivery" { + quota = 0 + require.NotContains(t, recorder.Body.String(), "partial output") + } else { + require.Contains(t, recorder.Body.String(), "partial output") + } + var settlement model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).First(&settlement).Error) + require.Equal(t, funding.source, settlement.Source) + require.Equal(t, model.BillingSettlementStatusApplied, settlement.Status) + require.EqualValues(t, quota-100, settlement.FundingDelta) + require.EqualValues(t, quota-100, settlement.AppliedFundingDelta) + require.EqualValues(t, quota-100, settlement.AppliedTokenDelta) + if ending != "no_delivery" { + var effect model.BillingSettlementEffect + require.Equal(t, model.BillingSettlementEffectApplied, settlement.EffectStatus) + require.NoError(t, common.UnmarshalJsonStr(settlement.EffectPayload, &effect)) + require.Equal(t, prompt, effect.PromptTokens) + require.Equal(t, completion, effect.CompletionTokens) + require.True(t, effect.QuotaIsActual) + require.EqualValues(t, quota, effect.Quota) + // This new earlier-expiring subscription must never own the replay. + require.NoError(t, db.Create(&model.UserSubscription{Id: 956, UserId: 951, PlanId: 954, Status: "active", AmountTotal: 1000, + StartTime: time.Now().Unix() - 1800, EndTime: time.Now().Unix() + 1800}).Error) + replayInfo := &relaycommon.RelayInfo{RequestId: requestID, UserId: 951, TokenId: 952, TokenKey: info.TokenKey, + OriginModelName: info.OriginModelName, ForcePreConsume: true, UserSetting: dto.UserSetting{BillingPreference: "wallet_only"}} + if funding.source == "wallet" { + replayInfo.UserSetting.BillingPreference = "subscription_only" + } + var beforeOutbox int64 + require.NoError(t, db.Model(&model.CacheInvalidationTask{}).Count(&beforeOutbox).Error) + replay, replayErr := service.NewBillingSession(c, replayInfo, 100) + require.Nil(t, replayErr) + require.Equal(t, funding.source, replayInfo.BillingSource) + require.NoError(t, replay.SettleWithEffect(quota, &effect)) + require.NoError(t, replay.SettleWithEffect(quota, &effect)) + replay.Refund(c) + var afterOutbox int64 + require.NoError(t, db.Model(&model.CacheInvalidationTask{}).Count(&afterOutbox).Error) + require.Equal(t, beforeOutbox, afterOutbox, "replay must not duplicate invalidation work") + var decoy model.UserSubscription + require.NoError(t, db.First(&decoy, 956).Error) + require.Zero(t, decoy.AmountUsed) + } else { + require.Empty(t, settlement.EffectPayload) + } + assertStreamFundingBalances(t, db, funding, requestID, quota, ending != "no_delivery", prompt, completion) + }) + } + } + } + } +} + +func assertStreamFundingBalances(t *testing.T, db *gorm.DB, funding streamFundingCase, requestID string, quota int, delivered bool, prompt, completion int) { + t.Helper() + var user model.User + var token model.Token + var channel model.Channel + var sub model.UserSubscription + require.NoError(t, db.First(&user, 951).Error) + require.NoError(t, db.First(&token, 952).Error) + require.NoError(t, db.First(&channel, 953).Error) + require.NoError(t, db.First(&sub, 955).Error) + wallet, subUsed := funding.wallet, funding.subscriptionUsed + if funding.source == "subscription" { + subUsed += int64(quota) + var record model.SubscriptionPreConsumeRecord + require.NoError(t, db.Where("request_id = ?", requestID).First(&record).Error) + require.Equal(t, sub.Id, record.UserSubscriptionId) + require.Equal(t, sub.LastResetTime, record.SubscriptionLastResetTime) + require.EqualValues(t, 100, record.PreConsumed) + wantStatus := "consumed" + if !delivered { + wantStatus = "refunded" + } + require.Equal(t, wantStatus, record.Status) + } else { + wallet -= int64(quota) + var records int64 + require.NoError(t, db.Model(&model.SubscriptionPreConsumeRecord{}).Count(&records).Error) + require.Zero(t, records) + } + require.EqualValues(t, wallet, user.Quota) + require.EqualValues(t, subUsed, sub.AmountUsed) + require.EqualValues(t, 1000-quota, token.RemainQuota) + require.EqualValues(t, quota, token.UsedQuota) + require.EqualValues(t, quota, user.UsedQuota) + require.EqualValues(t, quota, channel.UsedQuota) + wantLogs, wantRequests := 0, 0 + if delivered { + wantLogs = 1 + if quota != 0 { + wantRequests = 1 + } + } + require.EqualValues(t, wantRequests, user.RequestCount) + var logs, selections, finalizes int64 + require.NoError(t, db.Model(&model.Log{}).Where("type = ?", model.LogTypeConsume).Count(&logs).Error) + require.NoError(t, db.Model(&model.BillingPreConsumeSelection{}).Count(&selections).Error) + require.NoError(t, db.Model(&model.BillingSettlement{}).Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).Count(&finalizes).Error) + require.EqualValues(t, wantLogs, logs) + require.EqualValues(t, 1, selections) + require.EqualValues(t, 1, finalizes) + if delivered { + var log model.Log + require.NoError(t, db.Where("type = ?", model.LogTypeConsume).First(&log).Error) + require.Equal(t, quota, log.Quota) + require.Equal(t, prompt, log.PromptTokens) + require.Equal(t, completion, log.CompletionTokens) + } +} + +func TestStreamFundingAdmissionRespectsSubscriptionOverflow(t *testing.T) { + for _, mixed := range []bool{false, true} { + t.Run(fmt.Sprintf("mixed_subscriptions_%t", mixed), func(t *testing.T) { + funding := streamFundingCase{preference: "subscription_first", wallet: 1000, subscriptionUsed: 1000} + db := setupStreamFundingLedger(t, funding) + if mixed { + require.NoError(t, db.Create(&model.UserSubscription{Id: 956, UserId: 951, PlanId: 954, Status: "active", AmountTotal: 1000, + AmountUsed: 1000, AllowWalletOverflow: true, EndTime: time.Now().Unix() + 1800}).Error) + } + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { calls.Add(1) })) + t.Cleanup(upstream.Close) + c, info, _ := streamFundingContext(t, "responses_native", "strict-stream-overflow", upstream.URL, constant.ChannelTypeOpenAI, funding) + apiErr := service.PreConsumeBilling(c, 100, info) + require.NotNil(t, apiErr) + require.Equal(t, types.ErrorCodeInsufficientUserQuota, apiErr.GetErrorCode()) + require.Zero(t, calls.Load()) + var user model.User + var token model.Token + require.NoError(t, db.First(&user, 951).Error) + require.NoError(t, db.First(&token, 952).Error) + require.EqualValues(t, 1000, user.Quota) + require.EqualValues(t, 1000, token.RemainQuota) + for _, table := range []any{&model.Log{}, &model.BillingSettlement{}, &model.BillingPreConsumeSelection{}, &model.SubscriptionPreConsumeRecord{}, &model.CacheInvalidationTask{}} { + var count int64 + require.NoError(t, db.Model(table).Count(&count).Error) + require.Zero(t, count) + } + }) + } +} diff --git a/service/billing_playground_replay_test.go b/service/billing_playground_replay_test.go new file mode 100644 index 00000000..873d8607 --- /dev/null +++ b/service/billing_playground_replay_test.go @@ -0,0 +1,84 @@ +package service + +import ( + "errors" + "testing" + + "github.com/MAX-API-Next/MAX-API/model" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func TestPlaygroundWalletPaidReplayRejectsBeforeContextUse(t *testing.T) { + for _, missingContext := range []bool{false, true} { + name := "real_context" + if missingContext { + name = "nil_context" + } + t.Run(name, func(t *testing.T) { + truncate(t) + seedUser(t, 991, 100) + seedToken(t, 992, 991, "synthetic-playground-guard", 100) + ctx, _ := gin.CreateTestContext(nil) + info := &relaycommon.RelayInfo{RequestId: "paid-playground-guard", UserId: 991, TokenId: 992, + TokenKey: "synthetic-playground-guard", OriginModelName: "guard-model", ForcePreConsume: true} + paid := &BillingSession{relayInfo: info, funding: &WalletFunding{userId: 991}} + require.Nil(t, paid.preConsume(ctx, 10)) + replayInfo := &relaycommon.RelayInfo{RequestId: info.RequestId, UserId: info.UserId, TokenId: info.TokenId, + TokenKey: info.TokenKey, OriginModelName: info.OriginModelName, ForcePreConsume: true, IsPlayground: true} + replay := &BillingSession{relayInfo: replayInfo, funding: &WalletFunding{userId: 991}} + if missingContext { + ctx = nil + } + var apiErr *types.MaxAPIError + require.NotPanics(t, func() { apiErr = replay.preConsume(ctx, 10) }) + require.NotNil(t, apiErr) + require.ErrorIs(t, apiErr, model.ErrBillingSettlementOperationConflict) + require.True(t, types.IsSkipRetryError(apiErr)) + require.EqualValues(t, 90, getUserQuota(t, 991)) + require.EqualValues(t, 90, getTokenRemainQuota(t, 992)) + }) + } +} + +func TestPlaygroundWalletReplayLookupFailureFailsClosed(t *testing.T) { + for _, missingDatabase := range []bool{false, true} { + name := "query_failure" + if missingDatabase { + name = "missing_database" + } + t.Run(name, func(t *testing.T) { + truncate(t) + seedUser(t, 991, 100) + seedToken(t, 992, 991, "synthetic-playground-lookup", 100) + db := model.DB + const callbackName = "test:playground-selection-query-failure" + if missingDatabase { + model.DB = nil + t.Cleanup(func() { model.DB = db }) + } else { + require.NoError(t, db.Callback().Query().Before("gorm:query").Register(callbackName, func(tx *gorm.DB) { + if tx.Statement.Table == "billing_pre_consume_selections" { + tx.AddError(errors.New("synthetic selection query failure")) + } + })) + t.Cleanup(func() { _ = db.Callback().Query().Remove(callbackName) }) + } + session := &BillingSession{funding: &WalletFunding{userId: 991}, relayInfo: &relaycommon.RelayInfo{ + RequestId: "unknown-playground-identity", UserId: 991, IsPlayground: true, ForcePreConsume: true}} + var apiErr *types.MaxAPIError + require.NotPanics(t, func() { apiErr = session.preConsume(nil, 10) }) + require.NotNil(t, apiErr) + require.True(t, types.IsSkipRetryError(apiErr)) + model.DB = db + if !missingDatabase { + require.NoError(t, db.Callback().Query().Remove(callbackName)) + } + require.EqualValues(t, 100, getUserQuota(t, 991)) + require.EqualValues(t, 100, getTokenRemainQuota(t, 992)) + }) + } +} diff --git a/service/billing_session.go b/service/billing_session.go index f4084f89..b7d59cbf 100644 --- a/service/billing_session.go +++ b/service/billing_session.go @@ -736,6 +736,18 @@ func (s *BillingSession) Reserve(targetQuota int) error { // preConsume 执行预扣费：信任检查 -> 令牌预扣 -> 资金来源预扣。 // 任一步骤失败时原子回滚已完成的步骤。 func (s *BillingSession) preConsume(c *gin.Context, quota int) *types.MaxAPIError { + if _, wallet := s.funding.(*WalletFunding); wallet && s.relayInfo.IsPlayground { + // Playground wallets use legacy hooks without an operation key. Never + // let that path reinterpret an already persisted paid reservation. + _, persisted, err := model.ResolveBillingPreConsumeSource(s.relayInfo.RequestId) + if err != nil { + return mapAtomicPreConsumeError(err) + } + if persisted { + return mapAtomicPreConsumeError(fmt.Errorf("%w: playground wallet cannot reuse persisted request %s", + model.ErrBillingSettlementOperationConflict, s.relayInfo.RequestId)) + } + } effectiveQuota := quota // ---- 信任额度旁路 ---- diff --git a/service/log_info_generate.go b/service/log_info_generate.go index bc9d73dd..0fe78569 100644 --- a/service/log_info_generate.go +++ b/service/log_info_generate.go @@ -83,6 +83,9 @@ func GenerateTextOtherInfo(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, m other["model_price"] = modelPrice other["user_group_ratio"] = userGroupRatio other["frt"] = float64(relayInfo.FirstResponseTime.UnixMilli() - relayInfo.StartTime.UnixMilli()) + if !relayInfo.FirstResultTime.IsZero() && !relayInfo.StartTime.IsZero() { + other["first_result_ms"] = float64(relayInfo.FirstResultTime.UnixMilli() - relayInfo.StartTime.UnixMilli()) + } if relayInfo.ReasoningEffort != "" { other["reasoning_effort"] = relayInfo.ReasoningEffort } diff --git a/service/log_info_generate_test.go b/service/log_info_generate_test.go index 38baff88..2d2eb719 100644 --- a/service/log_info_generate_test.go +++ b/service/log_info_generate_test.go @@ -51,6 +51,19 @@ func TestGenerateTextOtherInfoDoesNotMarkSingleChannelAsRetry(t *testing.T) { require.NotContains(t, other, "retry_log") } +func TestGenerateTextOtherInfoSeparatesFirstFrameAndResultLatency(t *testing.T) { + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + start := time.Unix(100, 0) + info := &relaycommon.RelayInfo{StartTime: start, FirstResponseTime: start.Add(10 * time.Millisecond), + FirstResultTime: start.Add(200 * time.Millisecond), ChannelMeta: &relaycommon.ChannelMeta{}} + other := GenerateTextOtherInfo(ctx, info, 1, 1, 1, 0, 0, 0, 1) + require.Equal(t, float64(10), other["frt"]) + require.Equal(t, float64(200), other["first_result_ms"]) + info.FirstResultTime = time.Time{} + other = GenerateTextOtherInfo(ctx, info, 1, 1, 1, 0, 0, 0, 1) + require.NotContains(t, other, "first_result_ms", "missing output must not be logged as zero latency") +} + func TestAttachQuotaSaturationToOtherPreservesMalformedAdminInfo(t *testing.T) { clamp := &common.QuotaClamp{ Op: "test", diff --git a/service/openai_chat_responses_compat.go b/service/openai_chat_responses_compat.go index dda5db50..7ae53526 100644 --- a/service/openai_chat_responses_compat.go +++ b/service/openai_chat_responses_compat.go @@ -2,13 +2,40 @@ package service import ( "github.com/MAX-API-Next/MAX-API/dto" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" "github.com/MAX-API-Next/MAX-API/service/openaicompat" + "strings" ) func ChatCompletionsRequestToResponsesRequest(req *dto.GeneralOpenAIRequest) (*dto.OpenAIResponsesRequest, error) { return openaicompat.ChatCompletionsRequestToResponsesRequest(req) } +func ResponsesCustomToolNames(info *relaycommon.RelayInfo) map[string]struct{} { + if info == nil { + return nil + } + req, ok := info.Request.(*dto.OpenAIResponsesRequest) + if !ok || req == nil { + return nil + } + names := make(map[string]struct{}) + for _, tool := range req.GetToolsMap() { + if tool["type"] != "custom" { + continue + } + name, _ := tool["name"].(string) + if name = strings.TrimSpace(name); name != "" { + names[name] = struct{}{} + } + } + return names +} + +func ChatCompletionsResponseToResponsesResponseWithCustomTools(resp *dto.OpenAITextResponse, id string, customToolNames map[string]struct{}) (*dto.OpenAIResponsesResponse, *dto.Usage, error) { + return openaicompat.ChatCompletionsResponseToResponsesResponseWithCustomTools(resp, id, customToolNames) +} + func ResponsesRequestToChatCompletionsRequest(req *dto.OpenAIResponsesRequest) (*dto.GeneralOpenAIRequest, error) { return openaicompat.ResponsesRequestToChatCompletionsRequest(req) } diff --git a/service/openaicompat/chat_to_responses_response.go b/service/openaicompat/chat_to_responses_response.go index bc2cc789..cbe15c0c 100644 --- a/service/openaicompat/chat_to_responses_response.go +++ b/service/openaicompat/chat_to_responses_response.go @@ -3,14 +3,29 @@ package openaicompat import ( "errors" "fmt" + "net/http" "sort" "strings" "time" "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/types" ) +func NewResponsesStreamIncompleteError(reason string) *types.MaxAPIError { + reason = strings.TrimSpace(reason) + if reason == "" { + reason = responsesIncompleteReasonOther + } + return types.NewOpenAIError( + fmt.Errorf("responses stream ended before a terminal event: %s", reason), + types.ErrorCodeBadResponse, + http.StatusBadGateway, + types.ErrOptionWithSkipRetry(), + ) +} + const ( chatFinishReasonLength = "length" chatFinishReasonContentFilter = "content_filter" @@ -22,10 +37,11 @@ type ChatToResponsesStreamEvent struct { } type ChatToResponsesStreamState struct { - ID string - Model string - Created int64 - Usage *dto.Usage + ID string + Model string + Created int64 + Usage *dto.Usage + CustomToolNames map[string]struct{} status string incompleteDetails *dto.IncompleteDetails @@ -50,7 +66,9 @@ type chatToResponsesStreamTool struct { ID string Name string Arguments strings.Builder + Announced bool Done bool + Custom bool } type chatToResponsesOutputRef struct { @@ -59,6 +77,10 @@ type chatToResponsesOutputRef struct { } func ChatCompletionsResponseToResponsesResponse(resp *dto.OpenAITextResponse, id string) (*dto.OpenAIResponsesResponse, *dto.Usage, error) { + return ChatCompletionsResponseToResponsesResponseWithCustomTools(resp, id, nil) +} + +func ChatCompletionsResponseToResponsesResponseWithCustomTools(resp *dto.OpenAITextResponse, id string, customToolNames map[string]struct{}) (*dto.OpenAIResponsesResponse, *dto.Usage, error) { if resp == nil { return nil, nil, errors.New("response is nil") } @@ -109,7 +131,13 @@ func ChatCompletionsResponseToResponsesResponse(resp *dto.OpenAITextResponse, id }) } for i, toolCall := range choice.Message.ParseToolCalls() { - out.Output = append(out.Output, chatToolCallToResponsesOutput(toolCall, id, i, status)) + item := chatToolCallToResponsesOutput(toolCall, id, i, status) + if _, custom := customToolNames[toolCall.Function.Name]; custom && item.Type == responsesOutputTypeFunctionCall { + item.Type = responsesOutputTypeCustomToolCall + item.Input = customToolInputFromArguments(toolCall.Function.Arguments) + item.Arguments = nil + } + out.Output = append(out.Output, item) } return out, usage, nil @@ -226,11 +254,21 @@ func ChatCompletionsStreamChunkToResponsesEvents(chunk *dto.ChatCompletionsStrea return events, nil } +// FlushChatCompletionsStreamToResponsesOutput drains buffered output without +// constructing a terminal response. Hosts can account for successful delivery +// before freezing the terminal usage snapshot. +func FlushChatCompletionsStreamToResponsesOutput(state *ChatToResponsesStreamState) []ChatToResponsesStreamEvent { + if state == nil || state.finalized { + return nil + } + return state.doneDeltaEvents() +} + func FinalizeChatCompletionsStreamToResponses(state *ChatToResponsesStreamState) []ChatToResponsesStreamEvent { if state == nil || state.finalized { return nil } - events := state.doneDeltaEvents() + events := FlushChatCompletionsStreamToResponsesOutput(state) state.finalized = true eventType := responsesEventCompleted if state.status == "incomplete" { @@ -243,6 +281,22 @@ func FinalizeChatCompletionsStreamToResponses(state *ChatToResponsesStreamState) return events } +// MarkChatCompletionsStreamIncomplete records that the upstream transport +// ended before a provider terminal event was observed. Callers may still +// finalize the state to emit response.incomplete, but must not claim a normal +// response.completed result. +func (s *ChatToResponsesStreamState) MarkIncomplete(reason string) { + if s == nil || s.finalized || s.status == "incomplete" { + return + } + reason = strings.TrimSpace(reason) + if reason == "" { + reason = responsesIncompleteReasonOther + } + s.status = "incomplete" + s.incompleteDetails = &dto.IncompleteDetails{Reason: reason} +} + func (s *ChatToResponsesStreamState) UsageText() string { if s == nil { return "" @@ -314,42 +368,49 @@ func (s *ChatToResponsesStreamState) appendToolCallDelta(toolCall dto.ToolCallRe events := make([]ChatToResponsesStreamEvent, 0, 2) if tool == nil { tool = &chatToResponsesStreamTool{ - ChatIndex: chatIndex, - OutputIndex: s.nextIndex("tool", chatIndex), - ID: strings.TrimSpace(toolCall.ID), - Name: strings.TrimSpace(toolCall.Function.Name), + ChatIndex: chatIndex, + ID: strings.TrimSpace(toolCall.ID), + Name: strings.TrimSpace(toolCall.Function.Name), } if tool.ID == "" { tool.ID = fmt.Sprintf("%s_call_%d", s.ID, chatIndex) } s.toolsByIndex[chatIndex] = tool - events = append(events, responsesStreamEvent(responsesEventOutputItemAdded, dto.ResponsesStreamResponse{ - Type: responsesEventOutputItemAdded, - OutputIndex: intPtr(tool.OutputIndex), - ItemID: tool.ID, - Item: &dto.ResponsesOutput{ - Type: responsesOutputTypeFunctionCall, - ID: tool.ID, - Status: "in_progress", - CallId: tool.ID, - Name: tool.Name, - Arguments: []byte(`""`), - }, - })) } if strings.TrimSpace(toolCall.ID) != "" { + if tool.Announced && tool.ID != strings.TrimSpace(toolCall.ID) { + return nil, fmt.Errorf("tool-call index %d changed id after announcement", chatIndex) + } tool.ID = strings.TrimSpace(toolCall.ID) } if strings.TrimSpace(toolCall.Function.Name) != "" { + if tool.Announced && tool.Name != strings.TrimSpace(toolCall.Function.Name) { + return nil, fmt.Errorf("tool-call index %d changed name after announcement", chatIndex) + } tool.Name = strings.TrimSpace(toolCall.Function.Name) } - if toolCall.Function.Arguments != "" { - tool.Arguments.WriteString(toolCall.Function.Arguments) + delta := toolCall.Function.Arguments + tool.Arguments.WriteString(delta) + // Responses output items require a name. Hold the announcement until a + // later delta supplies it, while retaining any argument fragments. + if !tool.Announced && tool.Name != "" { + tool.Announced = true + _, tool.Custom = s.CustomToolNames[tool.Name] + tool.OutputIndex = s.nextIndex("tool", chatIndex) + events = append(events, responsesStreamEvent(responsesEventOutputItemAdded, dto.ResponsesStreamResponse{ + Type: responsesEventOutputItemAdded, + OutputIndex: intPtr(tool.OutputIndex), + ItemID: tool.ID, + Item: s.toolAnnouncement(tool), + })) + delta = tool.Arguments.String() + } + if tool.Announced && !tool.Custom && delta != "" { events = append(events, responsesStreamEvent(responsesEventFunctionArgsDelta, dto.ResponsesStreamResponse{ Type: responsesEventFunctionArgsDelta, OutputIndex: intPtr(tool.OutputIndex), ItemID: tool.ID, - Delta: toolCall.Function.Arguments, + Delta: delta, })) } return events, nil @@ -395,11 +456,44 @@ func (s *ChatToResponsesStreamState) doneDeltaEvents() []ChatToResponsesStreamEv continue } tool.Done = true - events = append(events, responsesStreamEvent(responsesEventFunctionArgsDone, dto.ResponsesStreamResponse{ - Type: responsesEventFunctionArgsDone, - OutputIndex: intPtr(tool.OutputIndex), - ItemID: tool.ID, - })) + if !tool.Announced { + if tool.Name == "" { + continue + } + tool.Announced = true + _, tool.Custom = s.CustomToolNames[tool.Name] + tool.OutputIndex = s.nextIndex("tool", tool.ChatIndex) + events = append(events, responsesStreamEvent(responsesEventOutputItemAdded, dto.ResponsesStreamResponse{ + Type: responsesEventOutputItemAdded, + OutputIndex: intPtr(tool.OutputIndex), + ItemID: tool.ID, + Item: s.toolAnnouncement(tool), + })) + if args := tool.Arguments.String(); !tool.Custom && args != "" { + events = append(events, responsesStreamEvent(responsesEventFunctionArgsDelta, dto.ResponsesStreamResponse{ + Type: responsesEventFunctionArgsDelta, + OutputIndex: intPtr(tool.OutputIndex), + ItemID: tool.ID, + Delta: args, + })) + } + } + if tool.Custom { + input := customToolInputFromArguments(tool.Arguments.String()) + if input != "" { + events = append(events, responsesStreamEvent("response.custom_tool_call_input.delta", dto.ResponsesStreamResponse{ + Type: "response.custom_tool_call_input.delta", OutputIndex: intPtr(tool.OutputIndex), ItemID: tool.ID, Delta: input, + })) + } + events = append(events, responsesStreamEvent("response.custom_tool_call_input.done", dto.ResponsesStreamResponse{ + Type: "response.custom_tool_call_input.done", OutputIndex: intPtr(tool.OutputIndex), ItemID: tool.ID, Input: &input, + })) + } else { + events = append(events, responsesStreamEvent(responsesEventFunctionArgsDone, dto.ResponsesStreamResponse{ + Type: responsesEventFunctionArgsDone, OutputIndex: intPtr(tool.OutputIndex), ItemID: tool.ID, + Arguments: common.GetPointer(tool.Arguments.String()), + })) + } events = append(events, responsesStreamEvent(responsesEventOutputItemDone, dto.ResponsesStreamResponse{ Type: responsesEventOutputItemDone, OutputIndex: intPtr(tool.OutputIndex), @@ -516,6 +610,10 @@ func (s *ChatToResponsesStreamState) reasoningOutput(status string) *dto.Respons } func (s *ChatToResponsesStreamState) toolOutput(tool *chatToResponsesStreamTool, status string) *dto.ResponsesOutput { + if tool.Custom { + return &dto.ResponsesOutput{Type: responsesOutputTypeCustomToolCall, ID: tool.ID, Status: status, + CallId: tool.ID, Name: tool.Name, Input: customToolInputFromArguments(tool.Arguments.String())} + } return &dto.ResponsesOutput{ Type: responsesOutputTypeFunctionCall, ID: tool.ID, @@ -526,6 +624,28 @@ func (s *ChatToResponsesStreamState) toolOutput(tool *chatToResponsesStreamTool, } } +func (s *ChatToResponsesStreamState) toolAnnouncement(tool *chatToResponsesStreamTool) *dto.ResponsesOutput { + item := s.toolOutput(tool, "in_progress") + if tool.Custom { + item.Input = "" + } else { + item.Arguments = []byte(`""`) + } + return item +} + +// Custom tools are represented upstream as functions with one string input. +// Keep malformed or interrupted arguments intact rather than silently losing them. +func customToolInputFromArguments(arguments string) string { + var wrapped struct { + Input *string `json:"input"` + } + if common.UnmarshalJsonStr(arguments, &wrapped) == nil && wrapped.Input != nil { + return *wrapped.Input + } + return arguments +} + func responseOutputStatus(resp *dto.OpenAIResponsesResponse) string { if resp == nil || responseStatusString(resp) != "incomplete" { return "completed" diff --git a/service/openaicompat/custom_tool_finalize_test.go b/service/openaicompat/custom_tool_finalize_test.go new file mode 100644 index 00000000..e17129f2 --- /dev/null +++ b/service/openaicompat/custom_tool_finalize_test.go @@ -0,0 +1,45 @@ +package openaicompat + +import ( + "testing" + + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/stretchr/testify/require" +) + +func TestFlushChatCompletionsStreamToResponsesOutputDefersTerminalUsage(t *testing.T) { + for _, incomplete := range []bool{false, true} { + for _, zero := range []bool{false, true} { + name := map[bool]string{false: "completed", true: "incomplete"}[incomplete] + "/" + map[bool]string{false: "estimated", true: "zero"}[zero] + t.Run(name, func(t *testing.T) { + state := NewChatToResponsesStreamState("resp_flush", "test") + state.CustomToolNames = map[string]struct{}{"apply_patch": {}} + _, err := ChatCompletionsStreamChunkToResponsesEvents(&dto.ChatCompletionsStreamResponse{ + Choices: []dto.ChatCompletionsStreamResponseChoice{{Delta: dto.ChatCompletionsStreamResponseChoiceDelta{ + ToolCalls: []dto.ToolCallResponse{{ID: "call_flush", Function: dto.FunctionResponse{Name: "apply_patch", Arguments: `{"input":"patch body"}`}}}, + }}}, + }, state) + require.NoError(t, err) + if incomplete { + state.MarkIncomplete("upstream_eof") + } + events := FlushChatCompletionsStreamToResponsesOutput(state) + require.Len(t, events, 3) + require.Equal(t, "response.custom_tool_call_input.delta", events[0].Type) + require.Equal(t, "patch body", events[0].Payload.Delta) + require.Empty(t, FlushChatCompletionsStreamToResponsesOutput(state), "buffered output must drain once") + usage := &dto.Usage{PromptTokens: 10, CompletionTokens: 3, TotalTokens: 13} + if zero { + usage = &dto.Usage{} + } + state.Usage = UsageFromChatUsage(usage) + terminal := FinalizeChatCompletionsStreamToResponses(state) + require.Len(t, terminal, 1, "terminal construction must not duplicate flushed input") + require.Equal(t, "response."+map[bool]string{false: "completed", true: "incomplete"}[incomplete], terminal[0].Type) + require.Equal(t, usage.CompletionTokens, terminal[0].Payload.Response.Usage.OutputTokens) + require.Empty(t, FinalizeChatCompletionsStreamToResponses(state)) + require.Empty(t, FlushChatCompletionsStreamToResponsesOutput(state)) + }) + } + } +} diff --git a/service/openaicompat/custom_tool_response_test.go b/service/openaicompat/custom_tool_response_test.go new file mode 100644 index 00000000..64a089dd --- /dev/null +++ b/service/openaicompat/custom_tool_response_test.go @@ -0,0 +1,62 @@ +package openaicompat + +import ( + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func TestResponsesCustomToolPreservesRawInputAndOrdinaryFunctions(t *testing.T) { + for _, input := range []string{"", "patch body", "引号 \" 和反斜线 \\ 及换行\nsecond line"} { + t.Run(input, func(t *testing.T) { + args, err := common.Marshal(map[string]string{"input": input}) + require.NoError(t, err) + tools := []dto.ToolCallResponse{ + {ID: "call_custom", Type: "function", Function: dto.FunctionResponse{Name: "apply_patch", Arguments: string(args)}}, + {ID: "call_function", Type: "function", Function: dto.FunctionResponse{Name: "lookup", Arguments: `{"input":"function argument"}`}}, + } + message := dto.Message{Role: "assistant"} + message.SetToolCalls(tools) + response, _, err := ChatCompletionsResponseToResponsesResponseWithCustomTools(&dto.OpenAITextResponse{Choices: []dto.OpenAITextResponseChoice{{Message: message, FinishReason: "tool_calls"}}}, "resp_test", map[string]struct{}{"apply_patch": {}}) + require.NoError(t, err) + require.Equal(t, input, response.Output[0].Input) + require.Equal(t, "custom_tool_call", response.Output[0].Type) + require.Equal(t, "call_custom", response.Output[0].CallId) + require.Equal(t, "function_call", response.Output[1].Type) + require.Equal(t, `{"input":"function argument"}`, response.Output[1].ArgumentsString()) + payload, err := common.Marshal(response) + require.NoError(t, err) + require.True(t, gjson.GetBytes(payload, "output.0.input").Exists()) + require.False(t, gjson.GetBytes(payload, "output.1.input").Exists()) + state := NewChatToResponsesStreamState("resp_test", "test") + state.CustomToolNames = map[string]struct{}{"apply_patch": {}} + var events []ChatToResponsesStreamEvent + for index, tool := range tools { + tool.SetIndex(index) + chunkEvents, err := ChatCompletionsStreamChunkToResponsesEvents(&dto.ChatCompletionsStreamResponse{Choices: []dto.ChatCompletionsStreamResponseChoice{{Delta: dto.ChatCompletionsStreamResponseChoiceDelta{ToolCalls: []dto.ToolCallResponse{tool}}}}}, state) + require.NoError(t, err) + events = append(events, chunkEvents...) + } + events = append(events, FinalizeChatCompletionsStreamToResponses(state)...) + customDone := 0 + functionDone := 0 + for _, event := range events { + if event.Type == "response.custom_tool_call_input.done" { + customDone++ + require.NotNil(t, event.Payload.Input) + require.Equal(t, input, *event.Payload.Input) + require.Equal(t, "call_custom", event.Payload.ItemID) + } + if event.Type == "response.function_call_arguments.done" { + functionDone++ + require.Equal(t, "call_function", event.Payload.ItemID) + } + } + require.Equal(t, 1, customDone) + require.Equal(t, 1, functionDone) + }) + } +} diff --git a/service/openaicompat/responses_compat_test.go b/service/openaicompat/responses_compat_test.go index 49b507eb..6b0207f8 100644 --- a/service/openaicompat/responses_compat_test.go +++ b/service/openaicompat/responses_compat_test.go @@ -235,7 +235,7 @@ func TestResponsesRequestToChatCompletionsRequestSkipsNonMessageOutputItems(t *t assert.Equal(t, "continue", got.Messages[1].StringContent()) } -func TestResponsesRequestToChatCompletionsRequestRejectsCustomToolCallItems(t *testing.T) { +func TestResponsesRequestToChatCompletionsRequestEncodesCustomToolHistory(t *testing.T) { got, err := ResponsesRequestToChatCompletionsRequest(&dto.OpenAIResponsesRequest{ Model: "gpt-test", Input: mustCompatRawMessage(t, []map[string]any{ @@ -245,12 +245,46 @@ func TestResponsesRequestToChatCompletionsRequestRejectsCustomToolCallItems(t *t "name": "apply_patch", "input": "patch body", }, + { + "type": "custom_tool_call_output", + "call_id": "call_custom", + "output": "ok", + }, }), }) - require.Error(t, err) - require.Nil(t, got) - assert.Contains(t, err.Error(), "custom_tool_call") + require.NoError(t, err) + require.Len(t, got.Messages, 2) + toolCalls := got.Messages[0].ParseToolCalls() + require.Len(t, toolCalls, 1) + assert.Equal(t, "function", toolCalls[0].Type) + assert.Equal(t, "call_custom", toolCalls[0].ID) + assert.Equal(t, "apply_patch", toolCalls[0].Function.Name) + assert.Equal(t, "patch body", gjson.Get(toolCalls[0].Function.Arguments, "input").String()) + assert.Equal(t, dto.Message{Role: "tool", ToolCallId: "call_custom", Content: "ok"}, got.Messages[1]) +} + +func TestResponsesRequestToChatCompletionsRequestEncodesCustomToolDefinition(t *testing.T) { + got, err := ResponsesRequestToChatCompletionsRequest(&dto.OpenAIResponsesRequest{ + Model: "gpt-test", + Tools: mustCompatRawMessage(t, []map[string]any{{ + "type": "custom", + "name": "apply_patch", + "description": "Apply a patch", + }}), + ToolChoice: mustCompatRawMessage(t, map[string]any{"type": "custom", "name": "apply_patch"}), + }) + require.NoError(t, err) + require.Len(t, got.Tools, 1) + assert.Equal(t, "function", got.Tools[0].Type) + assert.Equal(t, "apply_patch", got.Tools[0].Function.Name) + assert.Contains(t, got.Tools[0].Function.Description, "Apply a patch") + params := got.Tools[0].Function.Parameters.(map[string]any) + assert.Equal(t, []string{"input"}, params["required"]) + assert.Equal(t, map[string]any{ + "type": "function", + "function": map[string]any{"name": "apply_patch"}, + }, got.ToolChoice) } func TestResponsesRequestToChatCompletionsRequestRejectsMissingToolCallIDs(t *testing.T) { @@ -299,13 +333,13 @@ func TestResponsesRequestToChatCompletionsRequestRejectsUnsupportedToolTypes(t * Model: "gpt-test", Input: mustCompatRawMessage(t, "hello"), Tools: mustCompatRawMessage(t, []map[string]any{ - {"type": "custom", "name": "apply_patch"}, + {"type": "web_search_preview"}, }), }) require.Error(t, err) require.Nil(t, got) - assert.Contains(t, err.Error(), `tool type "custom"`) + assert.Contains(t, err.Error(), `tool type "web_search_preview"`) } func TestChatCompletionsResponseToResponsesResponsePreservesToolCallsAndIncompleteReason(t *testing.T) { @@ -496,6 +530,68 @@ func TestChatCompletionsStreamToResponsesEventsAggregatesUsage(t *testing.T) { assert.Equal(t, `"{\"q\":\"x\"}"`, string(events[9].Payload.Response.Output[1].Arguments)) } +func TestChatCompletionsStreamToResponsesEventsWaitsForToolName(t *testing.T) { + state := NewChatToResponsesStreamState("resp_unnamed", "gpt-test") + toolIndex := 0 + finishReason := "tool_calls" + + events := mustCompatResponsesEventsFromChatChunk(t, state, &dto.ChatCompletionsStreamResponse{ + Id: "chatcmpl_unnamed", + Choices: []dto.ChatCompletionsStreamResponseChoice{{Delta: dto.ChatCompletionsStreamResponseChoiceDelta{ + ToolCalls: []dto.ToolCallResponse{{ + Index: &toolIndex, + ID: "call_unnamed", + Function: dto.FunctionResponse{ + Arguments: `{"input":"ls"}`, + }, + }}, + }}}, + }) + require.Len(t, events, 1) + assert.Equal(t, responsesEventCreated, events[0].Type) + + nameEvents := mustCompatResponsesEventsFromChatChunk(t, state, &dto.ChatCompletionsStreamResponse{ + Choices: []dto.ChatCompletionsStreamResponseChoice{{Delta: dto.ChatCompletionsStreamResponseChoiceDelta{ + ToolCalls: []dto.ToolCallResponse{{ + Index: &toolIndex, + Function: dto.FunctionResponse{Name: "exec"}, + }}, + }}}, + }) + require.Len(t, nameEvents, 2) + assert.Equal(t, responsesEventOutputItemAdded, nameEvents[0].Type) + require.NotNil(t, nameEvents[0].Payload.Item) + assert.Equal(t, "exec", nameEvents[0].Payload.Item.Name) + assert.Equal(t, responsesEventFunctionArgsDelta, nameEvents[1].Type) + assert.Equal(t, `{"input":"ls"}`, nameEvents[1].Payload.Delta) + + finishEvents := mustCompatResponsesEventsFromChatChunk(t, state, &dto.ChatCompletionsStreamResponse{ + Choices: []dto.ChatCompletionsStreamResponseChoice{{FinishReason: &finishReason}}, + }) + require.Len(t, finishEvents, 2) + assert.Equal(t, responsesEventFunctionArgsDone, finishEvents[0].Type) + assert.Equal(t, responsesEventOutputItemDone, finishEvents[1].Type) +} + +func TestChatCompletionsStreamToResponsesMarksTransportEOFIncomplete(t *testing.T) { + state := NewChatToResponsesStreamState("resp_eof", "gpt-test") + mustCompatResponsesEventsFromChatChunk(t, state, &dto.ChatCompletionsStreamResponse{ + Choices: []dto.ChatCompletionsStreamResponseChoice{{ + Delta: dto.ChatCompletionsStreamResponseChoiceDelta{Content: lo.ToPtr("partial")}, + }}, + }) + + state.MarkIncomplete("upstream_eof") + events := FinalizeChatCompletionsStreamToResponses(state) + require.NotEmpty(t, events) + last := events[len(events)-1] + assert.Equal(t, responsesEventIncomplete, last.Type) + require.NotNil(t, last.Payload.Response) + assert.Equal(t, `"incomplete"`, string(last.Payload.Response.Status)) + require.NotNil(t, last.Payload.Response.IncompleteDetails) + assert.Equal(t, "upstream_eof", last.Payload.Response.IncompleteDetails.Reason) +} + func compatAssistantMessageWithTool(content string, id string, name string, args string) dto.Message { msg := dto.Message{Role: "assistant", Content: content} msg.SetToolCalls([]dto.ToolCallRequest{ diff --git a/service/openaicompat/responses_request_to_chat.go b/service/openaicompat/responses_request_to_chat.go index 19ac8b59..a9031b41 100644 --- a/service/openaicompat/responses_request_to_chat.go +++ b/service/openaicompat/responses_request_to_chat.go @@ -15,6 +15,7 @@ const ( responsesInputTypeFunctionCall = "function_call" responsesInputTypeFunctionCallOutput = "function_call_output" responsesInputTypeCustomToolCall = "custom_tool_call" + responsesInputTypeCustomToolOutput = "custom_tool_call_output" ) func ResponsesRequestToChatCompletionsRequest(req *dto.OpenAIResponsesRequest) (*dto.GeneralOpenAIRequest, error) { @@ -170,11 +171,15 @@ func responsesInputItemToChatMessages(item map[string]any, messages []dto.Messag } return appendToolCallToLastAssistant(messages, toolCall), nil case responsesInputTypeCustomToolCall: - return nil, errors.New("responses to chat conversion does not support custom_tool_call items") - case responsesInputTypeFunctionCallOutput: + toolCall, err := responsesCustomToolCallItemToChatToolCall(item) + if err != nil { + return nil, err + } + return appendToolCallToLastAssistant(messages, toolCall), nil + case responsesInputTypeFunctionCallOutput, responsesInputTypeCustomToolOutput: callID := strings.TrimSpace(common.Interface2String(item["call_id"])) if callID == "" { - return nil, errors.New("function_call_output item is missing call_id") + return nil, fmt.Errorf("%s item is missing call_id", itemType) } content := responseToolOutputToChatContent(item["output"]) return append(messages, dto.Message{Role: "tool", ToolCallId: callID, Content: content}), nil @@ -282,6 +287,35 @@ func responsesFunctionCallItemToChatToolCall(item map[string]any) (dto.ToolCallR }, nil } +// responsesCustomToolCallItemToChatToolCall carries a Responses free-form tool +// input through the Chat Completions function-call shape. The target adapters +// already understand function calls, so the raw text is retained verbatim in +// the single required "input" argument. +func responsesCustomToolCallItemToChatToolCall(item map[string]any) (dto.ToolCallRequest, error) { + name := strings.TrimSpace(common.Interface2String(item["name"])) + if name == "" { + return dto.ToolCallRequest{}, errors.New("custom_tool_call item is missing name") + } + arguments, err := common.Marshal(map[string]string{ + "input": responsesArgumentsString(item["input"]), + }) + if err != nil { + return dto.ToolCallRequest{}, err + } + callID := responsesCallID(item) + if callID == "" { + return dto.ToolCallRequest{}, errors.New("custom_tool_call item is missing call_id") + } + return dto.ToolCallRequest{ + ID: callID, + Type: "function", + Function: dto.FunctionRequest{ + Name: name, + Arguments: string(arguments), + }, + }, nil +} + func appendToolCallToLastAssistant(messages []dto.Message, toolCall dto.ToolCallRequest) []dto.Message { if len(messages) == 0 || messages[len(messages)-1].Role != "assistant" { messages = append(messages, dto.Message{Role: "assistant"}) @@ -303,9 +337,24 @@ func responsesRequestToolsToChat(raw json.RawMessage) ([]dto.ToolCallRequest, er return nil, fmt.Errorf("invalid tools: %w", err) } out := make([]dto.ToolCallRequest, 0, len(tools)) + // Both kinds become functions upstream. Reject ambiguous restoration before + // converting either order, rather than interpreting a function as raw input. + names := make(map[string]string) + for _, tool := range tools { + kind := strings.TrimSpace(common.Interface2String(tool["type"])) + if kind != "function" && kind != "custom" { + continue + } + name := strings.TrimSpace(common.Interface2String(tool["name"])) + if previous, exists := names[name]; exists && (kind == "custom" || previous == "custom") { + return nil, fmt.Errorf("ambiguous custom tool name %q", name) + } + names[name] = kind + } for _, tool := range tools { toolType := strings.TrimSpace(common.Interface2String(tool["type"])) - if toolType == "function" { + switch toolType { + case "function": out = append(out, dto.ToolCallRequest{ Type: "function", Function: dto.FunctionRequest{ @@ -314,13 +363,65 @@ func responsesRequestToolsToChat(raw json.RawMessage) ([]dto.ToolCallRequest, er Parameters: tool["parameters"], }, }) - continue + case "custom": + name := strings.TrimSpace(common.Interface2String(tool["name"])) + if name == "" { + return nil, errors.New("custom tool is missing name") + } + out = append(out, dto.ToolCallRequest{ + Type: "function", + Function: dto.FunctionRequest{ + Name: name, + Description: customToolDescription(tool), + Parameters: customToolInputSchema(), + }, + }) + default: + return nil, fmt.Errorf("responses to chat conversion does not support tool type %q", toolType) } - return nil, fmt.Errorf("responses to chat conversion does not support tool type %q", toolType) } return out, nil } +func customToolInputSchema() map[string]any { + return map[string]any{ + "type": "object", + "properties": map[string]any{ + "input": map[string]any{ + "type": "string", + "description": "Raw input for the tool.", + }, + }, + "required": []string{"input"}, + "additionalProperties": false, + } +} + +func customToolDescription(tool map[string]any) string { + description := strings.TrimSpace(common.Interface2String(tool["description"])) + const hint = `This tool takes freeform text. Put the complete raw text in the "input" argument.` + parts := []string{} + if description != "" { + parts = append(parts, description) + } + parts = append(parts, hint) + if format, ok := tool["format"].(map[string]any); ok && common.Interface2String(format["type"]) == "grammar" { + definition := common.Interface2String(format["definition"]) + if strings.TrimSpace(definition) != "" { + syntax := strings.TrimSpace(common.Interface2String(format["syntax"])) + label := syntax + " grammar" + switch syntax { + case "lark": + label = "Lark grammar" + case "regex": + label = "regular expression" + } + parts = append(parts, "The input must match this "+label+":\n"+definition) + } + } + return strings.Join(parts, "\n\n") +} + func responsesRequestToolChoiceToChat(raw json.RawMessage) (any, error) { if !rawJSONPresent(raw) { return nil, nil @@ -336,7 +437,7 @@ func responsesRequestToolChoiceToChat(raw json.RawMessage) (any, error) { if err := common.Unmarshal(raw, &choice); err != nil { return nil, fmt.Errorf("invalid tool_choice: %w", err) } - if common.Interface2String(choice["type"]) == "function" { + if choiceType := common.Interface2String(choice["type"]); choiceType == "function" || choiceType == "custom" { if name := strings.TrimSpace(common.Interface2String(choice["name"])); name != "" { return map[string]any{"type": "function", "function": map[string]any{"name": name}}, nil } diff --git a/service/openaicompat/responses_revalidation_test.go b/service/openaicompat/responses_revalidation_test.go new file mode 100644 index 00000000..3fb5b98e --- /dev/null +++ b/service/openaicompat/responses_revalidation_test.go @@ -0,0 +1,39 @@ +package openaicompat + +import ( + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/stretchr/testify/require" +) + +func TestResponsesToolNamesRejectAmbiguousCustomConversion(t *testing.T) { + for _, tools := range []string{ + `[{"type":"function","name":"lookup"},{"type":"custom","name":"lookup"}]`, + `[{"type":"custom","name":"lookup"},{"type":"function","name":"lookup"}]`, + `[{"type":"custom","name":"lookup"},{"type":"custom","name":"lookup"}]`, + } { + _, err := responsesRequestToolsToChat([]byte(tools)) + require.Error(t, err, tools) + require.Contains(t, err.Error(), "lookup") + } + tools, err := responsesRequestToolsToChat([]byte(`[{"type":"function","name":"lookup"},{"type":"custom","name":"patch"}]`)) + require.NoError(t, err) + require.Len(t, tools, 2) +} + +func TestResponsesCustomToolPreservesGrammarConstraint(t *testing.T) { + for _, syntax := range []string{"lark", "regex"} { + for _, description := range []string{"", "Apply a patch"} { + definition := "start: \"你好\"\n | /[a-z]+/" + raw, err := common.Marshal([]map[string]any{{"type": "custom", "name": "patch", "description": description, + "format": map[string]any{"type": "grammar", "syntax": syntax, "definition": definition}}}) + require.NoError(t, err) + tools, err := responsesRequestToolsToChat(raw) + require.NoError(t, err) + require.Len(t, tools, 1) + require.Contains(t, tools[0].Function.Description, definition) + require.Contains(t, tools[0].Function.Description, description) + } + } +} diff --git a/service/openaicompat/responses_to_chat.go b/service/openaicompat/responses_to_chat.go index dd9d65d5..22574d83 100644 --- a/service/openaicompat/responses_to_chat.go +++ b/service/openaicompat/responses_to_chat.go @@ -26,6 +26,7 @@ const ( responsesOutputTypeReasoning = "reasoning" responsesIncompleteReasonContentFilter = "content_filter" responsesIncompleteReasonMaxTokens = "max_output_tokens" + responsesIncompleteReasonOther = "other" ) func ResponsesFinishReasonFromStatus(resp *dto.OpenAIResponsesResponse) (string, bool) { diff --git a/service/quota.go b/service/quota.go index 334c0ae1..454b6886 100644 --- a/service/quota.go +++ b/service/quota.go @@ -373,7 +373,8 @@ func postAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, u logContent += fmt.Sprintf(", Tool %q called %d time(s), cost %s", item.Name, item.CallCount, cost.String()) } } - if fallbackQuota, ok := streamFallbackQuota(relayInfo, quota); ok { + if fallbackQuota, ok := streamFallbackQuota(relayInfo, quota); ok && + !knownResponsesPartialUsage(relayInfo, usage, recordSuccess, !tieredOk || tieredResult != nil) { quota = fallbackQuota shouldUpdateUsageStats = true logContent += fmt.Sprintf(", stream ended abnormally (%s), billed pre-consumed quota %d", relayInfo.StreamStatus.EndReason, fallbackQuota) diff --git a/service/responses_partial_quota_test.go b/service/responses_partial_quota_test.go new file mode 100644 index 00000000..71030fcf --- /dev/null +++ b/service/responses_partial_quota_test.go @@ -0,0 +1,83 @@ +package service + +import ( + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/model" + "github.com/MAX-API-Next/MAX-API/pkg/billingexpr" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + relayconstant "github.com/MAX-API-Next/MAX-API/relay/constant" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestResponsesPartialZeroDoesNotBroadenReservationFallbackPolicy(t *testing.T) { + for _, scenario := range []struct { + name string + usage *dto.Usage + successful, chat, expressionError, audio bool + quota int + }{ + {name: "unknown", usage: &dto.Usage{}, quota: 100}, + {name: "unreported_zero", usage: &dto.Usage{BillingUsage: &dto.BillingUsage{Source: dto.BillingUsageSourceOAIChat}}, quota: 100}, + {name: "negative_counts", usage: &dto.Usage{PromptTokens: -1, CompletionTokens: -1, TotalTokens: -2, BillingUsage: dto.NewReportedOpenAIChatBillingUsage(&dto.Usage{PromptTokens: -1, CompletionTokens: -1, TotalTokens: -2})}, quota: 100}, + {name: "negative_canonical_counts", usage: &dto.Usage{BillingUsage: dto.NewReportedOpenAIChatBillingUsage(&dto.Usage{PromptTokens: -1, CompletionTokens: -1, TotalTokens: -2})}, quota: 100}, + {name: "reported_zero", usage: &dto.Usage{BillingUsage: dto.NewReportedOpenAIChatBillingUsage(&dto.Usage{})}, quota: 0}, + {name: "known_free", usage: &dto.Usage{PromptTokens: 10, CompletionTokens: 2, TotalTokens: 12}, quota: 0}, + {name: "success_policy", usage: &dto.Usage{PromptTokens: 10, CompletionTokens: 2, TotalTokens: 12}, successful: true, quota: 100}, + {name: "chat_policy", usage: &dto.Usage{PromptTokens: 10, CompletionTokens: 2, TotalTokens: 12}, chat: true, quota: 100}, + {name: "expression_error", usage: &dto.Usage{BillingUsage: dto.NewReportedOpenAIChatBillingUsage(&dto.Usage{})}, expressionError: true, quota: 100}, + {name: "audio_free", usage: &dto.Usage{PromptTokens: 10, CompletionTokens: 2, TotalTokens: 12, PromptTokensDetails: dto.InputTokenDetails{AudioTokens: 10}, CompletionTokenDetails: dto.OutputTokenDetails{AudioTokens: 2}}, audio: true, quota: 0}, + {name: "audio_expression_error", usage: &dto.Usage{PromptTokens: 10, CompletionTokens: 2, TotalTokens: 12, PromptTokensDetails: dto.InputTokenDetails{AudioTokens: 10}, CompletionTokenDetails: dto.OutputTokenDetails{AudioTokens: 2}}, audio: true, expressionError: true, quota: 100}, + } { + t.Run(scenario.name, func(t *testing.T) { + truncate(t) + seedUser(t, 171, 1000) + seedToken(t, 172, 171, "synthetic-responses-policy-key", 1000) + seedChannel(t, 173) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + requestID := "responses-policy-" + scenario.name + c.Set(common.RequestIdKey, requestID) + info := &relaycommon.RelayInfo{RequestId: requestID, UserId: 171, TokenId: 172, TokenKey: "synthetic-responses-policy-key", ForcePreConsume: true, + IsStream: true, RelayMode: relayconstant.RelayModeResponses, StartTime: time.Now(), OriginModelName: "test-model", UsingGroup: "default", + ChannelMeta: &relaycommon.ChannelMeta{ChannelId: 173}, UserSetting: dto.UserSetting{BillingPreference: "wallet_only", QuotaWarningThreshold: 1}, + StreamStatus: relaycommon.NewStreamStatus(), PriceData: types.PriceData{ModelRatio: 0, CompletionRatio: 1, GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 1}}} + info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonScannerErr, nil) + if scenario.chat { + info.RelayMode = relayconstant.RelayModeChatCompletions + } + if scenario.expressionError { + expr := `unknown_function()` + info.TieredBillingSnapshot = &billingexpr.BillingSnapshot{BillingMode: "tiered_expr", ExprString: expr, ExprHash: billingexpr.ExprHashString(expr), EstimatedQuotaAfterGroup: 100} + } + require.Nil(t, PreConsumeBilling(c, 100, info)) + if scenario.audio { + postAudioConsumeQuota(c, info, scenario.usage, "partial upstream response", false) + } else if scenario.successful { + PostTextConsumeQuota(c, info, scenario.usage, nil) + } else { + PostPartialConsumeQuota(c, info, scenario.usage) + } + info.Billing.Refund(c) + var user model.User + require.NoError(t, model.DB.First(&user, 171).Error) + var token model.Token + require.NoError(t, model.DB.First(&token, 172).Error) + require.EqualValues(t, 1000-scenario.quota, user.Quota) + require.EqualValues(t, 1000-scenario.quota, token.RemainQuota) + var settlement model.BillingSettlement + require.NoError(t, model.DB.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).First(&settlement).Error) + require.EqualValues(t, scenario.quota-100, settlement.FundingDelta) + require.EqualValues(t, scenario.quota-100, settlement.TokenDelta) + require.Equal(t, model.BillingSettlementEffectApplied, settlement.EffectStatus) + require.Equal(t, scenario.quota, getLastLog(t).Quota) + }) + } +} diff --git a/service/text_quota.go b/service/text_quota.go index 88e7aae2..6fd4214b 100644 --- a/service/text_quota.go +++ b/service/text_quota.go @@ -413,6 +413,22 @@ func streamFallbackQuota(relayInfo *relaycommon.RelayInfo, quota int) (int, bool return fallbackQuota, true } +// Only failed Responses generations may bypass the abnormal-stream reservation +// fallback using known usage and a successfully resolved price. Other callers +// retain their existing policy, including unknown usage and expression errors. +func knownResponsesPartialUsage(info *relaycommon.RelayInfo, usage *dto.Usage, recordSuccess, priceResolved bool) bool { + if recordSuccess || !priceResolved || info == nil || !info.IsStream || info.RelayMode != relayconstant.RelayModeResponses || usage == nil { + return false + } + counts := effectiveBillingUsage(usage) + if usage.PromptTokens < 0 || usage.CompletionTokens < 0 || usage.TotalTokens < 0 || + counts.PromptTokens < 0 || counts.CompletionTokens < 0 || counts.TotalTokens < 0 { + return false + } + return counts.PromptTokens > 0 || counts.CompletionTokens > 0 || counts.TotalTokens > 0 || + (usage.BillingUsage != nil && usage.BillingUsage.TokenCountsReported) +} + func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage, extraContent []string) { postTextConsumeQuota(ctx, relayInfo, usage, extraContent, true) } @@ -487,7 +503,8 @@ func postTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us if originUsage == nil { fallbackBaseQuota = 0 } - if fallbackQuota, ok := streamFallbackQuota(relayInfo, fallbackBaseQuota); ok { + if fallbackQuota, ok := streamFallbackQuota(relayInfo, fallbackBaseQuota); ok && + !knownResponsesPartialUsage(relayInfo, originUsage, recordSuccess, !tieredBillingApplied || tieredResult != nil) { summary.Quota = fallbackQuota shouldUpdateUsageStats = true extraContent = append(extraContent, fmt.Sprintf("stream ended abnormally (%s), billed pre-consumed quota %d", relayInfo.StreamStatus.EndReason, fallbackQuota)) diff --git a/web/default/src/features/smart-ops/model-performance.test.tsx b/web/default/src/features/smart-ops/model-performance.test.tsx index b793f2e4..3ff79862 100644 --- a/web/default/src/features/smart-ops/model-performance.test.tsx +++ b/web/default/src/features/smart-ops/model-performance.test.tsx @@ -35,7 +35,6 @@ import { downsampleUptimeSeries, getUptimeAxisDomain, } from '@/features/pricing/components/model-details-chart-utils' -import { ModelPerformanceDetailsContent } from '@/features/pricing/components/model-details-performance' import { UptimeSparkline } from '@/features/pricing/components/model-details-uptime-sparkline' import type { ModelPerformanceData } from './types' @@ -43,9 +42,40 @@ const testEnv = createReactTestEnvironment() let ModelPerformance: typeof import('./model-performance').ModelPerformance let ModelPerformanceCollectionStateAlert: typeof import('./model-performance').ModelPerformanceCollectionStateAlert let ModelPerformanceDetailGrid: typeof import('./model-performance').ModelPerformanceDetailGrid +let ModelPerformanceDetailsContent: typeof import('@/features/pricing/components/model-details-performance').ModelPerformanceDetailsContent +let canvasContextSpy: + | { mock: { calls: unknown[][] }; mockRestore: () => void } + | undefined before(async () => { await testEnv.setup() + // Load the Bun-only test API at runtime; the app typecheck uses Node types. + const testRuntimeModule = 'bun:test' + const { mock, spyOn } = (await import(testRuntimeModule)) as { + mock: { + module: (name: string, factory: () => Record) => void + } + spyOn: ( + target: object, + method: string + ) => { mock: { calls: unknown[][] }; mockRestore: () => void } + } + canvasContextSpy = spyOn( + testEnv.dom.window.HTMLCanvasElement.prototype, + 'getContext' + ) + // These assertions cover page content and aggregation. Canvas rendering is + // outside the JSDOM test seam and belongs in browser integration coverage. + mock.module('@visactor/react-vchart', () => ({ + VChart: () =>
, + })) + // useChartTheme imports the engine separately from the React component. + mock.module('@visactor/vchart', () => ({ + ThemeManager: { setCurrentTheme: () => undefined }, + })) + ModelPerformanceDetailsContent = ( + await import('@/features/pricing/components/model-details-performance') + ).ModelPerformanceDetailsContent const module = await import('./model-performance') ModelPerformance = module.ModelPerformance ModelPerformanceCollectionStateAlert = @@ -53,7 +83,10 @@ before(async () => { ModelPerformanceDetailGrid = module.ModelPerformanceDetailGrid }) -after(() => testEnv.teardown()) +after(() => { + canvasContextSpy?.mockRestore() + testEnv.teardown() +}) const modelPerformanceData: ModelPerformanceData = { storage_mode: 'legacy_log', @@ -546,6 +579,18 @@ describe('ModelPerformance manual queries', () => { assert.ok(detailText.includes('Per-group performance')) assert.ok(detailText.includes('Latency trend (last 24h)')) assert.ok(detailText.includes('Availability (last 24h)')) + await waitFor(() => { + assert.equal( + view.container.querySelectorAll('[data-testid="performance-chart"]') + .length, + 2 + ) + }) + assert.equal( + canvasContextSpy?.mock.calls.length, + 0, + 'page-content tests must not initialize the browser Canvas engine' + ) assert.ok(detailText.includes('default')) assert.ok(detailText.includes(formatThroughput(48.5))) } finally { From e98d56799069484fcfd7bf41f21925ca10163c7b Mon Sep 17 00:00:00 2001 From: CSCITech Date: Sat, 10 Oct 2026 18:14:18 +0800 Subject: [PATCH 02/24] fix: preserve Responses failure and tool identities --- relay/channel/claude/relay-claude.go | 3 + .../channel/gemini/adaptor_responses_test.go | 2 + relay/channel/gemini/relay-gemini.go | 7 + relay/channel/gemini/relay_responses.go | 6 +- relay/channel/gemini/tool_response_id_test.go | 142 ++++++++++++++++++ relay/channel/openai/responses_via_chat.go | 6 +- relay/metadata_funding_settlement_test.go | 13 +- ...sponses_custom_tool_eof_settlement_test.go | 4 +- .../chat_to_responses_response.go | 8 +- .../openaicompat/custom_tool_finalize_test.go | 98 ++++++++++++ 10 files changed, 283 insertions(+), 6 deletions(-) create mode 100644 relay/channel/gemini/tool_response_id_test.go diff --git a/relay/channel/claude/relay-claude.go b/relay/channel/claude/relay-claude.go index 833b2210..55497651 100644 --- a/relay/channel/claude/relay-claude.go +++ b/relay/channel/claude/relay-claude.go @@ -1093,6 +1093,9 @@ func ClaudeStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon. // after the client has gone away. if streamStatus == nil || streamStatus.EndReason == relaycommon.StreamEndReasonEOF { if finalErr := HandleStreamFinalResponse(c, info, claudeInfo); finalErr != nil { + if !claudeInfo.responsesOutputDelivered { + return nil, finalErr + } return claudeInfo.Usage, finalErr } } diff --git a/relay/channel/gemini/adaptor_responses_test.go b/relay/channel/gemini/adaptor_responses_test.go index 5b577ad3..010d392e 100644 --- a/relay/channel/gemini/adaptor_responses_test.go +++ b/relay/channel/gemini/adaptor_responses_test.go @@ -104,6 +104,8 @@ func TestConvertOpenAIResponsesRequestToGeminiPreservesCustomToolCalls(t *testin require.Len(t, got.Contents[1].Parts, 2) assert.Equal(t, "ok", got.Contents[1].Parts[0].FunctionResponse.Response["content"]) assert.Equal(t, "legacy custom output", got.Contents[1].Parts[1].FunctionResponse.Response["content"]) + assert.JSONEq(t, `"call_custom"`, string(got.Contents[1].Parts[0].FunctionResponse.ID)) + assert.JSONEq(t, `"call_custom"`, string(got.Contents[1].Parts[1].FunctionResponse.ID)) assert.Equal(t, "user", got.Contents[2].Role) assert.Equal(t, "next turn", got.Contents[2].Parts[0].Text) } diff --git a/relay/channel/gemini/relay-gemini.go b/relay/channel/gemini/relay-gemini.go index 36d42d7b..509b5f16 100644 --- a/relay/channel/gemini/relay-gemini.go +++ b/relay/channel/gemini/relay-gemini.go @@ -385,6 +385,13 @@ func CovertOpenAI2Gemini(c *gin.Context, textRequest dto.GeneralOpenAIRequest, i Name: name, Response: contentMap, } + if message.ToolCallId != "" { + id, err := common.Marshal(message.ToolCallId) + if err != nil { + return nil, fmt.Errorf("encode Gemini function response id: %w", err) + } + functionResp.ID = id + } *parts = append(*parts, dto.GeminiPart{ FunctionResponse: functionResp, diff --git a/relay/channel/gemini/relay_responses.go b/relay/channel/gemini/relay_responses.go index 9342c07b..22cf0056 100644 --- a/relay/channel/gemini/relay_responses.go +++ b/relay/channel/gemini/relay_responses.go @@ -281,7 +281,11 @@ func GeminiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, r state.MarkIncomplete(reason) if info.StreamStatus == nil || info.StreamStatus.EndReason == relaycommon.StreamEndReasonEOF { if !finalizeStream() { - return usage, types.NewError(streamErr, streamErr.GetErrorCode(), types.ErrOptionWithSkipRetry()) + finalErr := types.NewError(streamErr, streamErr.GetErrorCode(), types.ErrOptionWithSkipRetry()) + if !outputDelivered { + return nil, finalErr + } + return usage, finalErr } } if !outputDelivered { diff --git a/relay/channel/gemini/tool_response_id_test.go b/relay/channel/gemini/tool_response_id_test.go new file mode 100644 index 00000000..c75f880e --- /dev/null +++ b/relay/channel/gemini/tool_response_id_test.go @@ -0,0 +1,142 @@ +package gemini + +import ( + "fmt" + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/dto" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/stretchr/testify/require" +) + +func TestGeminiToolResponseIDsSurviveOutboundJSON(t *testing.T) { + for _, route := range []string{"chat", "responses_function", "responses_custom"} { + for _, stream := range []bool{false, true} { + for _, prefix := range []string{"call", "调用\"\\\n"} { + t.Run(fmt.Sprintf("%s/stream_%t/prefix_%q", route, stream, prefix), func(t *testing.T) { + ids := []string{prefix + "_object", prefix + "_array", prefix + "_text"} + outputs := []string{`{"value":0}`, `["result",1]`, "plain result"} + order := []int{2, 0, 1} + info := &relaycommon.RelayInfo{IsStream: stream, OriginModelName: "gemini-test", + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gemini-test"}} + var converted any + var err error + if route == "chat" { + calls := make([]dto.ToolCallRequest, 0, len(ids)) + for i, id := range ids { + calls = append(calls, dto.ToolCallRequest{ID: id, Type: "function", + Function: dto.FunctionRequest{Name: "lookup", Arguments: fmt.Sprintf(`{"index":%d}`, i)}}) + } + assistant := dto.Message{Role: "assistant"} + assistant.SetToolCalls(calls) + messages := []dto.Message{{Role: "user", Content: "run tools"}, assistant} + for _, i := range order { + messages = append(messages, dto.Message{Role: "tool", ToolCallId: ids[i], Content: outputs[i]}) + } + converted, err = (&Adaptor{}).ConvertOpenAIRequest(nil, info, &dto.GeneralOpenAIRequest{ + Model: "gemini-test", Stream: common.GetPointer(stream), Messages: messages, + }) + } else { + input := []map[string]any{{"role": "user", "content": "run tools"}} + custom := route == "responses_custom" + for i, id := range ids { + call := map[string]any{"type": "function_call", "call_id": id, "name": "lookup", + "arguments": fmt.Sprintf(`{"index":%d}`, i)} + if custom { + call = map[string]any{"type": "custom_tool_call", "call_id": id, "name": "lookup", "input": fmt.Sprintf("input %d", i)} + } + input = append(input, call) + } + for _, i := range order { + outputType := "function_call_output" + if custom { + outputType = "custom_tool_call_output" + } + input = append(input, map[string]any{"type": outputType, "call_id": ids[i], "output": outputs[i]}) + } + tool := map[string]any{"type": "function", "name": "lookup", "parameters": map[string]any{"type": "object"}} + if custom { + tool = map[string]any{"type": "custom", "name": "lookup"} + } + converted, err = (&Adaptor{}).ConvertOpenAIResponsesRequest(nil, info, dto.OpenAIResponsesRequest{ + Model: "gemini-test", Stream: common.GetPointer(stream), + Input: mustGeminiRawMessage(t, input), Tools: mustGeminiRawMessage(t, []map[string]any{tool}), + }) + } + require.NoError(t, err) + wire, err := common.Marshal(converted) + require.NoError(t, err) + var request struct { + Contents []struct { + Role string + Parts []struct { + FunctionCall *struct{ ID, Name string } + FunctionResponse *struct { + ID, Name string + Response map[string]any + } + } + } + } + require.NoError(t, common.Unmarshal(wire, &request)) + callIDs, resultIDs := []string{}, []string{} + results := []map[string]any{} + for _, content := range request.Contents { + for _, part := range content.Parts { + if part.FunctionCall != nil { + require.Equal(t, "model", content.Role) + require.Equal(t, "lookup", part.FunctionCall.Name) + callIDs = append(callIDs, part.FunctionCall.ID) + } + if part.FunctionResponse != nil { + require.Equal(t, "user", content.Role) + require.Equal(t, "lookup", part.FunctionResponse.Name) + resultIDs = append(resultIDs, part.FunctionResponse.ID) + results = append(results, part.FunctionResponse.Response) + } + } + } + require.Equal(t, ids, callIDs) + require.Equal(t, []string{ids[2], ids[0], ids[1]}, resultIDs, "results must keep their own call ID in completion order") + require.Equal(t, []map[string]any{{"content": "plain result"}, {"value": float64(0)}, {"result": []any{"result", float64(1)}}}, results) + }) + } + } + } +} + +func TestGeminiLegacyFunctionResponseIDPresence(t *testing.T) { + for _, role := range []string{"tool", "function"} { + for _, id := range []string{"", "0"} { + t.Run(role+"/id_"+id, func(t *testing.T) { + info := &relaycommon.RelayInfo{OriginModelName: "gemini-test", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gemini-test"}} + converted, err := (&Adaptor{}).ConvertOpenAIRequest(nil, info, &dto.GeneralOpenAIRequest{ + Model: "gemini-test", Messages: []dto.Message{ + {Role: "user", Content: "legacy conversation"}, + {Role: role, Name: common.GetPointer("lookup"), ToolCallId: id, Content: "legacy result"}, + }, + }) + require.NoError(t, err) + wire, err := common.Marshal(converted) + require.NoError(t, err) + var request dto.GeminiChatRequest + require.NoError(t, common.Unmarshal(wire, &request)) + result := request.Contents[0].Parts[1].FunctionResponse + require.NotNil(t, result) + require.Equal(t, "lookup", result.Name) + require.Equal(t, "legacy result", result.Response["content"]) + if id == "" { + // omitempty must omit ID rather than emitting null or an invented value. + var fields map[string]any + payload, err := common.Marshal(result) + require.NoError(t, err) + require.NoError(t, common.Unmarshal(payload, &fields)) + require.NotContains(t, fields, "id") + } else { + require.JSONEq(t, `"0"`, string(result.ID)) + } + }) + } + } +} diff --git a/relay/channel/openai/responses_via_chat.go b/relay/channel/openai/responses_via_chat.go index 0b1089bc..16c1c521 100644 --- a/relay/channel/openai/responses_via_chat.go +++ b/relay/channel/openai/responses_via_chat.go @@ -204,7 +204,11 @@ func OaiChatToResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo state.MarkIncomplete(reason) if info.StreamStatus == nil || info.StreamStatus.EndReason == relaycommon.StreamEndReasonEOF { if !finalizeStream() { - return usage, types.NewError(streamErr, streamErr.GetErrorCode(), types.ErrOptionWithSkipRetry()) + finalErr := types.NewError(streamErr, streamErr.GetErrorCode(), types.ErrOptionWithSkipRetry()) + if !outputDelivered { + return nil, finalErr + } + return usage, finalErr } } if !outputDelivered { diff --git a/relay/metadata_funding_settlement_test.go b/relay/metadata_funding_settlement_test.go index 6b98c717..52009e19 100644 --- a/relay/metadata_funding_settlement_test.go +++ b/relay/metadata_funding_settlement_test.go @@ -94,7 +94,10 @@ func TestMetadataOnlyFailureRetainsOriginalFunding(t *testing.T) { for _, route := range []string{"responses_native", "chat_native", "chat_to_responses", "responses_to_chat_custom", "responses_to_chat_global", "claude_to_responses", "gemini_to_responses"} { for _, funding := range streamFundingCases() { for _, usage := range []string{"known", "zero", "missing"} { - for _, ending := range []string{"eof", "transport", "malformed", "cancel", "timeout"} { + for _, ending := range []string{"eof", "transport", "malformed", "cancel", "timeout", "eof_terminal_write_failure"} { + if ending == "eof_terminal_write_failure" && route != "chat_to_responses" && route != "claude_to_responses" && route != "gemini_to_responses" { + continue + } // Native Chat's normal EOF keeps its legacy empty-completion // pricing policy. Its abnormal endings belong in this matrix. if route == "chat_native" && ending == "eof" { @@ -150,6 +153,10 @@ func TestMetadataOnlyFailureRetainsOriginalFunding(t *testing.T) { info.UserSetting.BillingPreference = "subscription_only" } var cancelWriter *metadataCancelWriter + if ending == "eof_terminal_write_failure" { + failureContext, _ := gin.CreateTestContext(&customToolEOFWriter{ResponseRecorder: recorder, ending: "terminal_write_failure"}) + c.Writer = failureContext.Writer + } if ending == "cancel" { ctx, cancel := context.WithCancel(c.Request.Context()) t.Cleanup(cancel) @@ -176,6 +183,10 @@ func TestMetadataOnlyFailureRetainsOriginalFunding(t *testing.T) { require.Equal(t, types.ErrorCodeChannelResponseTimeExceeded, apiErr.GetErrorCode()) require.Equal(t, relaycommon.StreamEndReasonTimeout, info.StreamStatus.EndReason) } + if ending == "eof_terminal_write_failure" { + require.Equal(t, relaycommon.StreamEndReasonEOF, info.StreamStatus.EndReason, "failure occurs after the scanner records EOF") + require.Equal(t, types.ErrorCodeBadResponse, apiErr.GetErrorCode()) + } if cancelWriter != nil { require.True(t, cancelWriter.triggered.Load()) require.ErrorIs(t, c.Request.Context().Err(), context.Canceled) diff --git a/relay/responses_custom_tool_eof_settlement_test.go b/relay/responses_custom_tool_eof_settlement_test.go index b6ee4410..d1377385 100644 --- a/relay/responses_custom_tool_eof_settlement_test.go +++ b/relay/responses_custom_tool_eof_settlement_test.go @@ -100,14 +100,14 @@ func TestResponsesCustomToolEOFFinalizationUsesDurableLedger(t *testing.T) { require.NoError(t, db.Create(&model.UserSubscription{Id: 905, UserId: 901, PlanId: 904, AmountTotal: 1000, Status: "active", StartTime: time.Now().Add(-time.Hour).Unix(), EndTime: time.Now().Add(time.Hour).Unix()}).Error) } - prompt, completion := 10, 0 + prompt := 10 deliveredInput := strings.Join(inputs, "") if ending == "first_input_write_failure" || ending == "no_delivery" { deliveredInput = "" } else if ending == "second_input_write_failure" || ending == "finish_second_input_write_failure" || ending == "input_flush_failure" { deliveredInput = inputs[0] } - completion = service.EstimateTokenByModel("gpt-test", deliveredInput) + completion := service.EstimateTokenByModel("gpt-test", deliveredInput) if usageKind == "known" { prompt, completion = 4, 2 } else if usageKind == "zero" { diff --git a/service/openaicompat/chat_to_responses_response.go b/service/openaicompat/chat_to_responses_response.go index cbe15c0c..d5749ca0 100644 --- a/service/openaicompat/chat_to_responses_response.go +++ b/service/openaicompat/chat_to_responses_response.go @@ -418,6 +418,12 @@ func (s *ChatToResponsesStreamState) appendToolCallDelta(toolCall dto.ToolCallRe func (s *ChatToResponsesStreamState) doneDeltaEvents() []ChatToResponsesStreamEvent { events := make([]ChatToResponsesStreamEvent, 0) + tools := s.sortedTools() + for _, tool := range tools { + if tool.Name == "" { + s.MarkIncomplete("missing_tool_name") + } + } status := s.outputStatus() if s.textStarted && !s.textDone { s.textDone = true @@ -451,7 +457,7 @@ func (s *ChatToResponsesStreamState) doneDeltaEvents() []ChatToResponsesStreamEv Item: s.reasoningOutput(status), })) } - for _, tool := range s.sortedTools() { + for _, tool := range tools { if tool.Done { continue } diff --git a/service/openaicompat/custom_tool_finalize_test.go b/service/openaicompat/custom_tool_finalize_test.go index e17129f2..36b7fec3 100644 --- a/service/openaicompat/custom_tool_finalize_test.go +++ b/service/openaicompat/custom_tool_finalize_test.go @@ -1,8 +1,10 @@ package openaicompat import ( + "fmt" "testing" + "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/dto" "github.com/stretchr/testify/require" ) @@ -43,3 +45,99 @@ func TestFlushChatCompletionsStreamToResponsesOutputDefersTerminalUsage(t *testi } } } + +func TestChatToResponsesUnnamedToolCallCannotComplete(t *testing.T) { + for _, finish := range []string{"", "tool_calls", "length"} { + for _, flushFirst := range []bool{false, true} { + for _, lateName := range []bool{false, true} { + for _, mixedOutput := range []bool{false, true} { + t.Run(fmt.Sprintf("finish_%q/flush_%t/late_name_%t/mixed_%t", finish, flushFirst, lateName, mixedOutput), func(t *testing.T) { + state := NewChatToResponsesStreamState("resp_name", "test") + state.CustomToolNames = map[string]struct{}{"named_custom": {}} + delta := dto.ChatCompletionsStreamResponseChoiceDelta{ToolCalls: []dto.ToolCallResponse{ + {Index: common.GetPointer(0), ID: "call_unnamed", Function: dto.FunctionResponse{Arguments: `{"q":"buffered"}`}}, + }} + if mixedOutput { + delta.SetContentString("kept text") + delta.ToolCalls = append(delta.ToolCalls, dto.ToolCallResponse{Index: common.GetPointer(1), ID: "call_named", Function: dto.FunctionResponse{Name: "named_custom", Arguments: `{"input":"kept input"}`}}) + } + _, err := ChatCompletionsStreamChunkToResponsesEvents(&dto.ChatCompletionsStreamResponse{ + Choices: []dto.ChatCompletionsStreamResponseChoice{{Delta: delta}}, + }, state) + require.NoError(t, err) + last := dto.ChatCompletionsStreamResponseChoice{} + if lateName { + last.Delta.ToolCalls = []dto.ToolCallResponse{{Index: common.GetPointer(0), Function: dto.FunctionResponse{Name: "lookup"}}} + } + if finish != "" { + last.FinishReason = common.GetPointer(finish) + } + _, err = ChatCompletionsStreamChunkToResponsesEvents(&dto.ChatCompletionsStreamResponse{Choices: []dto.ChatCompletionsStreamResponseChoice{last}}, state) + require.NoError(t, err) + if finish != "" { + _, err = ChatCompletionsStreamChunkToResponsesEvents(&dto.ChatCompletionsStreamResponse{ + Choices: []dto.ChatCompletionsStreamResponseChoice{{FinishReason: common.GetPointer(finish)}}, + }, state) + require.NoError(t, err, "repeated finish cannot hide an unnamed call") + } + if flushFirst { + FlushChatCompletionsStreamToResponsesOutput(state) + } + state.Usage = UsageFromChatUsage(&dto.Usage{PromptTokens: 10, CompletionTokens: 2, TotalTokens: 12}) + terminal := FinalizeChatCompletionsStreamToResponses(state) + require.NotEmpty(t, terminal) + final := terminal[len(terminal)-1] + expectedStatus := "completed" + if !lateName || finish == "length" { + expectedStatus = "incomplete" + reason := "missing_tool_name" + if finish == "length" { + reason = responsesIncompleteReasonMaxTokens + } + require.NotNil(t, final.Payload.Response.IncompleteDetails) + require.Equal(t, reason, final.Payload.Response.IncompleteDetails.Reason) + } + require.Equal(t, "response."+expectedStatus, final.Type) + payload, err := common.Marshal(final.Payload) + require.NoError(t, err) + var wire struct { + Response struct { + Status string + Output []struct { + Type, Name, Status, Input, Arguments string + CallID string `json:"call_id"` + } + } + } + require.NoError(t, common.Unmarshal(payload, &wire)) + require.Equal(t, expectedStatus, wire.Response.Status) + expectedCount := 0 + if mixedOutput { + expectedCount = 2 + } + if lateName { + expectedCount++ + } + require.Len(t, wire.Response.Output, expectedCount, "valid output remains in the terminal") + for _, item := range wire.Response.Output { + require.Equal(t, expectedStatus, item.Status) + if item.Type == "function_call" { + require.Equal(t, "lookup", item.Name) + require.Equal(t, "call_unnamed", item.CallID) + require.JSONEq(t, `{"q":"buffered"}`, item.Arguments) + } else if item.Type == "custom_tool_call" { + require.Equal(t, "named_custom", item.Name) + require.Equal(t, "call_named", item.CallID) + require.Equal(t, "kept input", item.Input) + } + } + require.Equal(t, 10, final.Payload.Response.Usage.InputTokens) + require.Equal(t, 2, final.Payload.Response.Usage.OutputTokens) + require.Empty(t, FlushChatCompletionsStreamToResponsesOutput(state)) + require.Empty(t, FinalizeChatCompletionsStreamToResponses(state), "terminal output is emitted once") + }) + } + } + } + } +} From ba64fd71540c7c942bee588e843c30af6909c741 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Sat, 10 Oct 2026 18:56:22 +0800 Subject: [PATCH 03/24] fix: scope subscription reset guards to proven period evidence --- model/subscription.go | 11 ++- model/subscription_reset_billing_test.go | 105 +++++++++++++++++++++++ 2 files changed, 115 insertions(+), 1 deletion(-) diff --git a/model/subscription.go b/model/subscription.go index 6135c3d3..7e3ed20d 100644 --- a/model/subscription.go +++ b/model/subscription.go @@ -1099,9 +1099,18 @@ func resetUserSubscriptionTx(tx *gorm.DB, sub *UserSubscription, plan *Subscript // Only a recorded funding application makes such a reservation safe to reset; // effect-only replay does not debit the subscription again. func validateSubscriptionResetBillingTx(tx *gorm.DB, sub UserSubscription) error { + // Only retained, matching reservation evidence can prove that an intent + // belongs to an earlier period. Missing tombstones or ambiguous identities + // still block: neither cleanup nor a review disposition proves application. + earlierPeriod := tx.Model(&SubscriptionPreConsumeRecord{}).Select("1"). + Where("subscription_pre_consume_records.request_id = billing_settlements.subscription_pre_consume_request_id"). + Where("subscription_pre_consume_records.user_id = billing_settlements.user_id"). + Where("user_subscription_id = ? AND user_id = ? AND subscription_last_reset_time < ? AND status = ?", + sub.Id, sub.UserId, sub.LastResetTime, "consumed") var unresolved BillingSettlement query := tx.Select("id").Where("source = ? AND subscription_id = ?", BillingSettlementSourceSubscription, sub.Id). - Where("status NOT IN ?", []string{BillingSettlementStatusApplied, ""}).Limit(1).Find(&unresolved) + Where("status NOT IN ?", []string{BillingSettlementStatusApplied, ""}). + Where("NOT EXISTS (?)", earlierPeriod).Limit(1).Find(&unresolved) if query.Error != nil { return query.Error } diff --git a/model/subscription_reset_billing_test.go b/model/subscription_reset_billing_test.go index 4aac189a..d0e4d4b8 100644 --- a/model/subscription_reset_billing_test.go +++ b/model/subscription_reset_billing_test.go @@ -33,6 +33,111 @@ func resetBillingFinalize(sub UserSubscription, requestID string) BillingSettlem SubscriptionPreConsumeRequestID: requestID} } +func TestSubscriptionResetScopesProvenEarlierPeriod(t *testing.T) { + for _, state := range []string{BillingSettlementStatusManual, BillingSettlementStatusPending, "outcome_unknown"} { + for _, evidence := range []string{"earlier", "missing", "wrong_user", "future", "current"} { + for _, resetPlan := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/%s/plan_%t", state, evidence, resetPlan), func(t *testing.T) { + plan, sub, requestID := seedResetBillingReservation(t) + require.NoError(t, DB.Transaction(func(tx *gorm.DB) error { + var locked UserSubscription + if err := withRowLock(tx).First(&locked, sub.Id).Error; err != nil { + return err + } + return resetUserSubscriptionTx(tx, &locked, &plan, sub.LastResetTime+1, true) + })) + input := resetBillingFinalize(sub, requestID) + _, _, err := ApplyBillingSettlementOnce(input) + require.ErrorIs(t, err, ErrSubscriptionSettlementPeriodChanged) + require.NoError(t, DB.Model(&BillingSettlement{}).Where("operation_key = ?", input.OperationKey).Update("status", state).Error) + var current UserSubscription + require.NoError(t, DB.First(¤t, sub.Id).Error) + switch evidence { + case "missing": + require.NoError(t, DB.Where("request_id = ?", requestID).Delete(&SubscriptionPreConsumeRecord{}).Error) + case "wrong_user": + require.NoError(t, DB.Model(&SubscriptionPreConsumeRecord{}).Where("request_id = ?", requestID).Update("user_id", 9999).Error) + case "future", "current": + anchor := current.LastResetTime + if evidence == "future" { + anchor++ + } + require.NoError(t, DB.Model(&SubscriptionPreConsumeRecord{}).Where("request_id = ?", requestID).Update("subscription_last_reset_time", anchor).Error) + } + // New-period usage is independently reserved and durably finalized. + currentID := requestID + "-current" + _, err = PreConsumeTokenAndUserSubscription(currentID, sub.UserId, 4502, "synthetic-reset-billing", "gpt-test", 0, 200) + require.NoError(t, err) + currentInput := resetBillingFinalize(current, currentID) + currentInput.FundingDelta, currentInput.TokenDelta = 0, 0 + _, _, err = ApplyBillingSettlementOnce(currentInput) + require.NoError(t, err) + if resetPlan { + _, err = AdminResetPlanSubscriptions(plan.Id, false) + } else { + _, err = AdminResetUserSubscriptionsByPlan(sub.UserId, plan.Id, false) + } + wantUsed := int64(200) + if evidence == "earlier" { + require.NoError(t, err, "a proven earlier-period intent cannot mutate the current period") + wantUsed = 0 + for range 2 { + _, _, replayErr := ApplyBillingSettlementOnce(input) + require.Error(t, replayErr, "old-period replay remains unapplied") + } + } else { + require.ErrorIs(t, err, ErrSubscriptionResetPendingBilling) + } + var got UserSubscription + require.NoError(t, DB.First(&got, sub.Id).Error) + require.Equal(t, wantUsed, got.AmountUsed) + require.Equal(t, current.LastResetTime, got.LastResetTime) + var token Token + require.NoError(t, DB.First(&token, 4502).Error) + require.EqualValues(t, 700, token.RemainQuota) + var old BillingSettlement + require.NoError(t, DB.Where("operation_key = ?", input.OperationKey).First(&old).Error) + require.NotEqual(t, BillingSettlementStatusApplied, old.Status) + require.Zero(t, old.AppliedFundingDelta) + require.Zero(t, old.AppliedTokenDelta) + var count int64 + require.NoError(t, DB.Model(&BillingSettlement{}).Count(&count).Error) + require.EqualValues(t, 2, count) + }) + } + } + } +} + +func TestSubscriptionResetNeverKeepsCurrentZeroAnchorGuard(t *testing.T) { + for _, state := range []string{BillingSettlementStatusManual, BillingSettlementStatusPending, "outcome_unknown"} { + for _, advance := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/advance_%t", state, advance), func(t *testing.T) { + plan, sub, requestID := seedResetBillingReservation(t) + require.NoError(t, DB.Model(&SubscriptionPlan{}).Where("id = ?", plan.Id).Update("quota_reset_period", SubscriptionResetNever).Error) + InvalidateSubscriptionPlanCache(plan.Id) + t.Cleanup(func() { InvalidateSubscriptionPlanCache(plan.Id) }) + require.NoError(t, DB.Model(&UserSubscription{}).Where("id = ?", sub.Id).Updates(map[string]any{"last_reset_time": 0, "next_reset_time": 0}).Error) + require.NoError(t, DB.Model(&SubscriptionPreConsumeRecord{}).Where("request_id = ?", requestID).Update("subscription_last_reset_time", 0).Error) + input := resetBillingFinalize(sub, requestID) + input.FundingDelta, input.TokenDelta = 2000, 2000 + _, _, err := ApplyBillingSettlementOnce(input) + require.ErrorIs(t, err, ErrSubscriptionQuotaInsufficient) + require.NoError(t, DB.Model(&BillingSettlement{}).Where("operation_key = ?", input.OperationKey).Update("status", state).Error) + _, err = AdminResetUserSubscriptionsByPlan(sub.UserId, plan.Id, advance) + require.ErrorIs(t, err, ErrSubscriptionResetPendingBilling) + var got UserSubscription + require.NoError(t, DB.First(&got, sub.Id).Error) + require.EqualValues(t, 100, got.AmountUsed) + require.Zero(t, got.LastResetTime) + var token Token + require.NoError(t, DB.First(&token, 4502).Error) + require.EqualValues(t, 900, token.RemainQuota) + }) + } + } +} + func TestSubscriptionResetValidatesFundingEvidence(t *testing.T) { for _, tc := range []struct { name string From 1c2fbbf2b6c82199572bd89a1a9132367e450b1e Mon Sep 17 00:00:00 2001 From: CSCITech Date: Sat, 10 Oct 2026 23:25:19 +0800 Subject: [PATCH 04/24] fix: reject failed native Responses terminal aliases --- relay/channel/openai/relay_responses.go | 4 + .../openai/responses_terminal_repair_test.go | 50 ++++++++ relay/metadata_funding_settlement_test.go | 115 ++++++++++++++++++ 3 files changed, 169 insertions(+) diff --git a/relay/channel/openai/relay_responses.go b/relay/channel/openai/relay_responses.go index 379e7c15..1e4230c2 100644 --- a/relay/channel/openai/relay_responses.go +++ b/relay/channel/openai/relay_responses.go @@ -359,6 +359,10 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp case "response.incomplete", "response.done": terminalEventSeen = true if streamResponse.Response != nil { + if terminalErr := responsesTerminalError(streamResponse.Response, http.StatusInternalServerError, streamForwarded); terminalErr != nil { + terminalEventErr = terminalErr + break + } observeResponsesOutputs(info, streamResponse.Response.Output) } case "response.output_text.delta": diff --git a/relay/channel/openai/responses_terminal_repair_test.go b/relay/channel/openai/responses_terminal_repair_test.go index 168b96bb..bb789268 100644 --- a/relay/channel/openai/responses_terminal_repair_test.go +++ b/relay/channel/openai/responses_terminal_repair_test.go @@ -193,3 +193,53 @@ func TestResponsesToChatLegalTerminalsPreserveProviderFailure(t *testing.T) { }) } } + +func TestNativeResponsesLegalTerminalsPreserveReportedUsage(t *testing.T) { + oldEmptyRetry := common.EmptyCompletionRetryEnabled + common.EmptyCompletionRetryEnabled = false + t.Cleanup(func() { common.EmptyCompletionRetryEnabled = oldEmptyRetry }) + for _, event := range []string{"response.completed", "response.done", "response.incomplete"} { + for _, zero := range []bool{false, true} { + for _, delivered := range []bool{false, true} { + name := event + if zero { + name += "/zero" + } else { + name += "/known" + } + if delivered { + name += "/output" + } else { + name += "/metadata" + } + t.Run(name, func(t *testing.T) { + c, info := newResponsesUsageTestContext() + status := "completed" + if event == "response.incomplete" { + status = "incomplete" + } + prompt, completion := 4, 2 + if zero { + prompt, completion = 0, 0 + } + terminal, err := common.Marshal(map[string]any{"type": event, "response": map[string]any{ + "id": "resp_native_legal", "status": status, "output": []any{}, + "usage": map[string]int{"input_tokens": prompt, "output_tokens": completion, "total_tokens": prompt + completion}, + }}) + require.NoError(t, err) + body := "data: " + string(terminal) + "\n\n" + if delivered { + body = "data: {\"type\":\"response.output_text.delta\",\"delta\":\"hello\"}\n\n" + body + } + usage, apiErr := OaiResponsesStreamHandler(c, info, &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(body))}) + require.Nil(t, apiErr) + require.NotNil(t, usage) + require.Equal(t, prompt, usage.PromptTokens) + require.Equal(t, completion, usage.CompletionTokens) + require.Equal(t, prompt+completion, usage.TotalTokens) + require.Equal(t, delivered, info.HasRecordedChannelFirstResult()) + }) + } + } + } +} diff --git a/relay/metadata_funding_settlement_test.go b/relay/metadata_funding_settlement_test.go index 52009e19..bec67668 100644 --- a/relay/metadata_funding_settlement_test.go +++ b/relay/metadata_funding_settlement_test.go @@ -255,3 +255,118 @@ func TestMetadataOnlyCompletedResponsesPreservesReportedUsage(t *testing.T) { }) } } + +// A terminal event name cannot override a failed response status. Exercise the +// native adapter and original funding ledger, including durable failure replay. +func TestNativeResponsesTerminalFailureRetainsOriginalFunding(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitHttpClient() + service.InitTokenEncoders() + oldEmptyRetry := common.EmptyCompletionRetryEnabled + common.EmptyCompletionRetryEnabled = false + t.Cleanup(func() { common.EmptyCompletionRetryEnabled = oldEmptyRetry }) + const text = "partial terminal output" + for _, event := range []string{"response.completed", "response.done", "response.incomplete"} { + for _, status := range []string{"failed", "cancelled", "canceled"} { + for _, funding := range streamFundingCases() { + for _, usageKind := range []string{"known", "zero", "missing"} { + for _, delivered := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/%s/%s/%s/delivered_%t", event, status, funding.name, usageKind, delivered), func(t *testing.T) { + db := setupStreamFundingLedger(t, funding) + _, frames := metadataFundingFrames(t, "responses_native", usageKind) + if delivered { + data, err := common.Marshal(map[string]any{"type": "response.output_text.delta", "delta": text}) + require.NoError(t, err) + frames = append(frames, string(data)) + } + prompt, completion := 4, 2 + if usageKind == "zero" { + prompt, completion = 0, 0 + } else if usageKind == "missing" { + prompt, completion = 10, service.CountTextToken(text, "gpt-test") + } + response := map[string]any{"id": "resp_terminal_failure", "status": status, "output": []any{}, + "error": map[string]any{"type": "server_error", "message": "synthetic terminal provider failure"}} + if usageKind != "missing" { + response["usage"] = map[string]int{"input_tokens": prompt, "output_tokens": completion, "total_tokens": prompt + completion} + } + terminal, err := common.Marshal(map[string]any{"type": event, "response": response}) + require.NoError(t, err) + frames = append(frames, string(terminal)) + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, "data: "+strings.Join(frames, "\n\ndata: ")+"\n\n") + })) + t.Cleanup(upstream.Close) + fingerprint := sha256.Sum256([]byte(t.Name())) + requestID := fmt.Sprintf("native-terminal-%x", fingerprint[:16]) + c, info, _ := streamFundingContext(t, "responses_native", requestID, upstream.URL, constant.ChannelTypeOpenAI, funding) + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + // A new subscription and changed preference must not receive + // either this failure refund or the delivered partial charge. + require.NoError(t, db.Create(&model.UserSubscription{Id: 956, UserId: 951, PlanId: 954, Status: "active", AmountTotal: 1000, + StartTime: time.Now().Unix() - 1800, EndTime: time.Now().Unix() + 1800}).Error) + info.UserSetting.BillingPreference = "wallet_only" + if funding.source == "wallet" { + info.UserSetting.BillingPreference = "subscription_only" + } + apiErr := ResponsesHelper(c, info) + if apiErr != nil { + service.HandleFailedBilling(c, info, apiErr) + service.HandleFailedBilling(c, info, apiErr) + } + quota := prompt + completion + if !delivered { + quota = 0 + } + var settlement model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).First(&settlement).Error) + require.Equal(t, funding.source, settlement.Source) + require.Equal(t, model.BillingSettlementStatusApplied, settlement.Status) + require.EqualValues(t, quota-100, settlement.AppliedFundingDelta) + require.EqualValues(t, quota-100, settlement.AppliedTokenDelta) + if delivered { + var effect model.BillingSettlementEffect + require.NoError(t, common.UnmarshalJsonStr(settlement.EffectPayload, &effect)) + require.EqualValues(t, quota, effect.Quota) + require.True(t, effect.QuotaIsActual) + replayInfo := &relaycommon.RelayInfo{RequestId: requestID, UserId: info.UserId, TokenId: info.TokenId, TokenKey: info.TokenKey, + OriginModelName: info.OriginModelName, ForcePreConsume: true, UserSetting: info.UserSetting} + replay, replayErr := service.NewBillingSession(c, replayInfo, 100) + require.Nil(t, replayErr) + require.Equal(t, funding.source, replayInfo.BillingSource) + require.NoError(t, replay.SettleWithEffect(quota, &effect)) + require.NoError(t, replay.SettleWithEffect(quota, &effect)) + replay.Refund(c) + } else { + require.Empty(t, settlement.EffectPayload, "a failed metadata-only request has no consumption projection") + for range 2 { + _, already, err := model.ApplyBillingSettlementOnce(model.BillingSettlementInput{ + OperationKey: settlement.OperationKey, Source: settlement.Source, UserID: settlement.UserID, + SubscriptionID: settlement.SubscriptionID, TokenID: settlement.TokenID, TokenKey: info.TokenKey, + FundingDelta: -100, TokenDelta: -100, SubscriptionPreConsumeRequestID: settlement.SubscriptionPreConsumeRequestID, + FinalizeSubscriptionPreConsume: settlement.FinalizeSubscriptionPreConsume, AllowMissingToken: true, + }) + require.NoError(t, err) + require.True(t, already) + } + } + model.ProcessPendingBillingSettlementsOnce() + assertStreamFundingBalances(t, db, funding, requestID, quota, delivered, prompt, completion) + var decoy model.UserSubscription + require.NoError(t, db.First(&decoy, 956).Error) + require.Zero(t, decoy.AmountUsed) + require.EqualValues(t, 1, calls.Load()) + require.Equal(t, delivered, info.HasRecordedChannelFirstResult()) + require.NotNil(t, apiErr, "a failed terminal status must never be returned as success") + require.Equal(t, "synthetic terminal provider failure", apiErr.Error()) + require.True(t, types.IsSkipRetryError(apiErr), "the earlier created metadata already committed the response") + }) + } + } + } + } + } +} From 5eeab558f6eb3bc3704386cb964f0dd5238e20d1 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Sun, 11 Oct 2026 02:38:27 +0800 Subject: [PATCH 05/24] fix: preserve final usage after terminal metadata write failures --- relay/channel/gemini/relay_responses.go | 86 ++++---- relay/channel/openai/responses_via_chat.go | 6 + relay/responses_failure_boundary_test.go | 88 ++++++++ ...esponses_terminal_usage_settlement_test.go | 192 ++++++++++++++++++ .../chat_to_responses_response.go | 15 ++ 5 files changed, 347 insertions(+), 40 deletions(-) create mode 100644 relay/responses_terminal_usage_settlement_test.go diff --git a/relay/channel/gemini/relay_responses.go b/relay/channel/gemini/relay_responses.go index 22cf0056..36c9b741 100644 --- a/relay/channel/gemini/relay_responses.go +++ b/relay/channel/gemini/relay_responses.go @@ -113,12 +113,15 @@ func GeminiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, r } return true } - sendChunk := func(chunk *dto.ChatCompletionsStreamResponse) bool { + sendChunk := func(chunk *dto.ChatCompletionsStreamResponse, retainMetadataUsage func()) bool { events, err := openaicompat.ChatCompletionsStreamChunkToResponsesEvents(chunk, state) if err != nil { streamErr = types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError) return false } + if terminalResponseSeen && retainMetadataUsage != nil && openaicompat.ChatToResponsesEventsAreMetadataOnly(events) { + retainMetadataUsage() + } for _, event := range events { if !sendEvent(event) { return false @@ -169,51 +172,54 @@ func GeminiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, r } } - if !sendChunk(response) { - return false - } - inputReported := gjson.Get(data, "usageMetadata.promptTokenCount").Type == gjson.Number - outputReported := gjson.Get(data, "usageMetadata.candidatesTokenCount").Type == gjson.Number - if inputReported || outputReported || dto.HasGeminiUsageMetadataTokens(geminiResponse.GetUsageMetadata()) { - if metadata := geminiResponse.GetUsageMetadata(); metadata != nil { - fallbackPrompt := 0 - if !inputReported { - fallbackPrompt = info.GetEstimatePromptTokens() - } - mapped := buildUsageFromGeminiMetadata(*metadata, fallbackPrompt) - // A numeric zero completion count must not be inferred from total. - if outputReported { - mapped.CompletionTokens = metadata.CandidatesTokenCount + metadata.ThoughtsTokenCount - } - if !outputReported { - promptTokens := mapped.PromptTokens - patchGeminiZeroCompletionUsage(c, info, &mapped, deliveredText.String(), 0) - if inputReported { - mapped.PromptTokens = promptTokens - mapped.TotalTokens = mapped.PromptTokens + mapped.CompletionTokens - if mapped.BillingUsage != nil && mapped.BillingUsage.Estimated { - mapped.BillingUsage.GeminiUsageMetadata.PromptTokenCount = metadata.PromptTokenCount - mapped.BillingUsage = dto.NewEstimatedGeminiChatBillingUsage(&mapped) + retainUsage := func() { + inputReported := gjson.Get(data, "usageMetadata.promptTokenCount").Type == gjson.Number + outputReported := gjson.Get(data, "usageMetadata.candidatesTokenCount").Type == gjson.Number + if inputReported || outputReported || dto.HasGeminiUsageMetadataTokens(geminiResponse.GetUsageMetadata()) { + if metadata := geminiResponse.GetUsageMetadata(); metadata != nil { + fallbackPrompt := 0 + if !inputReported { + fallbackPrompt = info.GetEstimatePromptTokens() + } + mapped := buildUsageFromGeminiMetadata(*metadata, fallbackPrompt) + // A numeric zero completion count must not be inferred from total. + if outputReported { + mapped.CompletionTokens = metadata.CandidatesTokenCount + metadata.ThoughtsTokenCount + } + if !outputReported { + promptTokens := mapped.PromptTokens + patchGeminiZeroCompletionUsage(c, info, &mapped, deliveredText.String(), 0) + if inputReported { + mapped.PromptTokens = promptTokens + mapped.TotalTokens = mapped.PromptTokens + mapped.CompletionTokens + if mapped.BillingUsage != nil && mapped.BillingUsage.Estimated { + mapped.BillingUsage.GeminiUsageMetadata.PromptTokenCount = metadata.PromptTokenCount + mapped.BillingUsage = dto.NewEstimatedGeminiChatBillingUsage(&mapped) + } } } + if mapped.BillingUsage == nil { + copyMetadata := *metadata + mapped.BillingUsage = &dto.BillingUsage{Source: dto.BillingUsageSourceGeminiChat, Semantic: dto.BillingUsageSemanticGemini, GeminiUsageMetadata: ©Metadata} + } + if mapped.PromptTokens != metadata.PromptTokenCount+metadata.ToolUsePromptTokenCount { + mapped.TotalTokens = mapped.PromptTokens + mapped.CompletionTokens + mapped.BillingUsage = dto.NewEstimatedGeminiChatBillingUsage(&mapped) + } + if mapped.BillingUsage != nil { + mapped.BillingUsage.TokenCountsReported = inputReported && outputReported + } + reportedUsage = &mapped + reportedOutputUsage = outputReported } - if mapped.BillingUsage == nil { - copyMetadata := *metadata - mapped.BillingUsage = &dto.BillingUsage{Source: dto.BillingUsageSourceGeminiChat, Semantic: dto.BillingUsageSemanticGemini, GeminiUsageMetadata: ©Metadata} - } - if mapped.PromptTokens != metadata.PromptTokenCount+metadata.ToolUsePromptTokenCount { - mapped.TotalTokens = mapped.PromptTokens + mapped.CompletionTokens - mapped.BillingUsage = dto.NewEstimatedGeminiChatBillingUsage(&mapped) - } - if mapped.BillingUsage != nil { - mapped.BillingUsage.TokenCountsReported = inputReported && outputReported - } - reportedUsage = &mapped - reportedOutputUsage = outputReported } } + if !sendChunk(response, retainUsage) { + return false + } + retainUsage() if isStop { - return sendChunk(helper.GenerateStopResponse(responseID, created, info.UpstreamModelName, finishReason)) + return sendChunk(helper.GenerateStopResponse(responseID, created, info.UpstreamModelName, finishReason), nil) } return true }) diff --git a/relay/channel/openai/responses_via_chat.go b/relay/channel/openai/responses_via_chat.go index 16c1c521..66d3ec4b 100644 --- a/relay/channel/openai/responses_via_chat.go +++ b/relay/channel/openai/responses_via_chat.go @@ -141,6 +141,12 @@ func OaiChatToResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo sr.Stop(streamErr) return } + if terminalResponseSeen && reported && openaicompat.ChatToResponsesEventsAreMetadataOnly(events) { + // Final usage for output already delivered survives closing-event + // write failures. Buffered generation still waits for delivery below. + deliveredUsage = openaicompat.UsageFromChatUsage(state.Usage) + usageReported = true + } for _, event := range events { if !sendEvent(event) { sr.Stop(streamErr) diff --git a/relay/responses_failure_boundary_test.go b/relay/responses_failure_boundary_test.go index 7549c55e..5975a67f 100644 --- a/relay/responses_failure_boundary_test.go +++ b/relay/responses_failure_boundary_test.go @@ -1,6 +1,7 @@ package relay import ( + "fmt" "io" "net/http" "net/http/httptest" @@ -166,3 +167,90 @@ func TestResponsesProviderExplicitZeroSurvivesFallback(t *testing.T) { } } } + +func TestResponsesTerminalBufferedToolInputKeepsDeliveredUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, provider := range []struct { + name string + frames []string + handle func(*gin.Context, *relaycommon.RelayInfo, *http.Response) (*dto.Usage, *types.MaxAPIError) + }{ + {"chat", []string{ + `{"choices":[{"delta":{"content":"visible"}}],"usage":{"prompt_tokens":10,"completion_tokens":2,"total_tokens":12}}`, + `{"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_buffered","function":{"name":"patch","arguments":"{\"input\":\"undelivered buffered input\"}"}}]}}]}`, + `{"choices":[{"delta":{},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":1000,"completion_tokens":999,"total_tokens":1999}}`, + }, openai.OaiChatToResponsesStreamHandler}, + {"gemini", []string{ + `{"candidates":[{"content":{"parts":[{"text":"visible"}]}}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":2,"totalTokenCount":12}}`, + `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call_buffered","name":"patch","args":{"input":"undelivered buffered input"}}}]}}]}`, + `{"candidates":[{"content":{"parts":[]},"finishReason":"MAX_TOKENS"}],"usageMetadata":{"promptTokenCount":1000,"candidatesTokenCount":999,"totalTokenCount":1999}}`, + }, gemini.GeminiResponsesStreamHandler}, + } { + t.Run(provider.name, func(t *testing.T) { + writer := &responsesSelectiveFailureWriter{httptest.NewRecorder(), "event: response.custom_tool_call_input.delta\n"} + c, _ := gin.CreateTestContext(writer) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + info := &relaycommon.RelayInfo{IsStream: true, RelayFormat: types.RelayFormatOpenAIResponses, DisablePing: true, + OriginModelName: "gpt-test", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-test"}, + Request: &dto.OpenAIResponsesRequest{Tools: []byte(`[{"type":"custom","name":"patch"}]`)}} + info.SetEstimatePromptTokens(10) + usage, apiErr := provider.handle(c, info, &http.Response{StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader("data: " + strings.Join(provider.frames, "\n\ndata: ") + "\n\n"))}) + require.NotNil(t, apiErr) + require.True(t, types.IsSkipRetryError(apiErr)) + require.NotNil(t, usage) + require.Equal(t, 10, usage.PromptTokens) + require.Equal(t, 2, usage.CompletionTokens) + require.NotContains(t, writer.Body.String(), "undelivered buffered input") + }) + } +} + +func TestResponsesTerminalUsageSurvivesToolAndReasoningMetadataFailure(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, provider := range []struct { + name string + outputs []string + final string + handle func(*gin.Context, *relaycommon.RelayInfo, *http.Response) (*dto.Usage, *types.MaxAPIError) + }{ + {"chat", []string{ + `{"choices":[{"delta":{"reasoning_content":"thinking"}}]}`, + `{"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_lookup","function":{"name":"lookup","arguments":"{\"q\":\"x\"}"}}]}}]}`, + `{"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_patch","function":{"name":"patch","arguments":"{\"input\":\"\"}"}}]}}]}`, + }, `{"choices":[{"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":%d,"completion_tokens":%d,"total_tokens":%d}}`, openai.OaiChatToResponsesStreamHandler}, + {"gemini", []string{ + `{"candidates":[{"content":{"parts":[{"thought":true,"text":"thinking"}]}}]}`, + `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call_lookup","name":"lookup","args":{"q":"x"}}}]}}]}`, + `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call_patch","name":"patch","args":{"input":""}}}]}}]}`, + }, `{"candidates":[{"content":{"parts":[]},"finishReason":"MAX_TOKENS"}],"usageMetadata":{"promptTokenCount":%d,"candidatesTokenCount":%d,"totalTokenCount":%d}}`, gemini.GeminiResponsesStreamHandler}, + } { + for outputIndex, event := range []string{"response.reasoning_summary_text.done", "response.function_call_arguments.done", "response.custom_tool_call_input.done"} { + for _, zero := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/%s/zero_%t", provider.name, event, zero), func(t *testing.T) { + prompt, completion := 4, 2 + if zero { + prompt, completion = 0, 0 + } + writer := &responsesSelectiveFailureWriter{httptest.NewRecorder(), "event: " + event + "\n"} + c, _ := gin.CreateTestContext(writer) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + info := &relaycommon.RelayInfo{IsStream: true, RelayFormat: types.RelayFormatOpenAIResponses, DisablePing: true, + OriginModelName: "gpt-test", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-test"}, + Request: &dto.OpenAIResponsesRequest{Tools: []byte(`[{"type":"custom","name":"patch"}]`)}} + info.SetEstimatePromptTokens(10) + frames := []string{provider.outputs[outputIndex], fmt.Sprintf(provider.final, prompt, completion, prompt+completion)} + usage, apiErr := provider.handle(c, info, &http.Response{StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader("data: " + strings.Join(frames, "\n\ndata: ") + "\n\n"))}) + require.NotNil(t, apiErr) + require.True(t, types.IsSkipRetryError(apiErr)) + require.NotNil(t, usage) + require.Equal(t, prompt, usage.PromptTokens) + require.Equal(t, completion, usage.CompletionTokens) + require.NotNil(t, usage.BillingUsage) + require.True(t, usage.BillingUsage.TokenCountsReported) + }) + } + } + } +} diff --git a/relay/responses_terminal_usage_settlement_test.go b/relay/responses_terminal_usage_settlement_test.go new file mode 100644 index 00000000..6567ff1b --- /dev/null +++ b/relay/responses_terminal_usage_settlement_test.go @@ -0,0 +1,192 @@ +package relay + +import ( + "crypto/sha256" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/constant" + "github.com/MAX-API-Next/MAX-API/model" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/service" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type terminalUsageFailureWriter struct { + *httptest.ResponseRecorder + failure string + triggered bool +} + +func (w *terminalUsageFailureWriter) Write(data []byte) (int, error) { + if strings.HasPrefix(string(data), "event: response.output_text.done\n") { + w.triggered = true + switch w.failure { + case "write": + return 0, io.ErrClosedPipe + case "short": + return len(data) - 1, nil + } + } + return w.ResponseRecorder.Write(data) +} + +func (w *terminalUsageFailureWriter) FlushError() error { + if w.triggered && w.failure == "flush" { + return io.ErrClosedPipe + } + w.ResponseRecorder.Flush() + return nil +} + +// Provider accounting has already arrived when a metadata-only terminal write +// fails. Run through real HTTP, ResponsesHelper and the durable funding/effect +// path: a disconnect cannot turn final zero into a positive estimated charge. +func TestResponsesTerminalMetadataFailurePreservesFinalUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitHttpClient() + service.InitTokenEncoders() + for _, provider := range []struct{ name, finishReason string }{ + {"chat", "stop"}, {"gemini_stop", "STOP"}, {"gemini_length", "MAX_TOKENS"}, {"gemini_safety", "SAFETY"}, + } { + for _, funding := range streamFundingCases() { + for _, usageKind := range []string{"known", "zero", "at_reservation", "above_reservation", "missing", "empty"} { + for _, previousUsage := range []bool{false, true} { + for _, failure := range []string{"write", "short", "flush"} { + t.Run(fmt.Sprintf("%s/%s/%s/previous_%t/%s", provider.name, funding.name, usageKind, previousUsage, failure), func(t *testing.T) { + db := setupStreamFundingLedger(t, funding) + prompt, completion := 4, 2 + switch usageKind { + case "zero": + prompt, completion = 0, 0 + case "at_reservation": + prompt, completion = 80, 20 + case "above_reservation": + prompt, completion = 80, 50 + case "missing", "empty": + prompt, completion = 10, service.EstimateTokenByModel("gpt-test", "visible") + if previousUsage { + completion = 2 + } + } + channelType, route := constant.ChannelTypeAdvancedCustom, "chat_to_responses" + first := map[string]any{"choices": []any{map[string]any{"delta": map[string]any{"content": "visible"}}}} + last := map[string]any{"choices": []any{map[string]any{"delta": map[string]any{}, "finish_reason": "stop"}}} + usageField := "usage" + counts := map[string]int{"prompt_tokens": prompt, "completion_tokens": completion, "total_tokens": prompt + completion} + previousCounts := map[string]int{"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12} + if strings.HasPrefix(provider.name, "gemini_") { + channelType, route = constant.ChannelTypeGemini, "gemini_to_responses" + first = map[string]any{"candidates": []any{map[string]any{"content": map[string]any{"parts": []any{map[string]any{"text": "visible"}}}}}} + last = map[string]any{"candidates": []any{map[string]any{"content": map[string]any{"parts": []any{}}, "finishReason": provider.finishReason}}} + usageField = "usageMetadata" + counts = map[string]int{"promptTokenCount": prompt, "candidatesTokenCount": completion, "totalTokenCount": prompt + completion} + previousCounts = map[string]int{"promptTokenCount": 10, "candidatesTokenCount": 2, "totalTokenCount": 12} + } + if previousUsage { + first[usageField] = previousCounts + } + if usageKind == "empty" { + last[usageField] = map[string]any{} + } else if usageKind != "missing" { + last[usageField] = counts + } + frames := []string{} + for _, frame := range []map[string]any{first, last} { + encoded, err := common.Marshal(frame) + require.NoError(t, err) + frames = append(frames, string(encoded)) + } + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, "data: "+strings.Join(frames, "\n\ndata: ")+"\n\n") + })) + t.Cleanup(upstream.Close) + fingerprint := sha256.Sum256([]byte(t.Name())) + requestID := fmt.Sprintf("terminal-usage-%x", fingerprint[:16]) + c, info, recorder := streamFundingContext(t, route, requestID, upstream.URL, channelType, funding) + writer := &terminalUsageFailureWriter{ResponseRecorder: recorder, failure: failure} + failureContext, _ := gin.CreateTestContext(writer) + c.Writer = failureContext.Writer + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + // A preference change and earlier-expiring plan cannot take + // ownership of an already reserved generation. + require.NoError(t, db.Create(&model.UserSubscription{Id: 956, UserId: 951, PlanId: 954, Status: "active", AmountTotal: 1000, + StartTime: time.Now().Unix() - 1800, EndTime: time.Now().Unix() + 1800}).Error) + info.UserSetting.BillingPreference = "wallet_only" + if funding.source == "wallet" { + info.UserSetting.BillingPreference = "subscription_only" + } + apiErr := ResponsesHelper(c, info) + require.NotNil(t, apiErr) + require.True(t, types.IsSkipRetryError(apiErr)) + require.True(t, writer.triggered, "failure must follow delivery and receipt of final usage") + require.Contains(t, recorder.Body.String(), `"delta":"visible"`) + require.NotContains(t, recorder.Body.String(), "event: response.completed") + service.HandleFailedBilling(c, info, apiErr) + service.HandleFailedBilling(c, info, apiErr) + quota := prompt + completion + var settlement model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).First(&settlement).Error) + require.Equal(t, funding.source, settlement.Source) + require.Equal(t, model.BillingSettlementStatusApplied, settlement.Status) + require.EqualValues(t, quota-100, settlement.AppliedFundingDelta, "final usage must survive metadata transport failure") + require.EqualValues(t, quota-100, settlement.AppliedTokenDelta) + var effect model.BillingSettlementEffect + require.NoError(t, common.UnmarshalJsonStr(settlement.EffectPayload, &effect)) + require.Equal(t, prompt, effect.PromptTokens) + require.Equal(t, completion, effect.CompletionTokens) + require.EqualValues(t, quota, effect.Quota) + require.True(t, effect.QuotaIsActual) + assertStreamFundingBalances(t, db, funding, requestID, quota, true, prompt, completion) + var beforeOutbox int64 + require.NoError(t, db.Model(&model.CacheInvalidationTask{}).Count(&beforeOutbox).Error) + sessions := []*service.BillingSession{} + for range 2 { + replayInfo := &relaycommon.RelayInfo{RequestId: requestID, UserId: info.UserId, TokenId: info.TokenId, TokenKey: info.TokenKey, + OriginModelName: info.OriginModelName, ForcePreConsume: true, UserSetting: info.UserSetting} + session, replayErr := service.NewBillingSession(c, replayInfo, 100) + require.Nil(t, replayErr) + require.Equal(t, funding.source, replayInfo.BillingSource) + sessions = append(sessions, session) + } + var workers sync.WaitGroup + results := make(chan error, len(sessions)) + for _, session := range sessions { + workers.Add(1) + go func() { + defer workers.Done() + results <- session.SettleWithEffect(quota, &effect) + session.Refund(c) + }() + } + workers.Wait() + close(results) + for err := range results { + require.NoError(t, err) + } + model.ProcessPendingBillingSettlementsOnce() + model.ProcessPendingBillingSettlementsOnce() + var afterOutbox int64 + require.NoError(t, db.Model(&model.CacheInvalidationTask{}).Count(&afterOutbox).Error) + require.Equal(t, beforeOutbox, afterOutbox) + var decoy model.UserSubscription + require.NoError(t, db.First(&decoy, 956).Error) + require.Zero(t, decoy.AmountUsed) + assertStreamFundingBalances(t, db, funding, requestID, quota, true, prompt, completion) + }) + } + } + } + } + } +} diff --git a/service/openaicompat/chat_to_responses_response.go b/service/openaicompat/chat_to_responses_response.go index d5749ca0..bb1cbe5c 100644 --- a/service/openaicompat/chat_to_responses_response.go +++ b/service/openaicompat/chat_to_responses_response.go @@ -36,6 +36,21 @@ type ChatToResponsesStreamEvent struct { Payload dto.ResponsesStreamResponse } +// ChatToResponsesEventsAreMetadataOnly excludes new output, including custom +// input drained at finish time. Closing an already delivered item does not +// make its provider usage depend on another successful downstream write. +func ChatToResponsesEventsAreMetadataOnly(events []ChatToResponsesStreamEvent) bool { + for _, event := range events { + switch event.Type { + case responsesEventCreated, responsesEventOutputItemDone, responsesEventFunctionArgsDone, + responsesEventReasoningSummaryDone, "response.output_text.done", "response.custom_tool_call_input.done": + default: + return false + } + } + return true +} + type ChatToResponsesStreamState struct { ID string Model string From 36addf8d2856a9a3ce63f4e6b8158c347b41485d Mon Sep 17 00:00:00 2001 From: CSCITech Date: Sun, 11 Oct 2026 11:07:39 +0800 Subject: [PATCH 06/24] fix: stabilize Gemini call identities and partial usage accounting --- .../gemini/responses_revalidation_test.go | 67 ++++++++ relay/channel/gemini/stream_function_calls.go | 77 +++++++--- relay/channel/openai/responses_usage_test.go | 45 ++++++ relay/channel/openai/responses_via_chat.go | 22 ++- ...sponses_review_boundary_settlement_test.go | 143 ++++++++++++++++++ 5 files changed, 330 insertions(+), 24 deletions(-) create mode 100644 relay/responses_review_boundary_settlement_test.go diff --git a/relay/channel/gemini/responses_revalidation_test.go b/relay/channel/gemini/responses_revalidation_test.go index 87401d30..ede7c179 100644 --- a/relay/channel/gemini/responses_revalidation_test.go +++ b/relay/channel/gemini/responses_revalidation_test.go @@ -8,6 +8,7 @@ import ( "strings" "testing" + "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/dto" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" "github.com/MAX-API-Next/MAX-API/types" @@ -16,6 +17,72 @@ import ( "github.com/tidwall/gjson" ) +func TestGeminiPendingCallIDSurvivesPartReordering(t *testing.T) { + for _, scenario := range []struct { + name, initial, final string + want map[string]string + conflict bool + }{ + {"shift_right", `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-a","name":"lookup","args":{"q":"a"},"willContinue":true}}]}}]}`, + `{"candidates":[{"content":{"parts":[{"text":"visible"},{"functionCall":{"id":"call-a","willContinue":false}}]},"finishReason":"STOP"}]}`, map[string]string{"0/call-a": "a"}, false}, + {"shift_left", `{"candidates":[{"content":{"parts":[{"text":"visible"},{"functionCall":{"id":"call-a","name":"lookup","args":{"q":"a"},"willContinue":true}}]}}]}`, + `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-a","willContinue":false}}]},"finishReason":"STOP"}]}`, map[string]string{"0/call-a": "a"}, false}, + {"swap_pending", `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-a","name":"lookup","args":{"q":"a"},"willContinue":true}},{"functionCall":{"id":"call-b","name":"lookup","args":{"q":"b"},"willContinue":true}}]}}]}`, + `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-b","willContinue":false}},{"functionCall":{"id":"call-a","willContinue":false}}]},"finishReason":"STOP"}]}`, map[string]string{"0/call-a": "a", "0/call-b": "b"}, false}, + {"candidate_isolation", `{"candidates":[{"index":0,"content":{"parts":[{"functionCall":{"id":"call-a","name":"lookup","args":{"q":"first"},"willContinue":true}}]}},{"index":1,"content":{"parts":[{"functionCall":{"id":"call-a","name":"lookup","args":{"q":"second"},"willContinue":true}}]}}]}`, + `{"candidates":[{"index":1,"content":{"parts":[{"text":"visible"},{"functionCall":{"id":"call-a","willContinue":false}}]},"finishReason":"STOP"},{"index":0,"content":{"parts":[{"text":"visible"},{"functionCall":{"id":"call-a","willContinue":false}}]},"finishReason":"STOP"}]}`, map[string]string{"0/call-a": "first", "1/call-a": "second"}, false}, + {"anonymous_position", `{"candidates":[{"content":{"parts":[{"functionCall":{"name":"lookup","args":{"q":"anonymous"},"willContinue":true}}]}}]}`, + `{"candidates":[{"content":{"parts":[{"functionCall":{"willContinue":false}}]},"finishReason":"STOP"}]}`, map[string]string{"0/": "anonymous"}, false}, + {"moved_conflict", `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-a","name":"lookup","args":{"q":"a"},"willContinue":true}}]}}]}`, + `{"candidates":[{"content":{"parts":[{"text":"visible"},{"functionCall":{"id":"call-a","name":"changed","willContinue":false}}]},"finishReason":"STOP"}]}`, nil, true}, + {"moved_then_anonymous", `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-a","name":"lookup","args":{"q":"a"},"willContinue":true}}]}}]}` + "\n" + + `{"candidates":[{"content":{"parts":[{"text":"visible"},{"functionCall":{"id":"call-a","willContinue":true}}]}}]}`, + `{"candidates":[{"content":{"parts":[{"text":"visible"},{"functionCall":{"willContinue":false}}]},"finishReason":"STOP"}]}`, map[string]string{"0/call-a": "a"}, false}, + {"completed_repeat_over_pending", `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-a","name":"lookup","args":{"q":"a"}}}]}}]}` + "\n" + + `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-b","name":"lookup","args":{"q":"b"},"willContinue":true}}]}}]}`, + `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-a"}},{"functionCall":{"id":"call-b","willContinue":false}}]},"finishReason":"STOP"}]}`, map[string]string{"0/call-b": "b"}, false}, + {"independent_named_slots", `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-a","name":"lookup","args":{"q":"a"},"willContinue":true}}]}}]}` + "\n" + + `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-b","name":"lookup","args":{"q":"b"},"willContinue":true}}]}}]}`, + `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-a","willContinue":false}},{"functionCall":{"id":"call-b","willContinue":false}}]},"finishReason":"STOP"}]}`, map[string]string{"0/call-a": "a", "0/call-b": "b"}, false}, + {"ambiguous_anonymous", `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-a","name":"lookup","args":{"q":"a"},"willContinue":true}}]}}]}` + "\n" + + `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-b","name":"lookup","args":{"q":"b"},"willContinue":true}}]}}]}`, + `{"candidates":[{"content":{"parts":[{"functionCall":{"willContinue":false}}]},"finishReason":"STOP"}]}`, nil, true}, + } { + t.Run(scenario.name, func(t *testing.T) { + state := newGeminiStreamFunctionCalls() + var final dto.GeminiChatResponse + for _, frame := range strings.Split(scenario.initial, "\n") { + var initial dto.GeminiChatResponse + require.NoError(t, common.UnmarshalJsonStr(frame, &initial)) + require.NoError(t, state.prepare(&initial)) + } + require.NoError(t, common.UnmarshalJsonStr(scenario.final, &final)) + err := state.prepare(&final) + if scenario.conflict { + require.Error(t, err) + return + } + require.NoError(t, err) + require.Empty(t, state.pending) + calls := 0 + for _, candidate := range final.Candidates { + for _, part := range candidate.Content.Parts { + if part.FunctionCall == nil { + continue + } + calls++ + call := part.FunctionCall + require.Equal(t, "lookup", call.FunctionName) + want, exists := scenario.want[fmt.Sprintf("%d/%s", candidate.Index, call.ID)] + require.True(t, exists) + require.Equal(t, want, call.Arguments.(map[string]any)["q"]) + } + } + require.Equal(t, len(scenario.want), calls) + }) + } +} + func runGeminiResponsesRevalidation(t *testing.T, frames []string, custom bool) (*dto.Usage, *types.MaxAPIError, string) { t.Helper() gin.SetMode(gin.TestMode) diff --git a/relay/channel/gemini/stream_function_calls.go b/relay/channel/gemini/stream_function_calls.go index 1ecd3bc4..29f69d42 100644 --- a/relay/channel/gemini/stream_function_calls.go +++ b/relay/channel/gemini/stream_function_calls.go @@ -11,6 +11,12 @@ import ( type geminiCallPosition struct { candidate int64 part int + id string +} + +type geminiPendingFunctionCall struct { + part int + call dto.FunctionCall } type geminiCallIdentity struct { candidate int64 @@ -20,12 +26,12 @@ type geminiCallIdentity struct { // Gemini args are complete JSON snapshots, unlike Chat arguments deltas. // Buffer willContinue sequences and emit each stable call ID only once. type geminiStreamFunctionCalls struct { - pending map[geminiCallPosition]dto.FunctionCall + pending map[geminiCallPosition]geminiPendingFunctionCall completed map[geminiCallIdentity]dto.FunctionCall } func newGeminiStreamFunctionCalls() *geminiStreamFunctionCalls { - return &geminiStreamFunctionCalls{pending: make(map[geminiCallPosition]dto.FunctionCall), completed: make(map[geminiCallIdentity]dto.FunctionCall)} + return &geminiStreamFunctionCalls{pending: make(map[geminiCallPosition]geminiPendingFunctionCall), completed: make(map[geminiCallIdentity]dto.FunctionCall)} } func (s *geminiStreamFunctionCalls) prepare(response *dto.GeminiChatResponse) error { @@ -37,27 +43,9 @@ func (s *geminiStreamFunctionCalls) prepare(response *dto.GeminiChatResponse) er parts = append(parts, part) continue } - position := geminiCallPosition{candidate.Index, partIndex} call := *part.FunctionCall call.ID = strings.TrimSpace(call.ID) call.FunctionName = strings.TrimSpace(call.FunctionName) - if previous, exists := s.pending[position]; exists { - if call.ID != "" && previous.ID != "" && call.ID != previous.ID { - return fmt.Errorf("Gemini function call changed id while incomplete") - } - if call.FunctionName != "" && previous.FunctionName != "" && call.FunctionName != previous.FunctionName { - return fmt.Errorf("Gemini function call changed name while incomplete") - } - if call.ID == "" { - call.ID = previous.ID - } - if call.FunctionName == "" { - call.FunctionName = previous.FunctionName - } - if call.Arguments == nil { - call.Arguments = previous.Arguments - } - } identity := geminiCallIdentity{candidate.Index, call.ID} if previous, exists := s.completed[identity]; call.ID != "" && exists { if call.FunctionName != "" && call.FunctionName != previous.FunctionName { @@ -78,8 +66,55 @@ func (s *geminiStreamFunctionCalls) prepare(response *dto.GeminiChatResponse) er } continue } + position := geminiCallPosition{candidate: candidate.Index, part: partIndex} + // Named calls use a stable key; anonymous continuations use the + // latest physical slot only when it identifies exactly one call. + if call.ID != "" { + position.part, position.id = -1, call.ID + if _, exists := s.pending[position]; !exists { + anonymous := geminiCallPosition{candidate: candidate.Index, part: partIndex} + if _, exists := s.pending[anonymous]; exists { + position = anonymous + } + } + } else { + found := false + for pendingPosition, pending := range s.pending { + if pendingPosition.candidate == candidate.Index && pending.part == partIndex { + if found { + return fmt.Errorf("Gemini anonymous function call has ambiguous identity") + } + position = pendingPosition + found = true + } + } + } + if pending, exists := s.pending[position]; exists { + previous := pending.call + if call.ID != "" && previous.ID != "" && call.ID != previous.ID { + return fmt.Errorf("Gemini function call changed id while incomplete") + } + if call.FunctionName != "" && previous.FunctionName != "" && call.FunctionName != previous.FunctionName { + return fmt.Errorf("Gemini function call changed name while incomplete") + } + if call.ID == "" { + call.ID = previous.ID + } + if call.FunctionName == "" { + call.FunctionName = previous.FunctionName + } + if call.Arguments == nil { + call.Arguments = previous.Arguments + } + } + identity = geminiCallIdentity{candidate.Index, call.ID} if call.WillContinue != nil && *call.WillContinue { - s.pending[position] = call + delete(s.pending, position) + key := geminiCallPosition{candidate: candidate.Index, part: partIndex} + if call.ID != "" { + key.part, key.id = -1, call.ID + } + s.pending[key] = geminiPendingFunctionCall{part: partIndex, call: call} continue } if call.FunctionName == "" { diff --git a/relay/channel/openai/responses_usage_test.go b/relay/channel/openai/responses_usage_test.go index 04d3492d..10fe25a8 100644 --- a/relay/channel/openai/responses_usage_test.go +++ b/relay/channel/openai/responses_usage_test.go @@ -5,17 +5,62 @@ import ( "io" "net/http" "net/http/httptest" + "strconv" "strings" "testing" "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/dto" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/service" "github.com/MAX-API-Next/MAX-API/types" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) +func TestChatToResponsesNonStreamPartialUsageDetails(t *testing.T) { + service.InitTokenEncoders() + for _, scenario := range []struct { + name, counters string + prompt, completion int + cached int + estimated bool + }{ + {"prompt", `"prompt_tokens":4`, 4, service.EstimateTokenByModel("gpt-test", "visible"), 1, true}, + {"completion", `"completion_tokens":2`, 10, 2, 1, true}, + {"prompt_zero", `"prompt_tokens":0`, 0, service.EstimateTokenByModel("gpt-test", "visible"), 0, true}, + {"completion_zero", `"completion_tokens":0`, 10, 0, 1, true}, + {"complete", `"prompt_tokens":4,"completion_tokens":2`, 4, 2, 1, false}, + {"complete_zero", `"prompt_tokens":0,"completion_tokens":0`, 0, 0, 0, false}, + } { + t.Run(scenario.name, func(t *testing.T) { + writer := httptest.NewRecorder() + c, _ := gin.CreateTestContext(writer) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + info := &relaycommon.RelayInfo{OriginModelName: "gpt-test", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-test"}} + info.SetEstimatePromptTokens(10) + payload := `{"choices":[{"message":{"role":"assistant","content":"visible"},"finish_reason":"stop"}],"usage":{` + scenario.counters + `,"prompt_tokens_details":{"cached_tokens":` + strconv.Itoa(scenario.cached) + `}}}` + usage, apiErr := OaiChatToResponsesHandler(c, info, &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": {"application/json"}}, Body: io.NopCloser(strings.NewReader(payload))}) + require.Nil(t, apiErr) + require.Equal(t, scenario.prompt, usage.PromptTokens) + require.Equal(t, scenario.completion, usage.CompletionTokens) + require.NotNil(t, usage.BillingUsage) + require.Equal(t, dto.BillingUsageSourceOAIChat, usage.BillingUsage.Source) + require.Equal(t, scenario.estimated, usage.BillingUsage.Estimated) + require.Equal(t, !scenario.estimated, usage.BillingUsage.TokenCountsReported) + require.Equal(t, scenario.cached, usage.BillingUsage.OpenAIUsage.PromptTokensDetails.CachedTokens) + var response dto.OpenAIResponsesResponse + require.NoError(t, common.Unmarshal(writer.Body.Bytes(), &response)) + require.Equal(t, scenario.prompt, response.Usage.InputTokens) + require.Equal(t, scenario.completion, response.Usage.OutputTokens) + if scenario.cached > 0 { + require.NotNil(t, response.Usage.InputTokensDetails) + require.Equal(t, scenario.cached, response.Usage.InputTokensDetails.CachedTokens) + } + }) + } +} + func detailedResponsesUsage() *dto.Usage { return &dto.Usage{ InputTokens: 100, diff --git a/relay/channel/openai/responses_via_chat.go b/relay/channel/openai/responses_via_chat.go index 66d3ec4b..418c23bd 100644 --- a/relay/channel/openai/responses_via_chat.go +++ b/relay/channel/openai/responses_via_chat.go @@ -15,6 +15,7 @@ import ( "github.com/MAX-API-Next/MAX-API/service/openaicompat" "github.com/MAX-API-Next/MAX-API/types" "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" ) func OaiChatToResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.MaxAPIError) { @@ -42,11 +43,26 @@ func OaiChatToResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp if err != nil { return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) } - if (usage == nil || usage.TotalTokens == 0) && !helper.HasUsageTokenFields(string(body), "usage", "prompt_tokens", "completion_tokens") { + promptReported := gjson.GetBytes(body, "usage.prompt_tokens").Type == gjson.Number + completionReported := gjson.GetBytes(body, "usage.completion_tokens").Type == gjson.Number + chatUsage := chatResp.Usage + chatUsage.BillingUsage = nil + if !promptReported || !completionReported { text := service.ExtractOutputTextFromResponses(responsesResp) - usage = service.ResponseText2Usage(c, text, info.UpstreamModelName, info.GetEstimatePromptTokens()) - responsesResp.Usage = openaicompat.UsageFromChatUsage(usage) + estimated := service.ResponseText2Usage(c, text, info.UpstreamModelName, info.GetEstimatePromptTokens()) + if !promptReported { + chatUsage.PromptTokens = estimated.PromptTokens + } + if !completionReported { + chatUsage.CompletionTokens = estimated.CompletionTokens + } + chatUsage.TotalTokens = chatUsage.PromptTokens + chatUsage.CompletionTokens } + chatUsage.BillingUsage = dto.NewReportedOpenAIChatBillingUsage(&chatUsage) + chatUsage.BillingUsage.TokenCountsReported = promptReported && completionReported + chatUsage.BillingUsage.Estimated = !chatUsage.BillingUsage.TokenCountsReported + usage = openaicompat.UsageFromChatUsage(&chatUsage) + responsesResp.Usage = usage if service.ResponseAuditEnabled() { service.SetRelayResponseAuditContent(info, service.ExtractOutputTextFromResponses(responsesResp)) } diff --git a/relay/responses_review_boundary_settlement_test.go b/relay/responses_review_boundary_settlement_test.go new file mode 100644 index 00000000..7003b74c --- /dev/null +++ b/relay/responses_review_boundary_settlement_test.go @@ -0,0 +1,143 @@ +package relay + +import ( + "crypto/sha256" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/constant" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/model" + "github.com/MAX-API-Next/MAX-API/service" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestResponsesNonStreamPartialUsageRetainsOriginalFunding(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitHttpClient() + service.InitTokenEncoders() + estimatedOutput := service.EstimateTokenByModel("gpt-test", "visible") + for _, scenario := range []struct { + name, usage string + prompt, completion int + }{ + {"prompt_only", `,"usage":{"prompt_tokens":4,"total_tokens":999}`, 4, estimatedOutput}, + {"completion_only", `,"usage":{"completion_tokens":2}`, 10, 2}, + {"prompt_zero", `,"usage":{"prompt_tokens":0}`, 0, estimatedOutput}, + {"completion_zero", `,"usage":{"completion_tokens":0}`, 10, 0}, + {"prompt_null", `,"usage":{"prompt_tokens":null,"completion_tokens":2}`, 10, 2}, + {"completion_null", `,"usage":{"prompt_tokens":4,"completion_tokens":null}`, 4, estimatedOutput}, + {"empty", `,"usage":{}`, 10, estimatedOutput}, + {"missing", "", 10, estimatedOutput}, + {"total_only", `,"usage":{"total_tokens":999}`, 10, estimatedOutput}, + {"complete", `,"usage":{"prompt_tokens":4,"completion_tokens":2,"total_tokens":6}`, 4, 2}, + {"complete_zero", `,"usage":{"prompt_tokens":0,"completion_tokens":0,"total_tokens":0}`, 0, 0}, + } { + for _, funding := range streamFundingCases() { + t.Run(scenario.name+"/"+funding.name, func(t *testing.T) { + db := setupStreamFundingLedger(t, funding) + payload := `{"choices":[{"message":{"role":"assistant","content":"visible"},"finish_reason":"stop"}]` + scenario.usage + `}` + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = fmt.Fprint(w, payload) + })) + defer upstream.Close() + requestID := fmt.Sprintf("partial-counter-%x", sha256.Sum256([]byte(t.Name()))) + c, info, recorder := streamFundingContext(t, "chat_to_responses", requestID, upstream.URL, constant.ChannelTypeAdvancedCustom, funding) + info.IsStream = false + info.Request.(*dto.OpenAIResponsesRequest).Stream = common.GetPointer(false) + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + session := info.Billing.(*service.BillingSession) + if funding.source == "subscription" { + info.UserSetting.BillingPreference = "wallet_only" + } else { + info.UserSetting.BillingPreference = "subscription_only" + } + apiErr := ResponsesHelper(c, info) + if apiErr != nil { + service.HandleFailedBilling(c, info, apiErr) + } + quota := scenario.prompt + scenario.completion + var settlement model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).First(&settlement).Error) + require.EqualValues(t, quota-100, settlement.AppliedFundingDelta) + require.EqualValues(t, quota-100, settlement.AppliedTokenDelta) + require.Nil(t, apiErr) + require.Equal(t, funding.source, settlement.Source) + require.Equal(t, model.BillingSettlementStatusApplied, settlement.Status) + var response dto.OpenAIResponsesResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.NotNil(t, response.Usage) + require.Equal(t, scenario.prompt, response.Usage.InputTokens) + require.Equal(t, scenario.completion, response.Usage.OutputTokens) + require.Equal(t, quota, response.Usage.TotalTokens) + var effect model.BillingSettlementEffect + require.NoError(t, common.UnmarshalJsonStr(settlement.EffectPayload, &effect)) + require.NoError(t, session.SettleWithEffect(quota, &effect)) + session.Refund(c) + model.ProcessPendingBillingSettlementsOnce() + assertStreamFundingBalances(t, db, funding, requestID, quota, true, scenario.prompt, scenario.completion) + }) + } + } +} + +func TestGeminiMovedFunctionCallRetainsOriginalFunding(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitHttpClient() + service.InitTokenEncoders() + for _, route := range []string{"chat_native", "gemini_to_responses"} { + for _, funding := range streamFundingCases() { + t.Run(route+"/"+funding.name, func(t *testing.T) { + db := setupStreamFundingLedger(t, funding) + frames := []string{ + `{"candidates":[{"index":0,"content":{"parts":[{"functionCall":{"id":"call-moved","name":"lookup","args":{"q":"max"},"willContinue":true}}]}}]}`, + `{"candidates":[{"index":0,"content":{"parts":[{"text":"visible"},{"functionCall":{"id":"call-moved","willContinue":false}}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":4,"candidatesTokenCount":2,"totalTokenCount":6}}`, + `{"candidates":[{"index":0,"content":{"parts":[{"functionCall":{"id":"call-moved","name":"lookup","args":{"q":"max"}}}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":4,"candidatesTokenCount":2,"totalTokenCount":6}}`, + } + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, "data: "+strings.Join(frames, "\n\ndata: ")+"\n\n") + })) + defer upstream.Close() + requestID := fmt.Sprintf("moved-call-%x", sha256.Sum256([]byte(t.Name()))) + c, info, recorder := streamFundingContext(t, route, requestID, upstream.URL, constant.ChannelTypeGemini, funding) + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + session := info.Billing.(*service.BillingSession) + info.UserSetting.BillingPreference = "wallet_only" + if funding.source == "wallet" { + info.UserSetting.BillingPreference = "subscription_only" + } + var apiErr *types.MaxAPIError + if route == "chat_native" { + apiErr = TextHelper(c, info) + } else { + apiErr = ResponsesHelper(c, info) + } + if apiErr != nil { + service.HandleFailedBilling(c, info, apiErr) + } + var settlement model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).First(&settlement).Error) + require.EqualValues(t, -94, settlement.AppliedFundingDelta) + require.EqualValues(t, -94, settlement.AppliedTokenDelta) + require.Nil(t, apiErr) + require.Equal(t, funding.source, settlement.Source) + require.Contains(t, recorder.Body.String(), "call-moved") + require.Contains(t, recorder.Body.String(), "lookup") + var effect model.BillingSettlementEffect + require.NoError(t, common.UnmarshalJsonStr(settlement.EffectPayload, &effect)) + require.NoError(t, session.SettleWithEffect(6, &effect)) + session.Refund(c) + model.ProcessPendingBillingSettlementsOnce() + assertStreamFundingBalances(t, db, funding, requestID, 6, true, 4, 2) + }) + } + } +} From 9becf17ec3083cde28cddf6b236fc557dcbd0b52 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Sun, 11 Oct 2026 11:57:04 +0800 Subject: [PATCH 07/24] fix: validate unnamed custom tools before name conflicts --- service/openaicompat/responses_request_to_chat.go | 6 +++--- .../openaicompat/responses_revalidation_test.go | 15 +++++++++++++++ 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/service/openaicompat/responses_request_to_chat.go b/service/openaicompat/responses_request_to_chat.go index a9031b41..ce06cb11 100644 --- a/service/openaicompat/responses_request_to_chat.go +++ b/service/openaicompat/responses_request_to_chat.go @@ -346,6 +346,9 @@ func responsesRequestToolsToChat(raw json.RawMessage) ([]dto.ToolCallRequest, er continue } name := strings.TrimSpace(common.Interface2String(tool["name"])) + if kind == "custom" && name == "" { + return nil, errors.New("custom tool is missing name") + } if previous, exists := names[name]; exists && (kind == "custom" || previous == "custom") { return nil, fmt.Errorf("ambiguous custom tool name %q", name) } @@ -365,9 +368,6 @@ func responsesRequestToolsToChat(raw json.RawMessage) ([]dto.ToolCallRequest, er }) case "custom": name := strings.TrimSpace(common.Interface2String(tool["name"])) - if name == "" { - return nil, errors.New("custom tool is missing name") - } out = append(out, dto.ToolCallRequest{ Type: "function", Function: dto.FunctionRequest{ diff --git a/service/openaicompat/responses_revalidation_test.go b/service/openaicompat/responses_revalidation_test.go index 3fb5b98e..77ac6e16 100644 --- a/service/openaicompat/responses_revalidation_test.go +++ b/service/openaicompat/responses_revalidation_test.go @@ -20,6 +20,21 @@ func TestResponsesToolNamesRejectAmbiguousCustomConversion(t *testing.T) { tools, err := responsesRequestToolsToChat([]byte(`[{"type":"function","name":"lookup"},{"type":"custom","name":"patch"}]`)) require.NoError(t, err) require.Len(t, tools, 2) + for _, scenario := range []struct { + name, tools string + }{ + {"function_first", `[{"type":"function"},{"type":"custom"}]`}, + {"custom_first", `[{"type":"custom"},{"type":"function"}]`}, + {"whitespace_function_first", `[{"type":"function","name":" "},{"type":"custom","name":"\t"}]`}, + {"whitespace_custom_first", `[{"type":"custom","name":"\t"},{"type":"function","name":" "}]`}, + {"two_custom", `[{"type":"custom"},{"type":"custom"}]`}, + {"single_custom", `[{"type":"custom"}]`}, + } { + t.Run(scenario.name, func(t *testing.T) { + _, err := responsesRequestToolsToChat([]byte(scenario.tools)) + require.EqualError(t, err, "custom tool is missing name") + }) + } } func TestResponsesCustomToolPreservesGrammarConstraint(t *testing.T) { From dbba1823ad2092df85edc082b29c5e3311c02f74 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Sun, 11 Oct 2026 13:45:23 +0800 Subject: [PATCH 08/24] fix: refund undelivered Gemini conversion failures --- model/subscription.go | 4 +- model/subscription_reset_billing_test.go | 2 + relay/channel/gemini/relay-gemini.go | 81 +++++++++- ...sponses_review_boundary_settlement_test.go | 147 ++++++++++++++++++ 4 files changed, 228 insertions(+), 6 deletions(-) diff --git a/model/subscription.go b/model/subscription.go index 7e3ed20d..a2c1afa8 100644 --- a/model/subscription.go +++ b/model/subscription.go @@ -1115,7 +1115,7 @@ func validateSubscriptionResetBillingTx(tx *gorm.DB, sub UserSubscription) error return query.Error } if query.RowsAffected > 0 { - return ErrSubscriptionResetPendingBilling + return fmt.Errorf("%w (subscription_id=%d)", ErrSubscriptionResetPendingBilling, sub.Id) } const batchSize = 200 @@ -1156,7 +1156,7 @@ func validateSubscriptionResetBillingTx(tx *gorm.DB, sub UserSubscription) error } for _, record := range records { if record.UserId != sub.UserId || !applied[record.RequestId] { - return ErrSubscriptionResetPendingBilling + return fmt.Errorf("%w (subscription_id=%d)", ErrSubscriptionResetPendingBilling, sub.Id) } } lastID = records[len(records)-1].Id diff --git a/model/subscription_reset_billing_test.go b/model/subscription_reset_billing_test.go index d0e4d4b8..71f610f4 100644 --- a/model/subscription_reset_billing_test.go +++ b/model/subscription_reset_billing_test.go @@ -215,6 +215,7 @@ func TestSubscriptionPlanResetRollsBackEarlierSubscriptions(t *testing.T) { require.NoError(t, DB.Create(&first).Error) _, err := AdminResetPlanSubscriptions(plan.Id, false) require.ErrorIs(t, err, ErrSubscriptionResetPendingBilling) + require.Contains(t, err.Error(), fmt.Sprintf("subscription_id=%d", sub.Id)) var got UserSubscription require.NoError(t, DB.First(&got, first.Id).Error) require.EqualValues(t, 25, got.AmountUsed, "an earlier reset must roll back when a later subscription is blocked") @@ -253,6 +254,7 @@ func TestSubscriptionResetRetainsManualFundingGuardAfterTombstoneCleanup(t *test require.EqualValues(t, 1, deleted) _, err = AdminResetUserSubscriptionsByPlan(sub.UserId, plan.Id, false) require.ErrorIs(t, err, ErrSubscriptionResetPendingBilling) + require.Contains(t, err.Error(), fmt.Sprintf("subscription_id=%d", sub.Id)) var got UserSubscription require.NoError(t, DB.First(&got, sub.Id).Error) require.EqualValues(t, 100, got.AmountUsed) diff --git a/relay/channel/gemini/relay-gemini.go b/relay/channel/gemini/relay-gemini.go index 509b5f16..e1efa6d5 100644 --- a/relay/channel/gemini/relay-gemini.go +++ b/relay/channel/gemini/relay-gemini.go @@ -1570,6 +1570,38 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http return usage, nil } +// Legacy SSE rendering can discard write errors. Observe each formatter call +// without counting a partial frame as generation delivery. +type geminiChatWriteObserver struct { + gin.ResponseWriter + writeErr error +} + +func (w *geminiChatWriteObserver) Write(data []byte) (int, error) { + n, err := w.ResponseWriter.Write(data) + w.recordWrite(n, len(data), err) + return n, err +} + +func (w *geminiChatWriteObserver) WriteString(data string) (int, error) { + n, err := w.ResponseWriter.WriteString(data) + w.recordWrite(n, len(data), err) + return n, err +} + +func (w *geminiChatWriteObserver) recordWrite(n, size int, err error) { + if err == nil && n != size { + err = io.ErrShortWrite + } + if w.writeErr == nil { + w.writeErr = err + } +} + +func (w *geminiChatWriteObserver) Unwrap() http.ResponseWriter { + return w.ResponseWriter +} + func GeminiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.MaxAPIError) { id := helper.GetResponseID(c) createAt := common.GetTimestamp() @@ -1578,6 +1610,35 @@ func GeminiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp * nextToolCallIndexByChoice := make(map[int]int) functionCalls := newGeminiStreamFunctionCalls() var conversionErr *types.MaxAPIError + generationDelivered := false + send := func(response *dto.ChatCompletionsStreamResponse) error { + // The scanner serializes callbacks and pings under its write mutex. + writer := &geminiChatWriteObserver{ResponseWriter: c.Writer} + before := writer.Size() + c.Writer = writer + defer func() { c.Writer = writer.ResponseWriter }() + err := handleStream(c, info, response) + if writer.writeErr == nil && writer.Size() > before { + // A later flush error does not erase a completely written frame. + for _, choice := range response.Choices { + if choice.Delta.GetContentString() != "" || choice.Delta.GetReasoningContent() != "" { + generationDelivered = true + } + for _, tool := range choice.Delta.ToolCalls { + if tool.Function.Name != "" || tool.Function.Arguments != "" { + generationDelivered = true + } + } + } + } + if writer.writeErr != nil { + err = writer.writeErr + } + if err != nil { + conversionErr = types.NewError(err, types.ErrorCodeBadResponseBody) + } + return err + } usage, err := geminiStreamHandler(c, info, resp, func(data string, geminiResponse *dto.GeminiChatResponse) bool { if prepareErr := functionCalls.prepare(geminiResponse); prepareErr != nil { @@ -1635,9 +1696,10 @@ func GeminiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp * emptyResponse.Choices[0].Delta.ToolCalls = copiedToolCalls } finishReason = constant.FinishReasonToolCalls - err := handleStream(c, info, emptyResponse) + err := send(emptyResponse) if err != nil { logger.LogError(c, err.Error()) + return false } response.ClearToolCalls() @@ -1645,20 +1707,25 @@ func GeminiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp * response.Choices[0].FinishReason = nil } } else { - err := handleStream(c, info, emptyResponse) + err := send(emptyResponse) if err != nil { logger.LogError(c, err.Error()) + return false } } } - err := handleStream(c, info, response) + err := send(response) if err != nil { logger.LogError(c, err.Error()) + return false } if isStop { if info.RelayFormat != types.RelayFormatClaude { - _ = handleStream(c, info, helper.GenerateStopResponse(id, createAt, info.UpstreamModelName, finishReason)) + if err := send(helper.GenerateStopResponse(id, createAt, info.UpstreamModelName, finishReason)); err != nil { + logger.LogError(c, err.Error()) + return false + } } } return true @@ -1667,6 +1734,9 @@ func GeminiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp * if err != nil { return usage, err } + if conversionErr == nil && info.StreamStatus.IsAbnormalEnd() { + conversionErr = types.NewError(fmt.Errorf("Gemini stream ended: %s", info.StreamStatus.Summary()), types.ErrorCodeBadResponseBody) + } if conversionErr == nil && len(functionCalls.pending) > 0 { conversionErr = types.NewError(fmt.Errorf("Gemini stream ended with an incomplete function call"), types.ErrorCodeBadResponseBody) } @@ -1674,6 +1744,9 @@ func GeminiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp * if c.Writer.Written() { conversionErr = types.NewError(conversionErr, conversionErr.GetErrorCode(), types.ErrOptionWithSkipRetry()) } + if !generationDelivered { + return nil, conversionErr + } return usage, conversionErr } diff --git a/relay/responses_review_boundary_settlement_test.go b/relay/responses_review_boundary_settlement_test.go index 7003b74c..686496aa 100644 --- a/relay/responses_review_boundary_settlement_test.go +++ b/relay/responses_review_boundary_settlement_test.go @@ -3,21 +3,168 @@ package relay import ( "crypto/sha256" "fmt" + "io" "net/http" "net/http/httptest" "strings" + "sync" "testing" "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/constant" "github.com/MAX-API-Next/MAX-API/dto" "github.com/MAX-API-Next/MAX-API/model" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" "github.com/MAX-API-Next/MAX-API/service" "github.com/MAX-API-Next/MAX-API/types" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) +type geminiChatFrameFailureWriter struct { + *httptest.ResponseRecorder + kind string + triggered bool +} + +func (w *geminiChatFrameFailureWriter) Write(data []byte) (int, error) { + if strings.Contains(string(data), `"tool_calls"`) { + switch w.kind { + case "write": + w.triggered = true + return 0, io.ErrClosedPipe + case "short": + w.triggered = true + return w.ResponseRecorder.Write(data[:3]) + } + } + return w.ResponseRecorder.Write(data) +} + +func (w *geminiChatFrameFailureWriter) WriteString(data string) (int, error) { + return w.Write([]byte(data)) +} + +func (w *geminiChatFrameFailureWriter) Flush() { + if w.kind == "flush" && strings.Contains(w.Body.String(), `"content":"visible"`) { + w.triggered = true + panic(io.ErrClosedPipe) + } + w.ResponseRecorder.Flush() +} + +func TestGeminiChatConversionFailureRetainsOnlyDeliveredFunding(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitHttpClient() + service.InitTokenEncoders() + pending := `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-pending","name":"lookup","args":{"q":"pending"},"willContinue":true}}]}}]}` + known := `{"promptTokenCount":4,"candidatesTokenCount":2,"totalTokenCount":6}` + withUsage := func(frame, usage string) string { + return strings.TrimSuffix(frame, "}") + `,"usageMetadata":` + usage + `}` + } + text := withUsage(`{"candidates":[{"content":{"parts":[{"text":"visible"}]}}]}`, known) + tool := withUsage(`{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-delivered","name":"lookup","args":{"q":"visible"}}}]}}]}`, known) + for _, scenario := range []struct { + name, frames, failure string + delivered bool + }{ + {"pending_missing", pending, "", false}, + {"pending_known", withUsage(pending, known), "", false}, + {"pending_zero", withUsage(pending, `{"promptTokenCount":0,"candidatesTokenCount":0,"totalTokenCount":0}`), "", false}, + {"pending_conflict", pending + "\n" + `{"candidates":[{"content":{"parts":[{"functionCall":{"id":"call-pending","name":"changed"}}]}}]}`, "", false}, + {"pending_stop", pending + "\n" + withUsage(`{"candidates":[{"content":{"parts":[]},"finishReason":"STOP"}]}`, known), "", false}, + {"pending_malformed", pending + "\n{broken", "", false}, + {"first_malformed", `{broken`, "", false}, + {"text_delivered", text + "\n" + pending, "", true}, + {"whitespace_delivered", withUsage(`{"candidates":[{"content":{"parts":[{"text":" "}]}}]}`, known) + "\n" + pending, "", true}, + {"reasoning_delivered", withUsage(`{"candidates":[{"content":{"parts":[{"text":"reasoning","thought":true}]}}]}`, known) + "\n" + pending, "", true}, + {"tool_delivered", tool + "\n" + pending, "", true}, + {"text_malformed", text + "\n{broken", "", true}, + {"tool_write", tool + "\n" + pending, "write", false}, + {"tool_short", tool + "\n" + pending, "short", false}, + {"text_flush", text + "\n" + pending, "flush", true}, + } { + for _, funding := range streamFundingCases() { + t.Run(scenario.name+"/"+funding.name, func(t *testing.T) { + db := setupStreamFundingLedger(t, funding) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, "data: "+strings.ReplaceAll(scenario.frames, "\n", "\n\ndata: ")+"\n\n") + })) + defer upstream.Close() + requestID := fmt.Sprintf("gemini-delivery-%x", sha256.Sum256([]byte(t.Name()))) + c, info, recorder := streamFundingContext(t, "chat_native", requestID, upstream.URL, constant.ChannelTypeGemini, funding) + writer := &geminiChatFrameFailureWriter{ResponseRecorder: recorder, kind: scenario.failure} + baseContext := c + c, _ = gin.CreateTestContext(writer) + c.Request, c.Keys = baseContext.Request, baseContext.Keys + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + session := info.Billing.(*service.BillingSession) + info.UserSetting.BillingPreference = "wallet_only" + if funding.source == "wallet" { + info.UserSetting.BillingPreference = "subscription_only" + } + apiErr := TextHelper(c, info) + if apiErr != nil { + service.HandleFailedBilling(c, info, apiErr) + service.HandleFailedBilling(c, info, apiErr) + } + quota, prompt, completion := 0, 0, 0 + if scenario.delivered { + quota, prompt, completion = 6, 4, 2 + } + var record model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).First(&record).Error) + require.EqualValues(t, quota-100, record.AppliedFundingDelta) + require.EqualValues(t, quota-100, record.AppliedTokenDelta) + require.Equal(t, funding.source, record.Source) + require.Equal(t, model.BillingSettlementStatusApplied, record.Status) + require.NotNil(t, apiErr) + require.Equal(t, scenario.name != "first_malformed", types.IsSkipRetryError(apiErr)) + if scenario.failure != "" { + require.True(t, writer.triggered) + } + if scenario.name == "tool_delivered" { + require.Contains(t, recorder.Body.String(), "call-delivered") + require.Contains(t, recorder.Body.String(), "lookup") + } + session.Refund(c) + model.ProcessPendingBillingSettlementsOnce() + assertStreamFundingBalances(t, db, funding, requestID, quota, scenario.delivered, prompt, completion) + var effect *model.BillingSettlementEffect + if scenario.delivered { + effect = &model.BillingSettlementEffect{} + require.NoError(t, common.UnmarshalJsonStr(record.EffectPayload, effect)) + replayInfo := &relaycommon.RelayInfo{RequestId: requestID, UserId: info.UserId, TokenId: info.TokenId, TokenKey: info.TokenKey, OriginModelName: info.OriginModelName, ForcePreConsume: true, UserSetting: info.UserSetting} + replay, replayErr := service.NewBillingSession(c, replayInfo, 100) + require.Nil(t, replayErr) + require.Equal(t, funding.source, replayInfo.BillingSource) + require.NoError(t, replay.SettleWithEffect(quota, effect)) + replay.Refund(c) + } + before := requestReplaySnapshot(t, db) + results := make(chan error, 2) + var workers sync.WaitGroup + for range 2 { + workers.Add(1) + go func() { + defer workers.Done() + _, _, err := model.ApplyBillingSettlementOnce(requestReplayInput(record, effect)) + results <- err + }() + } + workers.Wait() + close(results) + for err := range results { + require.NoError(t, err) + } + model.ProcessPendingBillingSettlementsOnce() + require.Equal(t, before, requestReplaySnapshot(t, db)) + }) + } + } +} + func TestResponsesNonStreamPartialUsageRetainsOriginalFunding(t *testing.T) { gin.SetMode(gin.TestMode) service.InitHttpClient() From 31081a2d86fd2081f73c4526a6c1e1159543e6bc Mon Sep 17 00:00:00 2001 From: CSCITech Date: Mon, 12 Oct 2026 00:42:44 +0800 Subject: [PATCH 09/24] test: verify billing recovery across process crashes --- service/billing_process_recovery_test.go | 783 +++++++++++++++++++++++ 1 file changed, 783 insertions(+) create mode 100644 service/billing_process_recovery_test.go diff --git a/service/billing_process_recovery_test.go b/service/billing_process_recovery_test.go new file mode 100644 index 00000000..ddb55ef7 --- /dev/null +++ b/service/billing_process_recovery_test.go @@ -0,0 +1,783 @@ +package service + +import ( + "bytes" + "context" + "crypto/sha256" + "database/sql" + "errors" + "fmt" + "net/http/httptest" + "os" + "os/exec" + "path/filepath" + "runtime" + "strings" + "sync" + "syscall" + "testing" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/model" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "gorm.io/gorm" + "gorm.io/gorm/logger" +) + +const billingProcessConfigEnv = "MAXAPI_BILLING_PROCESS_CONFIG" + +type billingProcessFunding struct { + name, preference, source string + wallet, subscriptionUsed int64 + allowOverflow bool +} + +func billingProcessFundingCases() []billingProcessFunding { + return []billingProcessFunding{ + {"wallet", "wallet_only", model.BillingSettlementSourceWallet, 1000, 0, false}, + {"subscription", "subscription_only", model.BillingSettlementSourceSubscription, 1000, 0, false}, + {"wallet_fallback", "subscription_first", model.BillingSettlementSourceWallet, 1000, 1000, true}, + {"subscription_fallback", "wallet_first", model.BillingSettlementSourceSubscription, 50, 0, false}, + } +} + +type billingProcessConfig struct { + Database string + Ready string + Release string + Role string + Phase string + Preference string + Source string + Quota int + Refund bool +} + +type billingProcessBarrier struct { + PID int + Phase string + Table string + InsideTransaction bool +} + +// Each child is this package's actual test executable, with fresh globals and +// connections. Only the child replaces model.DB. No production failpoints, +// graceful shutdown, transaction errors, or in-memory DBs simulate a crash. +func TestBillingProcessHelper(t *testing.T) { + configPath := os.Getenv(billingProcessConfigEnv) + if configPath == "" { + t.Skip("subprocess helper; invoked by the process recovery tests") + } + data, err := os.ReadFile(configPath) + require.NoError(t, err) + var cfg billingProcessConfig + require.NoError(t, common.Unmarshal(data, &cfg)) + require.Equal(t, filepath.Join(filepath.Dir(configPath), "ledger.sqlite"), cfg.Database) + db := openBillingProcessDB(t, cfg.Database) + model.DB, model.LOG_DB = db, db + common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled = true, false, false + common.LogConsumeEnabled, common.DataExportEnabled = true, false + gin.SetMode(gin.TestMode) + var user model.User + require.NoError(t, db.First(&user, 951).Error) + require.Equal(t, "synthetic-process-recovery", user.Username, "only the isolated fixture may be mutated") + + if cfg.Role == "work" { + c, info := billingProcessContext(cfg) + session, apiErr := NewBillingSession(c, info, 100) + require.Nil(t, apiErr) + require.Equal(t, cfg.Source, info.BillingSource) + require.Equal(t, 100, session.GetPreConsumedQuota()) + installBillingProcessBarrier(t, db, cfg) + if cfg.Refund { + input, err := session.PrepareRefundSettlement() + require.NoError(t, err) + require.NotNil(t, input) + _, _, err = model.ApplyBillingSettlementOnce(*input) + require.NoError(t, err) + } else { + prompt, completion := billingProcessUsage(cfg.Quota) + effect := newConsumeBillingSettlementEffect(info, model.RecordConsumeLogParams{ + Content: "synthetic process recovery", ChannelId: 953, + ModelName: "process-recovery-model", TokenId: 952, TokenName: "synthetic-process-token", + Group: "default", Quota: cfg.Quota, PromptTokens: prompt, CompletionTokens: completion, + }, info.RequestId, "", true) + require.NoError(t, session.SettleWithEffect(cfg.Quota, effect)) + } + if cfg.Phase == "complete" { + pauseBillingProcess(t, cfg, billingProcessBarrier{Phase: cfg.Phase}) + } + return + } + if cfg.Role == "reset" { + _, err := model.AdminResetUserSubscriptionsByPlan(951, 954, false) + require.NoError(t, err) + return + } + + if cfg.Release != "" { + writeBillingProcessJSON(t, cfg.Ready, billingProcessBarrier{PID: os.Getpid(), Phase: "ready"}) + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + ticker := time.NewTicker(10 * time.Millisecond) + defer ticker.Stop() + for { + if _, err := os.Stat(cfg.Release); err == nil { + break + } else { + require.True(t, errors.Is(err, os.ErrNotExist)) + } + select { + case <-ctx.Done(): + t.Fatal("parent did not release the recovery workers") + case <-ticker.C: + } + } + } + installBillingProcessBarrier(t, db, cfg) + if cfg.Role == "refund_competitor" { + c, info := billingProcessContext(cfg) + // The original session's refund intent competes with the already durable + // final usage intent. Do not create another request or reserve again. + var funding FundingSource = &WalletFunding{userId: 951, consumed: 100} + if cfg.Source == model.BillingSettlementSourceSubscription { + funding = &SubscriptionFunding{requestId: info.RequestId, userId: 951, subscriptionId: 955, preConsumed: 100} + } + session := &BillingSession{relayInfo: info, funding: funding, preConsumedQuota: 100, tokenConsumed: 100} + session.Refund(c) + // The current durable Refund completes synchronously, including conflicts. + session.mu.Lock() + refunded := session.refunded + session.mu.Unlock() + require.False(t, refunded, "a conflicting refund cannot replace the durable final usage") + return + } + // This is the production restart runner, including pending funding and + // effect-only repair. It reconstructs inputs from the first persisted intent. + model.ProcessPendingBillingSettlementsOnce() + model.ProcessPendingBillingSettlementsOnce() + record := billingProcessRecord(t, db) + t.Logf("recovery pid=%d status=%s effect=%s attempts=%d next_attempt=%d", os.Getpid(), record.Status, record.EffectStatus, record.Attempts, record.NextAttempt) + if record.Status == model.BillingSettlementStatusApplied { + var effect *model.BillingSettlementEffect + if record.EffectPayload != "" { + effect = &model.BillingSettlementEffect{} + require.NoError(t, common.UnmarshalJsonStr(record.EffectPayload, effect)) + } + applied, duplicate, err := model.ApplyBillingSettlementOnce(model.BillingSettlementInput{ + OperationKey: record.OperationKey, Source: record.Source, UserID: record.UserID, + SubscriptionID: record.SubscriptionID, TokenID: record.TokenID, + FundingDelta: record.FundingDelta, TokenDelta: record.TokenDelta, + TaskID: record.TaskID, TaskQuota: record.TaskQuota, TaskQuotaTarget: record.TaskQuotaTarget, + SubscriptionPreConsumeRequestID: record.SubscriptionPreConsumeRequestID, + FinalizeSubscriptionPreConsume: record.FinalizeSubscriptionPreConsume, AllowMissingToken: record.AllowMissingToken, + ManualOnFailure: record.ManualOnFailure, PreConsumeRequestID: record.PreConsumeRequestID, + PreConsumeModelName: record.PreConsumeModelName, PreConsumeRequestedQuota: record.PreConsumeRequestedQuota, + PreConsumeEffectiveQuota: record.PreConsumeEffectiveQuota, Effect: effect, + }) + require.NoError(t, err) + require.True(t, duplicate) + require.Equal(t, record.AppliedFundingDelta, applied, "a new process must return the original signed result, including zero") + } + if cfg.Phase == "complete" { + pauseBillingProcess(t, cfg, billingProcessBarrier{Phase: cfg.Phase}) + } +} + +func billingProcessContext(cfg billingProcessConfig) (*gin.Context, *relaycommon.RelayInfo) { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + info := &relaycommon.RelayInfo{RequestId: "process-recovery-request", UserId: 951, TokenId: 952, + TokenKey: "synthetic-process-recovery-key", OriginModelName: "process-recovery-model", ForcePreConsume: true, + UserSetting: dto.UserSetting{BillingPreference: cfg.Preference}} + return c, info +} + +func billingProcessUsage(quota int) (int, int) { + if quota == 0 { + return 0, 0 + } + return quota - 2, 2 +} + +func openBillingProcessDB(t *testing.T, path string) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open(path), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) + require.NoError(t, err) + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + // Bound SQLite writer waits. Use SQLite's durable journal; no + // synchronous=OFF or memory/WAL shortcuts weaken the interruption test. + require.NoError(t, db.Exec("PRAGMA busy_timeout = 30000").Error) + t.Cleanup(func() { require.NoError(t, sqlDB.Close()) }) + return db +} + +func seedBillingProcessDB(t *testing.T, funding billingProcessFunding) (string, *gorm.DB) { + t.Helper() + dir := t.TempDir() + db := openBillingProcessDB(t, filepath.Join(dir, "ledger.sqlite")) + require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}, &model.Channel{}, &model.Log{}, &model.BillingLogReceipt{}, + &model.BillingSettlement{}, &model.BillingPreConsumeSelection{}, &model.CacheInvalidationTask{}, + &model.UserSubscription{}, &model.SubscriptionPreConsumeRecord{}, &model.SubscriptionPlan{})) + require.NoError(t, db.Create(&model.User{Id: 951, Username: "synthetic-process-recovery", Quota: funding.wallet}).Error) + require.NoError(t, db.Create(&model.Token{Id: 952, UserId: 951, Key: "synthetic-process-recovery-key", + Name: "synthetic-process-token", RemainQuota: 1000, Status: common.TokenStatusEnabled}).Error) + require.NoError(t, db.Create(&model.Channel{Id: 953}).Error) + require.NoError(t, db.Create(&model.SubscriptionPlan{Id: 954, Title: "synthetic-process-plan", Enabled: true, + TotalAmount: 1000, AllowWalletOverflow: &funding.allowOverflow, QuotaResetPeriod: model.SubscriptionResetNever}).Error) + now := time.Now().Unix() + require.NoError(t, db.Create(&model.UserSubscription{Id: 955, UserId: 951, PlanId: 954, Status: "active", + AmountTotal: 1000, AmountUsed: funding.subscriptionUsed, AllowWalletOverflow: funding.allowOverflow, + StartTime: now - 3600, EndTime: now + 86400, LastResetTime: now - 3600}).Error) + return dir, db +} + +// Callbacks exist only in the test child. The marker is written after the SQL +// at that boundary succeeded, and the process remains blocked until killed. +func installBillingProcessBarrier(t *testing.T, db *gorm.DB, cfg billingProcessConfig) { + t.Helper() + if cfg.Phase == "" || cfg.Phase == "complete" { + return + } + var once sync.Once + stop := func(tx *gorm.DB) { + if tx.Error != nil { + return + } + once.Do(func() { + _, inTx := tx.Statement.ConnPool.(*sql.Tx) + pauseBillingProcess(t, cfg, billingProcessBarrier{Phase: cfg.Phase, Table: tx.Statement.Table, InsideTransaction: inTx}) + }) + } + create := func(tx *gorm.DB) { + switch cfg.Phase { + case "intent_committed": + if record, ok := tx.Statement.Dest.(*model.BillingSettlement); ok && record.OperationKey == model.BillingRequestFinalizeOperationKey("process-recovery-request") { + stop(tx) + } + case "funding_committed", "receipt_before_log": + if tx.Statement.Table == "billing_log_receipts" { + stop(tx) + } + case "log_before_commit": + if tx.Statement.Table == "logs" { + stop(tx) + } + } + } + update := func(tx *gorm.DB) { + if tx.Statement.Table != "billing_settlements" { + return + } + updates, ok := tx.Statement.Dest.(map[string]interface{}) + if !ok { + return + } + if cfg.Phase == "funding_before_commit" && updates["status"] == model.BillingSettlementStatusApplied || + cfg.Phase == "log_committed" && updates["effect_status"] == model.BillingSettlementEffectApplying || + cfg.Phase == "usage_before_commit" && updates["effect_status"] == model.BillingSettlementEffectApplied { + stop(tx) + } + } + const name = "test:billing-process-barrier" + if cfg.Phase == "intent_committed" { + require.NoError(t, db.Callback().Create().After("gorm:commit_or_rollback_transaction").Register(name, create)) + } else if cfg.Phase == "funding_committed" { + require.NoError(t, db.Callback().Create().Before("gorm:create").Register(name, create)) + } else if cfg.Phase == "receipt_before_log" || cfg.Phase == "log_before_commit" { + require.NoError(t, db.Callback().Create().After("gorm:create").Register(name, create)) + } else if cfg.Phase == "log_committed" { + require.NoError(t, db.Callback().Update().Before("gorm:update").Register(name, update)) + } else if cfg.Phase == "funding_before_commit" || cfg.Phase == "usage_before_commit" { + require.NoError(t, db.Callback().Update().After("gorm:update").Register(name, update)) + } else { + t.Fatalf("unknown interruption phase %q", cfg.Phase) + } +} + +func writeBillingProcessJSON(t *testing.T, path string, value any) { + t.Helper() + data, err := common.Marshal(value) + require.NoError(t, err) + // Rename prevents the parent from observing a partially written marker. + require.NoError(t, os.WriteFile(path+".writing", data, 0600)) + require.NoError(t, os.Rename(path+".writing", path)) +} + +func pauseBillingProcess(t *testing.T, cfg billingProcessConfig, event billingProcessBarrier) { + t.Helper() + event.PID = os.Getpid() + writeBillingProcessJSON(t, cfg.Ready, event) + <-make(chan struct{}) +} + +type billingProcessOutput struct { + mu sync.Mutex + bytes.Buffer +} + +func (b *billingProcessOutput) Write(p []byte) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + return b.Buffer.Write(p) +} + +func (b *billingProcessOutput) text() string { + b.mu.Lock() + defer b.mu.Unlock() + return b.Buffer.String() +} + +type billingTestProcess struct { + cmd *exec.Cmd + done chan error + exited chan struct{} + output *billingProcessOutput + ready string +} + +func startBillingTestProcess(t *testing.T, dir string, cfg billingProcessConfig) *billingTestProcess { + t.Helper() + exe, err := os.Executable() + require.NoError(t, err) + file, err := os.CreateTemp(dir, "process-*.json") + require.NoError(t, err) + require.NoError(t, file.Close()) + cfg.Database, cfg.Ready = filepath.Join(dir, "ledger.sqlite"), file.Name()+".ready" + writeBillingProcessJSON(t, file.Name(), cfg) + ctx, cancel := context.WithTimeout(context.Background(), 45*time.Second) + cmd := exec.CommandContext(ctx, exe, "-test.run=^TestBillingProcessHelper$", "-test.timeout=40s", "-test.v") + cmd.Env = append(os.Environ(), billingProcessConfigEnv+"="+file.Name()) + p := &billingTestProcess{cmd: cmd, done: make(chan error, 1), exited: make(chan struct{}), output: &billingProcessOutput{}, ready: cfg.Ready} + cmd.Stdout, cmd.Stderr = p.output, p.output + require.NoError(t, cmd.Start()) + go func() { + p.done <- cmd.Wait() + close(p.exited) + }() + t.Cleanup(func() { + cancel() + <-p.exited + }) + return p +} + +func awaitBillingProcessBarrier(t *testing.T, p *billingTestProcess) billingProcessBarrier { + t.Helper() + ticker := time.NewTicker(10 * time.Millisecond) + defer ticker.Stop() + timer := time.NewTimer(30 * time.Second) + defer timer.Stop() + for { + data, err := os.ReadFile(p.ready) + if err == nil { + var event billingProcessBarrier + require.NoError(t, common.Unmarshal(data, &event)) + require.Equal(t, p.cmd.Process.Pid, event.PID) + return event + } + require.True(t, billingProcessMarkerNotReady(err), "barrier read: %v", err) + select { + case err := <-p.done: + t.Fatalf("child exited before its barrier: %v\n%s", err, p.output.text()) + case <-timer.C: + t.Fatalf("child never reached its barrier\n%s", p.output.text()) + case <-ticker.C: + } + } +} + +func billingProcessMarkerNotReady(err error) bool { + // Windows can temporarily deny access while the atomic rename is observed. + // Retry only ERROR_SHARING_VIOLATION (32), never arbitrary I/O failures. + return errors.Is(err, os.ErrNotExist) || runtime.GOOS == "windows" && errors.Is(err, syscall.Errno(32)) +} + +func killBillingTestProcess(t *testing.T, p *billingTestProcess, phase string) { + t.Helper() + event := awaitBillingProcessBarrier(t, p) + require.Equal(t, phase, event.Phase) + if phase != "complete" && phase != "intent_committed" { + require.True(t, event.InsideTransaction, "the interruption must hit the actual SQL transaction") + } + require.NoError(t, p.cmd.Process.Kill()) + err := <-p.done + var exitErr *exec.ExitError + require.ErrorAs(t, err, &exitErr, "an OS-forced termination must not be a normal test exit") + require.False(t, exitErr.Success()) + t.Logf("forced termination pid=%d phase=%s table=%s inside_transaction=%t", event.PID, phase, event.Table, event.InsideTransaction) +} + +func finishBillingTestProcess(t *testing.T, p *billingTestProcess) { + t.Helper() + require.NoError(t, <-p.done, "child process failed\n%s", p.output.text()) +} + +func billingProcessRecord(t *testing.T, db *gorm.DB) model.BillingSettlement { + t.Helper() + var record model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey("process-recovery-request")).First(&record).Error) + return record +} + +func assertBillingProcessLedger(t *testing.T, db *gorm.DB, funding billingProcessFunding, quota int, phase string, refund bool) { + t.Helper() + beforeFunding := phase == "intent_committed" || phase == "funding_before_commit" + fundedQuota := quota + if beforeFunding { + fundedQuota = 100 + } + if refund && !beforeFunding { + fundedQuota = 0 + } + var user model.User + var token model.Token + var sub model.UserSubscription + var channel model.Channel + require.NoError(t, db.First(&user, 951).Error) + require.NoError(t, db.First(&token, 952).Error) + require.NoError(t, db.First(&sub, 955).Error) + require.NoError(t, db.First(&channel, 953).Error) + wallet, subUsed := funding.wallet, funding.subscriptionUsed + if funding.source == model.BillingSettlementSourceWallet { + wallet -= int64(fundedQuota) + } else { + subUsed += int64(fundedQuota) + } + require.EqualValues(t, wallet, user.Quota) + require.EqualValues(t, subUsed, sub.AmountUsed) + require.EqualValues(t, 1000-fundedQuota, token.RemainQuota) + require.EqualValues(t, fundedQuota, token.UsedQuota) + counted := phase == "complete" && !refund + wantUsed, wantRequests := 0, 0 + if counted { + wantUsed, wantRequests = quota, 1 + } + require.EqualValues(t, wantUsed, user.UsedQuota) + require.EqualValues(t, wantRequests, user.RequestCount) + require.EqualValues(t, wantUsed, channel.UsedQuota) + var logs []model.Log + var receipts []model.BillingLogReceipt + require.NoError(t, db.Find(&logs).Error) + require.NoError(t, db.Find(&receipts).Error) + wantLogs := 0 + if !refund && (phase == "log_committed" || phase == "usage_before_commit" || phase == "complete") { + wantLogs = 1 + } + require.Len(t, logs, wantLogs) + require.Len(t, receipts, wantLogs) + if wantLogs == 1 { + prompt, completion := billingProcessUsage(quota) + require.EqualValues(t, quota, logs[0].Quota) + require.Equal(t, prompt, logs[0].PromptTokens) + require.Equal(t, completion, logs[0].CompletionTokens) + require.Equal(t, "process-recovery-request", logs[0].RequestId) + require.Equal(t, model.BillingRequestFinalizeOperationKey("process-recovery-request"), receipts[0].OperationKey) + } + record := billingProcessRecord(t, db) + require.Equal(t, funding.source, record.Source) + require.EqualValues(t, 951, record.UserID) + require.EqualValues(t, 952, record.TokenID) + require.EqualValues(t, quota-100, record.FundingDelta) + require.EqualValues(t, quota-100, record.TokenDelta) + if beforeFunding { + require.Equal(t, model.BillingSettlementStatusPending, record.Status) + require.Zero(t, record.AppliedFundingDelta) + require.Zero(t, record.AppliedTokenDelta) + require.Empty(t, record.EffectStatus) + } else { + require.Equal(t, model.BillingSettlementStatusApplied, record.Status) + require.EqualValues(t, quota-100, record.AppliedFundingDelta) + require.EqualValues(t, quota-100, record.AppliedTokenDelta) + if !refund { + if counted { + require.Equal(t, model.BillingSettlementEffectApplied, record.EffectStatus) + } else { + require.Equal(t, model.BillingSettlementEffectPending, record.EffectStatus) + } + } else { + require.Empty(t, record.EffectStatus) + } + } + if refund { + require.Empty(t, record.EffectPayload) + } else { + var effect model.BillingSettlementEffect + require.NoError(t, common.UnmarshalJsonStr(record.EffectPayload, &effect)) + require.True(t, effect.QuotaIsActual) + require.EqualValues(t, quota, effect.Quota) + } + var selections []model.BillingPreConsumeSelection + require.NoError(t, db.Find(&selections).Error) + require.Len(t, selections, 1) + require.Equal(t, funding.source, selections[0].Source) + require.EqualValues(t, 100, selections[0].EffectiveQuota) + var preconsume []model.SubscriptionPreConsumeRecord + require.NoError(t, db.Find(&preconsume).Error) + var settlementCount int64 + require.NoError(t, db.Model(&model.BillingSettlement{}).Count(&settlementCount).Error) + if funding.source == model.BillingSettlementSourceSubscription { + require.EqualValues(t, 1, settlementCount) + require.Len(t, preconsume, 1) + require.EqualValues(t, 100, preconsume[0].PreConsumed) + require.Equal(t, "process-recovery-request", record.SubscriptionPreConsumeRequestID) + wantStatus := "consumed" + if refund && !beforeFunding { + wantStatus = "refunded" + } + require.Equal(t, wantStatus, preconsume[0].Status) + } else { + require.EqualValues(t, 2, settlementCount) + require.Empty(t, preconsume) + } + t.Logf("ledger phase=%s source=%s quota=%d wallet=%d subscription_used=%d token_remaining=%d token_used=%d user_used=%d requests=%d channel_used=%d logs=%d receipts=%d funding=%s effect=%s", + phase, record.Source, quota, user.Quota, sub.AmountUsed, token.RemainQuota, token.UsedQuota, user.UsedQuota, + user.RequestCount, channel.UsedQuota, len(logs), len(receipts), record.Status, record.EffectStatus) +} + +func billingProcessSnapshot(t *testing.T, db *gorm.DB) [32]byte { + t.Helper() + var snapshot struct { + Users []model.User + Tokens []model.Token + Channels []model.Channel + Subs []model.UserSubscription + Reserved []model.SubscriptionPreConsumeRecord + Selections []model.BillingPreConsumeSelection + Settlements []model.BillingSettlement + Logs []model.Log + Receipts []model.BillingLogReceipt + Outbox []model.CacheInvalidationTask + } + for _, target := range []any{&snapshot.Users, &snapshot.Tokens, &snapshot.Channels, &snapshot.Subs, &snapshot.Reserved, + &snapshot.Selections, &snapshot.Settlements, &snapshot.Logs, &snapshot.Receipts, &snapshot.Outbox} { + require.NoError(t, db.Order("id").Find(target).Error) + } + data, err := common.Marshal(snapshot) + require.NoError(t, err) + return sha256.Sum256(data) +} + +func TestBillingProcessCrashRecovery(t *testing.T) { + phases := []string{"intent_committed", "funding_before_commit", "funding_committed", "receipt_before_log", + "log_before_commit", "log_committed", "usage_before_commit", "complete"} + for _, funding := range billingProcessFundingCases() { + for _, quota := range []int{130, 100, 6, 0} { + for _, phase := range phases { + t.Run(fmt.Sprintf("%s/quota_%d/%s", funding.name, quota, phase), func(t *testing.T) { + dir, db := seedBillingProcessDB(t, funding) + cfg := billingProcessConfig{Role: "work", Phase: phase, Preference: funding.preference, Source: funding.source, Quota: quota} + killBillingTestProcess(t, startBillingTestProcess(t, dir, cfg), phase) + assertBillingProcessLedger(t, db, funding, quota, phase, false) + original := billingProcessRecord(t, db) + cfg.Role, cfg.Phase = "recover", "" + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + assertBillingProcessLedger(t, db, funding, quota, "complete", false) + applied := billingProcessRecord(t, db) + require.Equal(t, original.EffectPayload, applied.EffectPayload, "the first persisted usage, including explicit zero, must survive restart") + snapshot := billingProcessSnapshot(t, db) + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + require.Equal(t, snapshot, billingProcessSnapshot(t, db), "another OS process must not change any durable accounting row") + }) + } + } + } +} + +func TestBillingProcessRefundRecovery(t *testing.T) { + for _, funding := range billingProcessFundingCases() { + for _, phase := range []string{"intent_committed", "funding_before_commit", "complete"} { + t.Run(funding.name+"/"+phase, func(t *testing.T) { + dir, db := seedBillingProcessDB(t, funding) + cfg := billingProcessConfig{Role: "work", Phase: phase, Preference: funding.preference, Source: funding.source, Refund: true} + killBillingTestProcess(t, startBillingTestProcess(t, dir, cfg), phase) + assertBillingProcessLedger(t, db, funding, 0, phase, true) + cfg.Role, cfg.Phase = "recover", "" + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + assertBillingProcessLedger(t, db, funding, 0, "complete", true) + snapshot := billingProcessSnapshot(t, db) + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + require.Equal(t, snapshot, billingProcessSnapshot(t, db)) + }) + } + } +} + +func TestBillingProcessRepeatedCrashRecovery(t *testing.T) { + for _, funding := range billingProcessFundingCases() { + for _, quota := range []int{130, 0} { + t.Run(fmt.Sprintf("%s/quota_%d", funding.name, quota), func(t *testing.T) { + dir, db := seedBillingProcessDB(t, funding) + cfg := billingProcessConfig{Role: "work", Phase: "funding_before_commit", Preference: funding.preference, Source: funding.source, Quota: quota} + killBillingTestProcess(t, startBillingTestProcess(t, dir, cfg), cfg.Phase) + assertBillingProcessLedger(t, db, funding, quota, cfg.Phase, false) + cfg.Role = "recover" + for _, phase := range []string{"log_committed", "usage_before_commit"} { + cfg.Phase = phase + killBillingTestProcess(t, startBillingTestProcess(t, dir, cfg), phase) + assertBillingProcessLedger(t, db, funding, quota, phase, false) + } + cfg.Phase = "" + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + assertBillingProcessLedger(t, db, funding, quota, "complete", false) + snapshot := billingProcessSnapshot(t, db) + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + require.Equal(t, snapshot, billingProcessSnapshot(t, db)) + }) + } + } +} + +func TestBillingProcessConcurrentRecoveryAndRefund(t *testing.T) { + for _, funding := range billingProcessFundingCases() { + for _, quota := range []int{130, 0} { + t.Run(fmt.Sprintf("%s/quota_%d", funding.name, quota), func(t *testing.T) { + dir, db := seedBillingProcessDB(t, funding) + cfg := billingProcessConfig{Role: "work", Phase: "intent_committed", Preference: funding.preference, Source: funding.source, Quota: quota} + killBillingTestProcess(t, startBillingTestProcess(t, dir, cfg), cfg.Phase) + original := billingProcessRecord(t, db) + cfg.Phase, cfg.Release = "", filepath.Join(dir, "release") + workers := make([]*billingTestProcess, 0, 3) + for _, role := range []string{"recover", "recover", "refund_competitor"} { + cfg.Role = role + p := startBillingTestProcess(t, dir, cfg) + workers = append(workers, p) + awaitBillingProcessBarrier(t, p) + } + require.NoError(t, os.WriteFile(cfg.Release, []byte("release"), 0600)) + for _, p := range workers { + finishBillingTestProcess(t, p) + } + cfg.Role, cfg.Release = "recover", "" + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + current := billingProcessRecord(t, db) + t.Logf("post-contention funding=%s effect=%s attempts=%d next_attempt=%d cause=%s", current.Status, current.EffectStatus, current.Attempts, current.NextAttempt, current.LastError) + if current.Status == model.BillingSettlementStatusPending { + require.Positive(t, current.Attempts) + require.True(t, strings.Contains(current.LastError, "SQLITE_BUSY") || strings.Contains(current.LastError, "database is locked"), + "only the observed transient writer contention may be scheduled") + assertBillingProcessLedger(t, db, funding, quota, "intent_committed", false) + // Honor the production durable backoff, rather than treating an + // intentionally not-yet-due intent as a failed financial recovery. + // One scheduled attempt is enough after all competing writers exit; + // any further error must fail the test instead of retrying to green. + wait := time.Until(time.Unix(current.NextAttempt, 0)) + require.Positive(t, current.NextAttempt) + require.LessOrEqual(t, wait, 64*time.Second) + if wait > 0 { + timer := time.NewTimer(wait) + <-timer.C + } + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + } + assertBillingProcessLedger(t, db, funding, quota, "complete", false) + require.Equal(t, original.EffectPayload, billingProcessRecord(t, db).EffectPayload) + snapshot := billingProcessSnapshot(t, db) + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + require.Equal(t, snapshot, billingProcessSnapshot(t, db)) + }) + } + } +} + +func TestBillingProcessUnsupportedStateIsNotRetried(t *testing.T) { + for _, funding := range billingProcessFundingCases()[:2] { + t.Run(funding.name, func(t *testing.T) { + dir, db := seedBillingProcessDB(t, funding) + cfg := billingProcessConfig{Role: "work", Phase: "intent_committed", Preference: funding.preference, Source: funding.source, Quota: 130} + killBillingTestProcess(t, startBillingTestProcess(t, dir, cfg), cfg.Phase) + assertBillingProcessLedger(t, db, funding, cfg.Quota, cfg.Phase, false) + // Inject an unsupported durable state. This tests the runner's allowlist, + // not a real provider's outcome-unknown protocol or commit ambiguity. + record := billingProcessRecord(t, db) + require.NoError(t, db.Model(&record).Update("status", "outcome_unknown").Error) + snapshot := billingProcessSnapshot(t, db) + cfg.Role, cfg.Phase = "recover", "" + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + require.Equal(t, snapshot, billingProcessSnapshot(t, db), "unsupported state must remain durably inspectable without automatic balance changes") + }) + } +} + +func TestBillingProcessSubscriptionPeriodRecovery(t *testing.T) { + funding := billingProcessFundingCases()[1] + for _, phase := range []string{"intent_committed", "funding_committed"} { + t.Run(phase, func(t *testing.T) { + dir, db := seedBillingProcessDB(t, funding) + cfg := billingProcessConfig{Role: "work", Phase: phase, Preference: funding.preference, Source: funding.source, Quota: 130} + killBillingTestProcess(t, startBillingTestProcess(t, dir, cfg), phase) + original := billingProcessRecord(t, db) + if phase == "funding_committed" { + // Applied funding permits the real reset API. Only effects of the + // original period may recover afterward, without charging the new one. + cfg.Role, cfg.Phase = "reset", "" + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + } else { + // A completed reset fixture puts the original pending intent in an + // older period. It does not assert that the admin API allows such a + // reset while same-period funding is unresolved. + var sub model.UserSubscription + require.NoError(t, db.First(&sub, 955).Error) + require.NoError(t, db.Model(&sub).Updates(map[string]interface{}{ + "last_reset_time": sub.LastResetTime + 1, "amount_used": 0, + }).Error) + } + cfg.Role, cfg.Phase = "recover", "" + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + record := billingProcessRecord(t, db) + require.Equal(t, original.EffectPayload, record.EffectPayload) + var sub model.UserSubscription + var user model.User + var token model.Token + var channel model.Channel + require.NoError(t, db.First(&sub, 955).Error) + require.NoError(t, db.First(&user, 951).Error) + require.NoError(t, db.First(&token, 952).Error) + require.NoError(t, db.First(&channel, 953).Error) + require.Zero(t, sub.AmountUsed, "the new subscription period cannot pay for the original request") + require.EqualValues(t, 1000, user.Quota) + var logs []model.Log + var receipts []model.BillingLogReceipt + require.NoError(t, db.Find(&logs).Error) + require.NoError(t, db.Find(&receipts).Error) + if phase == "intent_committed" { + require.Equal(t, model.BillingSettlementStatusManual, record.Status) + require.Contains(t, record.LastError, model.ErrSubscriptionSettlementPeriodChanged.Error()) + require.Zero(t, record.AppliedFundingDelta) + require.Zero(t, record.AppliedTokenDelta) + require.EqualValues(t, 900, token.RemainQuota) + require.EqualValues(t, 100, token.UsedQuota) + require.Zero(t, user.UsedQuota) + require.Zero(t, user.RequestCount) + require.Zero(t, channel.UsedQuota) + require.Empty(t, logs) + require.Empty(t, receipts) + } else { + require.Equal(t, model.BillingSettlementStatusApplied, record.Status) + require.Equal(t, model.BillingSettlementEffectApplied, record.EffectStatus) + require.EqualValues(t, 30, record.AppliedFundingDelta) + require.EqualValues(t, 30, record.AppliedTokenDelta) + require.EqualValues(t, 870, token.RemainQuota) + require.EqualValues(t, 130, token.UsedQuota) + require.EqualValues(t, 130, user.UsedQuota) + require.EqualValues(t, 1, user.RequestCount) + require.EqualValues(t, 130, channel.UsedQuota) + require.Len(t, logs, 1) + require.Len(t, receipts, 1) + var other map[string]interface{} + require.NoError(t, common.UnmarshalJsonStr(logs[0].Other, &other)) + require.EqualValues(t, 130, other["subscription_used"], "logs must retain the original period's usage") + require.EqualValues(t, 870, other["subscription_remain"]) + } + snapshot := billingProcessSnapshot(t, db) + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + require.Equal(t, snapshot, billingProcessSnapshot(t, db)) + }) + } +} From 47eadf517b9633463aa39dd072f924028a1cdca6 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Mon, 12 Oct 2026 01:58:16 +0800 Subject: [PATCH 10/24] fix: estimate delivered native Responses output usage --- relay/channel/openai/relay_responses.go | 122 +++++++- ...s_native_missing_output_settlement_test.go | 296 ++++++++++++++++++ relay/subscription_stream_settlement_test.go | 4 +- 3 files changed, 419 insertions(+), 3 deletions(-) create mode 100644 relay/responses_native_missing_output_settlement_test.go diff --git a/relay/channel/openai/relay_responses.go b/relay/channel/openai/relay_responses.go index 1e4230c2..13945ff1 100644 --- a/relay/channel/openai/relay_responses.go +++ b/relay/channel/openai/relay_responses.go @@ -20,6 +20,122 @@ import ( const maxPendingResponsesStreamEvents = 32 +// Collect generated output only after its complete SSE frame was delivered. +// Keep tool snapshots separate from deltas so item.done cannot count the same +// arguments twice. This estimate never replaces upstream-reported token counts. +type responsesOutputEstimate struct { + text strings.Builder + tools []*responsesToolEstimate + toolsByAlias map[string]*responsesToolEstimate +} + +type responsesToolEstimate struct { + name string + deltas strings.Builder + snapshot string +} + +func (t *responsesToolEstimate) observeSnapshot(snapshot string) { + // Preserve already delivered content if a later malformed snapshot is + // empty or contradicts it; identical/full extensions are counted once. + if strings.HasPrefix(snapshot, t.snapshot) { + t.snapshot = snapshot + } +} + +func (e *responsesOutputEstimate) tool(event dto.ResponsesStreamResponse) *responsesToolEstimate { + aliases := make([]string, 0, 4) + if event.OutputIndex != nil { + aliases = append(aliases, fmt.Sprintf("index:%d", *event.OutputIndex)) + } + if event.ItemID != "" { + aliases = append(aliases, "id:"+event.ItemID) + } + if event.Item != nil { + if event.Item.ID != "" { + aliases = append(aliases, "id:"+event.Item.ID) + } + if event.Item.CallId != "" { + aliases = append(aliases, "call:"+event.Item.CallId) + } + } + if len(aliases) == 0 { + // Compatibility providers may omit both identifiers. Keep their single + // function/custom argument stream distinct instead of inventing IDs. + kind := dto.BuildInCallFunctionCall + if strings.HasPrefix(event.Type, "response.custom_tool_call_input.") || + (event.Item != nil && event.Item.Type == dto.BuildInCallCustomToolCall) { + kind = dto.BuildInCallCustomToolCall + } + aliases = append(aliases, "unidentified:"+kind) + } + if e.toolsByAlias == nil { + e.toolsByAlias = make(map[string]*responsesToolEstimate) + } + var tool *responsesToolEstimate + for _, alias := range aliases { + if tool = e.toolsByAlias[alias]; tool != nil { + break + } + } + if tool == nil { + tool = &responsesToolEstimate{} + e.tools = append(e.tools, tool) + } + for _, alias := range aliases { + e.toolsByAlias[alias] = tool + } + return tool +} + +func (e *responsesOutputEstimate) observe(event dto.ResponsesStreamResponse) { + switch event.Type { + case "response.output_text.delta", "response.reasoning_text.delta", "response.reasoning_summary_text.delta", "response.refusal.delta": + e.text.WriteString(event.Delta) + case "response.function_call_arguments.delta", "response.custom_tool_call_input.delta": + e.tool(event).deltas.WriteString(event.Delta) + case "response.function_call_arguments.done": + if event.Arguments != nil { + e.tool(event).observeSnapshot(*event.Arguments) + } + case "response.custom_tool_call_input.done": + if event.Input != nil { + e.tool(event).observeSnapshot(*event.Input) + } + case "response.output_item.added", "response.output_item.done": + if event.Item != nil && (event.Item.Type == dto.BuildInCallFunctionCall || event.Item.Type == dto.BuildInCallCustomToolCall) { + tool := e.tool(event) + if tool.name == "" { + tool.name = event.Item.Name + } + tool.observeSnapshot(event.Item.ArgumentsString()) + } + case "response.completed", "response.incomplete", "response.done": + if event.Response != nil { + for i, item := range event.Response.Output { + e.observe(dto.ResponsesStreamResponse{Type: "response.output_item.done", Item: &item, OutputIndex: &i}) + } + } + } +} + +func (e *responsesOutputEstimate) textForTokens() string { + var text strings.Builder + text.WriteString(e.text.String()) + for _, tool := range e.tools { + text.WriteString(tool.name) + arguments := tool.deltas.String() + // A full delivered snapshot may include a suffix that never arrived as + // a delta. Only extend matching prefixes; retain delivered deltas if a + // malformed snapshot contradicts them, without charging both copies. + if strings.HasPrefix(tool.snapshot, arguments) { + arguments = tool.snapshot + } + text.WriteString(arguments) + } + return text.String() +} + func applyResponsesUsage(dst *dto.Usage, src *dto.Usage) { if dst == nil || src == nil { return @@ -248,6 +364,7 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp var usage = &dto.Usage{} var responseTextBuilder strings.Builder + var outputEstimate responsesOutputEstimate var streamErr *types.MaxAPIError type pendingResponsesStreamEvent struct { response dto.ResponsesStreamResponse @@ -265,6 +382,7 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp delivered, writeErr := sendResponsesStreamData(c, event.response, event.data) if delivered { streamForwarded = true + outputEstimate.observe(event.response) if helper.ResponsesStreamEventHasOutput(event.response) || responsesStreamHasVisibleOutput(&event.response) { outputDelivered = true } @@ -283,8 +401,8 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp } finalizeUsage := func() { if !usageReported { - if usage.CompletionTokens == 0 && responseTextBuilder.Len() > 0 { - usage.CompletionTokens = service.CountTextToken(responseTextBuilder.String(), info.UpstreamModelName) + if usage.CompletionTokens == 0 { + usage.CompletionTokens = service.CountTextToken(outputEstimate.textForTokens(), info.UpstreamModelName) } if usage.PromptTokens == 0 && usage.CompletionTokens != 0 { usage.PromptTokens = info.GetEstimatePromptTokens() diff --git a/relay/responses_native_missing_output_settlement_test.go b/relay/responses_native_missing_output_settlement_test.go new file mode 100644 index 00000000..7e09d99b --- /dev/null +++ b/relay/responses_native_missing_output_settlement_test.go @@ -0,0 +1,296 @@ +package relay + +import ( + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/constant" + "github.com/MAX-API-Next/MAX-API/model" + "github.com/MAX-API-Next/MAX-API/pkg/billingexpr" + "github.com/MAX-API-Next/MAX-API/service" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type nativeMissingOutputCase struct { + name, estimate string + frames []map[string]any +} + +func nativeMissingOutputCases() []nativeMissingOutputCase { + delta := func(kind, text string) map[string]any { + return map[string]any{"type": "response." + kind + ".delta", "delta": text, "item_id": "item_0", "output_index": 0} + } + item := func(event, kind, id, name, payload string, index int) map[string]any { + output := map[string]any{"type": kind, "id": id, "name": name} + if kind == "custom_tool_call" { + output["input"] = payload + } else { + output["arguments"] = payload + } + return map[string]any{"type": "response.output_item." + event, "output_index": index, "item": output} + } + const arguments = `{"q":"Paris"}` + const input = "print('generated output')" + cases := []nativeMissingOutputCase{} + for _, kind := range []string{"output_text", "reasoning_text", "reasoning_summary_text", "refusal", "function_call_arguments", "custom_tool_call_input"} { + cases = append(cases, nativeMissingOutputCase{kind, "generated output", []map[string]any{delta(kind, "generated output")}}) + } + cases = append(cases, + nativeMissingOutputCase{"function_snapshot", "lookup" + arguments, []map[string]any{item("done", "function_call", "item_0", "lookup", arguments, 0)}}, + nativeMissingOutputCase{"function_snapshots_once", "lookup" + arguments, []map[string]any{ + item("added", "function_call", "item_0", "lookup", arguments, 0), item("done", "function_call", "item_0", "lookup", arguments, 0)}}, + nativeMissingOutputCase{"snapshot_retained_after_empty", "lookup" + arguments, []map[string]any{ + item("added", "function_call", "item_0", "lookup", arguments, 0), item("done", "function_call", "item_0", "lookup", "", 0)}}, + nativeMissingOutputCase{"function_object_arguments", "lookup" + arguments, []map[string]any{ + {"type": "response.output_item.done", "output_index": 0, "item": map[string]any{"type": "function_call", "id": "item_0", "name": "lookup", "arguments": map[string]any{"q": "Paris"}}}}}, + nativeMissingOutputCase{"terminal_function_snapshot", "lookup" + arguments, []map[string]any{ + {"type": "response.completed", "response": map[string]any{"id": "resp_missing", "status": "completed", "output": []any{item("done", "function_call", "item_0", "lookup", arguments, 0)["item"]}}}}}, + nativeMissingOutputCase{"function_delta_snapshot_once", "lookup" + arguments, []map[string]any{ + item("added", "function_call", "item_0", "lookup", "", 0), delta("function_call_arguments", arguments), item("done", "function_call", "item_0", "lookup", arguments, 0)}}, + nativeMissingOutputCase{"function_snapshot_completes_delta", "lookup" + arguments, []map[string]any{ + item("added", "function_call", "item_0", "lookup", "", 0), delta("function_call_arguments", `{"q":`), item("done", "function_call", "item_0", "lookup", arguments, 0)}}, + nativeMissingOutputCase{"custom_delta_snapshot_once", "python" + input, []map[string]any{ + item("added", "custom_tool_call", "item_0", "python", "", 0), delta("custom_tool_call_input", input), item("done", "custom_tool_call", "item_0", "python", input, 0)}}, + nativeMissingOutputCase{"tool_name_only", "lookup", []map[string]any{item("added", "function_call", "item_0", "lookup", "", 0)}}, + nativeMissingOutputCase{"function_arguments_done", "lookup" + arguments, []map[string]any{ + item("added", "function_call", "item_0", "lookup", "", 0), + {"type": "response.function_call_arguments.done", "item_id": "item_0", "output_index": 0, "arguments": arguments}}}, + nativeMissingOutputCase{"custom_input_done", "python" + input, []map[string]any{ + item("added", "custom_tool_call", "item_0", "python", "", 0), + {"type": "response.custom_tool_call_input.done", "item_id": "item_0", "output_index": 0, "input": input}}}, + nativeMissingOutputCase{"two_tools_zero_index", "lookup" + arguments + "python" + input, []map[string]any{ + item("added", "function_call", "item_0", "lookup", "", 0), delta("function_call_arguments", arguments), + item("added", "custom_tool_call", "item_1", "python", "", 1), + {"type": "response.custom_tool_call_input.delta", "item_id": "item_1", "delta": input}, + item("done", "function_call", "item_0", "lookup", arguments, 0), item("done", "custom_tool_call", "item_1", "python", input, 1)}}, + nativeMissingOutputCase{"mixed_text_tool", "visible text" + "lookup" + arguments, []map[string]any{ + delta("output_text", "visible text"), item("added", "function_call", "item_1", "lookup", "", 1), + {"type": "response.function_call_arguments.delta", "output_index": 1, "item_id": "item_1", "delta": arguments}}}, + nativeMissingOutputCase{"large_delta_above_reservation", strings.Repeat("generated output ", 150), []map[string]any{delta("custom_tool_call_input", strings.Repeat("generated output ", 150))}}, + ) + return cases +} + +// Fail one complete SSE frame, while allowing previously delivered frames. +type nativeMissingOutputWriter struct { + *httptest.ResponseRecorder + short, triggered bool +} + +func (w *nativeMissingOutputWriter) Write(data []byte) (int, error) { + if strings.Contains(string(data), "unwritten-generated-output") { + w.triggered = true + if w.short { + return w.ResponseRecorder.Write(data[:3]) + } + return 0, io.ErrClosedPipe + } + return w.ResponseRecorder.Write(data) +} + +func (w *nativeMissingOutputWriter) WriteString(data string) (int, error) { + return w.Write([]byte(data)) +} + +func TestNativeResponsesMissingOutputRetainsOriginalFunding(t *testing.T) { + initRequestReplayTests(t) + oldAudit := common.LogResponseContentEnabled + common.LogResponseContentEnabled = true + t.Cleanup(func() { common.LogResponseContentEnabled = oldAudit }) + for _, scenario := range nativeMissingOutputCases() { + for _, funding := range streamFundingCases() { + for _, usage := range []string{"missing", "known", "zero", "terminal_zero"} { + t.Run(scenario.name+"/"+funding.name+"/"+usage, func(t *testing.T) { + runNativeMissingOutputLedger(t, scenario, funding, usage, "eof", false) + }) + } + } + } + for _, ending := range []string{"malformed", "transport", "terminal_failed", "write", "short"} { + for _, funding := range streamFundingCases() { + for _, kind := range []string{"function_call_arguments", "custom_tool_call_input", "reasoning_summary_text", "refusal"} { + scenario := nativeMissingOutputCase{kind, "generated output", []map[string]any{ + {"type": "response." + kind + ".delta", "item_id": "item_0", "output_index": 0, "delta": "generated output"}, + }} + t.Run(kind+"/"+funding.name+"/missing/"+ending, func(t *testing.T) { + runNativeMissingOutputLedger(t, scenario, funding, "missing", ending, false) + }) + if ending == "write" || ending == "short" { + t.Run(kind+"/"+funding.name+"/no_delivery/"+ending, func(t *testing.T) { + runNativeMissingOutputLedger(t, scenario, funding, "missing", ending, true) + }) + } + } + } + } + for _, funding := range streamFundingCases() { + scenario := nativeMissingOutputCases()[4] + t.Run("missing_function/"+funding.name+"/at_reservation", func(t *testing.T) { + runNativeMissingOutputLedger(t, scenario, funding, "at_reservation", "eof", false) + }) + for _, pricing := range []string{"free", "tiered_zero"} { + scenario := nativeMissingOutputCases()[4] + t.Run("missing_function/"+funding.name+"/"+pricing, func(t *testing.T) { + runNativeMissingOutputLedger(t, scenario, funding, "missing", "eof", false, pricing) + }) + } + } +} + +// Real HTTP -> native adaptor -> Helper -> durable funding and effect seam. +// Assertions cover exact amounts, original funding, logs, all mirrors and replay. +func runNativeMissingOutputLedger(t *testing.T, scenario nativeMissingOutputCase, funding streamFundingCase, usage, ending string, noDelivery bool, pricing ...string) { + t.Helper() + db := setupStreamFundingLedger(t, funding) + response := map[string]any{"id": "resp_missing", "model": "gpt-test", "status": "in_progress", "output": []any{}} + if usage == "known" { + response["usage"] = map[string]any{"input_tokens": 4, "output_tokens": 2, "total_tokens": 6} + } else if usage == "at_reservation" { + response["usage"] = map[string]any{"input_tokens": 80, "output_tokens": 20, "total_tokens": 100} + } else if usage == "zero" { + response["usage"] = map[string]any{"input_tokens": 0, "output_tokens": 0, "total_tokens": 0} + } + encode := func(v any) string { + data, err := common.Marshal(v) + require.NoError(t, err) + return string(data) + } + frames := []string{encode(map[string]any{"type": "response.created", "response": response})} + hasTerminal := false + if !noDelivery { + for _, frame := range scenario.frames { + if frame["type"] == "response.completed" { + hasTerminal = true + if usage == "terminal_zero" { + // A terminal snapshot supplies usage in that one terminal; + // emitting a second response.completed violates the protocol. + var copy map[string]any + require.NoError(t, common.UnmarshalJsonStr(encode(frame), ©)) + copy["response"].(map[string]any)["usage"] = map[string]any{"input_tokens": 0, "output_tokens": 0, "total_tokens": 0} + frame = copy + } + } + frames = append(frames, encode(frame)) + } + } + if usage == "terminal_zero" && !hasTerminal { + frames = append(frames, `{"type":"response.completed","response":{"id":"resp_missing","status":"completed","usage":{"input_tokens":0,"output_tokens":0,"total_tokens":0}}}`) + } + switch ending { + case "malformed": + frames = append(frames, `{broken`) + case "terminal_failed": + frames = append(frames, `{"type":"response.failed","response":{"id":"resp_missing","status":"failed","error":{"code":"server_error","message":"synthetic failure"}}}`) + case "write", "short": + frames = append(frames, `{"type":"response.`+scenario.name+`.delta","item_id":"item_0","output_index":0,"delta":"unwritten-generated-output"}`) + } + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + if ending == "transport" { + w.Header().Set("Content-Length", "100000") + } + _, _ = fmt.Fprint(w, "data: "+strings.Join(frames, "\n\ndata: ")+"\n\n") + })) + t.Cleanup(upstream.Close) + requestID := "native-missing-" + fmt.Sprintf("%x", len(t.Name())) + "-" + scenario.name + c, info, recorder := streamFundingContext(t, "responses_native", requestID, upstream.URL, constant.ChannelTypeOpenAI, funding) + if len(pricing) > 0 { + if pricing[0] == "free" { + info.PriceData.ModelRatio = 0 + } else { + info.TieredBillingSnapshot = &billingexpr.BillingSnapshot{BillingMode: "tiered_expr", ExprString: "0", ExprHash: billingexpr.ExprHashString("0"), GroupRatio: 1, QuotaPerUnit: common.QuotaPerUnit, EstimatedQuotaAfterGroup: 100} + } + } + var writer *nativeMissingOutputWriter + if ending == "write" || ending == "short" { + writer = &nativeMissingOutputWriter{ResponseRecorder: recorder, short: ending == "short"} + failed, _ := gin.CreateTestContext(writer) + c.Writer = failed.Writer + } + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + apiErr := ResponsesHelper(c, info) + legalTerminal := usage == "terminal_zero" || hasTerminal + if legalTerminal && ending == "eof" { + require.Nil(t, apiErr) + } else { + require.NotNil(t, apiErr) + } + if writer != nil { + require.True(t, writer.triggered) + require.NotContains(t, recorder.Body.String(), "unwritten-generated-output") + } + service.HandleFailedBilling(c, info, apiErr) + service.HandleFailedBilling(c, info, apiErr) + require.EqualValues(t, 1, calls.Load()) + prompt, completion := 10, service.CountTextToken(scenario.estimate, "gpt-test") + if usage == "known" { + prompt, completion = 4, 2 + } else if usage == "at_reservation" { + prompt, completion = 80, 20 + } else if usage == "zero" || usage == "terminal_zero" || noDelivery { + prompt, completion = 0, 0 + } + quota := prompt + completion + if scenario.name == "large_delta_above_reservation" && usage == "missing" { + require.Greater(t, quota, 100, "exercise a positive funding delta, not just a refund") + } + if len(pricing) > 0 { + quota = 0 + } + var record model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).First(&record).Error) + require.Equal(t, funding.source, record.Source) + require.Equal(t, model.BillingSettlementStatusApplied, record.Status) + require.EqualValues(t, quota-100, record.FundingDelta) + require.EqualValues(t, quota-100, record.TokenDelta) + require.EqualValues(t, quota-100, record.AppliedFundingDelta) + require.EqualValues(t, quota-100, record.AppliedTokenDelta) + var effect *model.BillingSettlementEffect + if !noDelivery { + require.Equal(t, model.BillingSettlementEffectApplied, record.EffectStatus) + effect = &model.BillingSettlementEffect{} + require.NoError(t, common.UnmarshalJsonStr(record.EffectPayload, effect)) + require.Equal(t, prompt, effect.PromptTokens) + require.Equal(t, completion, effect.CompletionTokens) + require.EqualValues(t, quota, effect.Quota) + require.True(t, effect.QuotaIsActual) + } + assertStreamFundingBalances(t, db, funding, requestID, quota, !noDelivery, prompt, completion) + if ending == "eof" { + var audit strings.Builder + for _, frame := range scenario.frames { + if frame["type"] == "response.output_text.delta" { + audit.WriteString(frame["delta"].(string)) + } + } + require.Equal(t, strings.TrimSpace(audit.String()), info.AuditResponseContent, "audit text must exclude tool arguments, refusal and reasoning estimates") + } + before := requestReplaySnapshot(t, db) + for i := 0; i < 2; i++ { + delta, alreadyApplied, err := model.ApplyBillingSettlementOnce(requestReplayInput(record, effect)) + require.NoError(t, err) + require.True(t, alreadyApplied) + require.EqualValues(t, quota-100, delta) + if effect != nil { + require.NoError(t, model.ProcessBillingSettlementEffect(record.OperationKey)) + } + service.HandleFailedBilling(c, info, apiErr) + } + model.ProcessPendingBillingSettlementsOnce() + require.Equal(t, before, requestReplaySnapshot(t, db), "replay/recovery must preserve every ledger table and original payload") + var logs []model.Log + require.NoError(t, db.Find(&logs).Error) + if !noDelivery { + require.Len(t, logs, 1) + require.False(t, strings.Contains(logs[0].Content, "generated output"), "billing estimate must not copy tool or reasoning content into the audit log") + } +} diff --git a/relay/subscription_stream_settlement_test.go b/relay/subscription_stream_settlement_test.go index 9565245f..a94f30b3 100644 --- a/relay/subscription_stream_settlement_test.go +++ b/relay/subscription_stream_settlement_test.go @@ -340,7 +340,9 @@ func assertStreamFundingBalances(t *testing.T, db *gorm.DB, funding streamFundin wantLogs, wantRequests := 0, 0 if delivered { wantLogs = 1 - if quota != 0 { + // Free pricing and a zero-valued expression still count a request with + // measured generation; explicit numeric zero usage does not. + if quota != 0 || prompt+completion > 0 { wantRequests = 1 } } From f58f0823ed704fb6eeb48899cd4d2768d0c1543d Mon Sep 17 00:00:00 2001 From: CSCITech Date: Mon, 12 Oct 2026 12:42:08 +0800 Subject: [PATCH 11/24] fix: preserve upstream protocol in global Responses bridge --- relay/chat_completions_via_responses.go | 15 ++ ...chat_responses_protocol_settlement_test.go | 245 ++++++++++++++++++ relay/compatible_handler.go | 24 +- ...s_native_missing_output_settlement_test.go | 8 +- service/openai_chat_responses_mode.go | 4 + service/openaicompat/policy.go | 24 +- 6 files changed, 313 insertions(+), 7 deletions(-) create mode 100644 relay/chat_responses_protocol_settlement_test.go diff --git a/relay/chat_completions_via_responses.go b/relay/chat_completions_via_responses.go index ec8d79df..0f4a376b 100644 --- a/relay/chat_completions_via_responses.go +++ b/relay/chat_completions_via_responses.go @@ -1,6 +1,7 @@ package relay import ( + "fmt" "io" "net/http" "strings" @@ -70,6 +71,9 @@ func applySystemPromptIfNeeded(c *gin.Context, info *relaycommon.RelayInfo, requ } func chatCompletionsViaResponses(c *gin.Context, info *relaycommon.RelayInfo, adaptor channel.Adaptor, request *dto.GeneralOpenAIRequest) (*dto.Usage, *types.MaxAPIError) { + if info == nil || adaptor == nil || !service.SupportsOpenAIResponsesUpstream(info.ChannelType) { + return nil, types.NewError(fmt.Errorf("channel does not use OpenAI Responses upstream"), types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } chatJSON, err := common.Marshal(request) if err != nil { return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) @@ -112,6 +116,17 @@ func chatCompletionsViaResponses(c *gin.Context, info *relaycommon.RelayInfo, ad if err != nil { return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } + // An inconsistent adaptor must fail before HTTP, never send a native + // Messages/GenerateContent payload to the hardcoded Responses reply parser. + switch converted := convertedRequest.(type) { + case dto.OpenAIResponsesRequest: + case *dto.OpenAIResponsesRequest: + if converted == nil { + return nil, types.NewError(fmt.Errorf("nil OpenAI Responses upstream request"), types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } + default: + return nil, types.NewError(fmt.Errorf("unexpected Responses upstream request type %T", convertedRequest), types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } convertedRequest, err = filterResponsesRequestFields(convertedRequest, info, adaptor) if err != nil { return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) diff --git a/relay/chat_responses_protocol_settlement_test.go b/relay/chat_responses_protocol_settlement_test.go new file mode 100644 index 00000000..bb854d3d --- /dev/null +++ b/relay/chat_responses_protocol_settlement_test.go @@ -0,0 +1,245 @@ +package relay + +import ( + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/constant" + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/model" + "github.com/MAX-API-Next/MAX-API/relay/channel/claude" + relayconstant "github.com/MAX-API-Next/MAX-API/relay/constant" + "github.com/MAX-API-Next/MAX-API/service" + "github.com/MAX-API-Next/MAX-API/setting/model_setting" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/stretchr/testify/require" +) + +func TestGlobalChatResponsesProtocolRetainsOriginalFunding(t *testing.T) { + initRequestReplayTests(t) + providers := []struct { + name, wire, converter string + channelType int + }{ + {"claude", "claude", "", constant.ChannelTypeAnthropic}, + {"gemini", "gemini", "", constant.ChannelTypeGemini}, + {"openai", "responses", "", constant.ChannelTypeOpenAI}, + {"custom_none", "responses", dto.AdvancedCustomConverterNone, constant.ChannelTypeAdvancedCustom}, + {"custom_claude", "claude", dto.AdvancedCustomConverterOpenAIChatCompletionsToAnthropicMessages, constant.ChannelTypeAdvancedCustom}, + {"custom_gemini", "gemini", dto.AdvancedCustomConverterOpenAIChatCompletionsToGeminiGenerateContent, constant.ChannelTypeAdvancedCustom}, + {"custom_responses", "responses", dto.AdvancedCustomConverterOpenAIChatCompletionsToOpenAIResponses, constant.ChannelTypeAdvancedCustom}, + } + for _, provider := range providers { + modes := []bool{true} + if provider.name == "claude" || provider.name == "gemini" { + modes = []bool{false, true} + } + for _, enabled := range modes { + for _, stream := range []bool{false, true} { + for _, funding := range streamFundingCases() { + t.Run(fmt.Sprintf("%s/global_%t/stream_%t/%s", provider.name, enabled, stream, funding.name), func(t *testing.T) { + db := setupStreamFundingLedger(t, funding) + body := globalBridgeWireBody(t, provider.wire, stream) + type observedRequest struct { + path string + body map[string]any + err error + } + requests := make(chan observedRequest, 1) + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + data, err := io.ReadAll(r.Body) + var parsed map[string]any + if err == nil { + err = common.Unmarshal(data, &parsed) + } + requests <- observedRequest{r.URL.Path, parsed, err} + if stream { + w.Header().Set("Content-Type", "text/event-stream") + } else { + w.Header().Set("Content-Type", "application/json") + } + _, _ = io.WriteString(w, body) + })) + t.Cleanup(upstream.Close) + requestID := "global-protocol-" + provider.name + c, info, recorder := streamFundingContext(t, "chat_native", requestID, upstream.URL, provider.channelType, funding) + info.IsStream = stream + info.Request.(*dto.GeneralOpenAIRequest).Stream = common.GetPointer(stream) + model_setting.GetGlobalSettings().ChatCompletionsToResponsesPolicy = model_setting.ChatCompletionsToResponsesPolicy{Enabled: enabled, AllChannels: true, ModelPatterns: []string{"^gpt-test$"}} + if provider.channelType == constant.ChannelTypeAdvancedCustom { + path := "/v1/responses" + if provider.wire == "claude" { + path = "/v1/messages" + } else if provider.wire == "gemini" { + path = "/v1beta/models/{model}:generateContent" + } + c.Set(string(constant.ContextKeyChannelOtherSetting), dto.ChannelOtherSettings{AdvancedCustom: &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{IncomingPath: "/v1/chat/completions", UpstreamPath: path, Converter: provider.converter}}}}) + } + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + apiErr := TextHelper(c, info) + require.Nil(t, apiErr, "global policy must not send native replies to a Responses parser") + require.Contains(t, recorder.Body.String(), "partial output") + require.Contains(t, recorder.Body.String(), "chat.completion") + require.EqualValues(t, 1, calls.Load()) + select { + case request := <-requests: + require.NoError(t, request.err) + switch provider.wire { + case "claude": + require.Contains(t, request.path, "/messages") + require.Contains(t, request.body, "messages") + require.NotContains(t, request.body, "input") + case "gemini": + require.Contains(t, strings.ToLower(request.path), "generatecontent") + require.Contains(t, request.body, "contents") + require.NotContains(t, request.body, "input") + default: + require.Equal(t, "/v1/responses", request.path) + require.Contains(t, request.body, "input") + require.NotContains(t, request.body, "messages") + } + case <-time.After(3 * time.Second): + t.Fatal("no observed synthetic upstream request") + } + assertStreamFundingBalances(t, db, funding, requestID, 6, true, 4, 2) + var record model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).First(&record).Error) + require.Equal(t, funding.source, record.Source) + require.EqualValues(t, -94, record.AppliedFundingDelta) + require.EqualValues(t, -94, record.AppliedTokenDelta) + var effect model.BillingSettlementEffect + require.NoError(t, common.UnmarshalJsonStr(record.EffectPayload, &effect)) + before := requestReplaySnapshot(t, db) + for i := 0; i < 2; i++ { + delta, already, err := model.ApplyBillingSettlementOnce(requestReplayInput(record, &effect)) + require.NoError(t, err) + require.True(t, already) + require.EqualValues(t, -94, delta) + require.NoError(t, model.ProcessBillingSettlementEffect(record.OperationKey)) + service.HandleFailedBilling(c, info, apiErr) + } + model.ProcessPendingBillingSettlementsOnce() + require.Equal(t, before, requestReplaySnapshot(t, db)) + }) + } + } + } + } +} + +// Exercise real adaptor conversion rather than duplicating the capability list +// in a test. A future adaptor change must update its global bridge eligibility. +func TestGlobalResponsesBridgeEligibilityMatchesAdaptorProtocol(t *testing.T) { + initRequestReplayTests(t) + for channelType := 1; channelType < constant.ChannelTypeDummy; channelType++ { + t.Run(fmt.Sprintf("channel_%d", channelType), func(t *testing.T) { + funding := streamFundingCases()[3] + c, info, _ := streamFundingContext(t, "chat_native", "capability-probe", "https://synthetic.invalid", channelType, funding) + if channelType == constant.ChannelTypeAdvancedCustom { + c.Set(string(constant.ContextKeyChannelOtherSetting), dto.ChannelOtherSettings{AdvancedCustom: &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{IncomingPath: "/v1/chat/completions", UpstreamPath: "/v1/responses", Converter: dto.AdvancedCustomConverterNone}}}}) + } + info.InitChannelMeta(c) + info.RelayMode = relayconstant.RelayModeResponses + info.RequestURLPath = "/v1/responses" + adaptor := GetAdaptor(info.ApiType) + expected := false + if adaptor != nil { + adaptor.Init(info) + converted, err := adaptor.ConvertOpenAIResponsesRequest(c, info, dto.OpenAIResponsesRequest{Model: "gpt-test", Input: []byte(`"test"`)}) + if err == nil { + switch value := converted.(type) { + case dto.OpenAIResponsesRequest: + expected = true + case *dto.OpenAIResponsesRequest: + expected = value != nil + } + } + } + for _, policy := range []model_setting.ChatCompletionsToResponsesPolicy{ + {Enabled: true, AllChannels: true, ModelPatterns: []string{"^gpt-test$"}}, + {Enabled: true, ChannelIDs: []int{953}, ModelPatterns: []string{"^gpt-test$"}}, + {Enabled: true, ChannelTypes: []int{channelType}, ModelPatterns: []string{"^gpt-test$"}}, + } { + require.Equal(t, expected, service.ShouldChatCompletionsUseResponsesPolicy(policy, 953, channelType, "gpt-test")) + require.False(t, service.ShouldChatCompletionsUseResponsesPolicy(policy, 953, channelType, "other-model")) + policy.Enabled = false + require.False(t, service.ShouldChatCompletionsUseResponsesPolicy(policy, 953, channelType, "gpt-test")) + } + }) + } +} + +func TestGlobalResponsesBridgeRejectsNativeBeforeUpstreamAndRefunds(t *testing.T) { + initRequestReplayTests(t) + for _, funding := range streamFundingCases() { + for _, mismatch := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/mismatched_adaptor_%t", funding.name, mismatch), func(t *testing.T) { + db := setupStreamFundingLedger(t, funding) + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { calls.Add(1) })) + t.Cleanup(upstream.Close) + channelType := constant.ChannelTypeAnthropic + if mismatch { + channelType = constant.ChannelTypeOpenAI + } + requestID := "rejected-global-protocol" + c, info, recorder := streamFundingContext(t, "chat_native", requestID, upstream.URL, channelType, funding) + info.InitChannelMeta(c) + info.RequestURLPath = "/v1/chat/completions" + adaptor := &claude.Adaptor{} + adaptor.Init(info) + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + usage, apiErr := chatCompletionsViaResponses(c, info, adaptor, info.Request.(*dto.GeneralOpenAIRequest)) + require.Nil(t, usage) + require.NotNil(t, apiErr) + require.Equal(t, types.ErrorCodeConvertRequestFailed, apiErr.GetErrorCode()) + require.True(t, types.IsSkipRetryError(apiErr)) + require.Zero(t, calls.Load()) + require.Empty(t, recorder.Body.String()) + require.Equal(t, relayconstant.RelayModeChatCompletions, info.RelayMode) + require.Equal(t, "/v1/chat/completions", info.RequestURLPath) + service.HandleFailedBilling(c, info, apiErr) + service.HandleFailedBilling(c, info, apiErr) + assertStreamFundingBalances(t, db, funding, requestID, 0, false, 0, 0) + before := requestReplaySnapshot(t, db) + model.ProcessPendingBillingSettlementsOnce() + require.Equal(t, before, requestReplaySnapshot(t, db)) + }) + } + } +} + +func globalBridgeWireBody(t *testing.T, wire string, stream bool) string { + t.Helper() + if !stream { + switch wire { + case "claude": + return `{"id":"msg_funding","type":"message","role":"assistant","model":"gpt-test","content":[{"type":"text","text":"partial output"}],"stop_reason":"end_turn","usage":{"input_tokens":4,"output_tokens":2}}` + case "gemini": + return `{"candidates":[{"content":{"role":"model","parts":[{"text":"partial output"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":4,"candidatesTokenCount":2,"totalTokenCount":6}}` + default: + return `{"id":"resp_funding","status":"completed","model":"gpt-test","output":[{"type":"message","id":"msg_output","role":"assistant","content":[{"type":"output_text","text":"partial output"}]}],"usage":{"input_tokens":4,"output_tokens":2,"total_tokens":6}}` + } + } + var frames []string + switch wire { + case "claude": + _, frames = streamFundingFrames(t, "claude_to_responses", "known") + frames = append(frames, `{"type":"message_stop"}`) + case "gemini": + _, frames = streamFundingFrames(t, "gemini_to_responses", "known") + default: + _, frames = streamFundingFrames(t, "responses_native", "known") + frames = append(frames, `{"type":"response.completed","response":{"id":"resp_funding","status":"completed","model":"gpt-test","usage":{"input_tokens":4,"output_tokens":2,"total_tokens":6}}}`) + } + return "data: " + strings.Join(frames, "\n\ndata: ") + "\n\n" +} diff --git a/relay/compatible_handler.go b/relay/compatible_handler.go index 7f3d6f0b..482935c9 100644 --- a/relay/compatible_handler.go +++ b/relay/compatible_handler.go @@ -74,7 +74,7 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (maxAPIError *types if info.RelayMode == relayconstant.RelayModeChatCompletions && !passThroughGlobal && !info.ChannelSetting.PassThroughBodyEnabled && - service.ShouldChatCompletionsUseResponsesGlobal(info.ChannelId, info.ChannelType, info.OriginModelName) { + shouldUseGlobalResponsesBridge(c, info) { applySystemPromptIfNeeded(c, info, request) usage, maxApiErr := chatCompletionsViaResponses(c, info, adaptor, request) if maxApiErr != nil { @@ -253,3 +253,25 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (maxAPIError *types } return nil } + +func shouldUseGlobalResponsesBridge(c *gin.Context, info *relaycommon.RelayInfo) bool { + if info == nil || !service.ShouldChatCompletionsUseResponsesGlobal(info.ChannelId, info.ChannelType, info.OriginModelName) { + return false + } + if info.ChannelType != constant.ChannelTypeAdvancedCustom { + return true + } + // Explicit converters already select their upstream protocol and matching + // reply handler. A global override must not replace that selection. + config := info.ChannelOtherSettings.AdvancedCustom + if config == nil { + return false + } + path := strings.Split(info.RequestURLPath, "?")[0] + if c != nil && c.Request != nil && c.Request.URL != nil { + path = c.Request.URL.Path + } + route, ok := config.MatchPath(path) + converter := strings.TrimSpace(route.Converter) + return ok && (converter == "" || converter == dto.AdvancedCustomConverterNone) +} diff --git a/relay/responses_native_missing_output_settlement_test.go b/relay/responses_native_missing_output_settlement_test.go index 7e09d99b..2a23648d 100644 --- a/relay/responses_native_missing_output_settlement_test.go +++ b/relay/responses_native_missing_output_settlement_test.go @@ -171,10 +171,10 @@ func runNativeMissingOutputLedger(t *testing.T, scenario nativeMissingOutputCase if usage == "terminal_zero" { // A terminal snapshot supplies usage in that one terminal; // emitting a second response.completed violates the protocol. - var copy map[string]any - require.NoError(t, common.UnmarshalJsonStr(encode(frame), ©)) - copy["response"].(map[string]any)["usage"] = map[string]any{"input_tokens": 0, "output_tokens": 0, "total_tokens": 0} - frame = copy + var terminalFrame map[string]any + require.NoError(t, common.UnmarshalJsonStr(encode(frame), &terminalFrame)) + terminalFrame["response"].(map[string]any)["usage"] = map[string]any{"input_tokens": 0, "output_tokens": 0, "total_tokens": 0} + frame = terminalFrame } } frames = append(frames, encode(frame)) diff --git a/service/openai_chat_responses_mode.go b/service/openai_chat_responses_mode.go index 79c4f2aa..6e727023 100644 --- a/service/openai_chat_responses_mode.go +++ b/service/openai_chat_responses_mode.go @@ -12,3 +12,7 @@ func ShouldChatCompletionsUseResponsesPolicy(policy model_setting.ChatCompletion func ShouldChatCompletionsUseResponsesGlobal(channelID int, channelType int, model string) bool { return openaicompat.ShouldChatCompletionsUseResponsesGlobal(channelID, channelType, model) } + +func SupportsOpenAIResponsesUpstream(channelType int) bool { + return openaicompat.SupportsOpenAIResponsesUpstream(channelType) +} diff --git a/service/openaicompat/policy.go b/service/openaicompat/policy.go index a9367de4..1c8406ed 100644 --- a/service/openaicompat/policy.go +++ b/service/openaicompat/policy.go @@ -1,9 +1,29 @@ package openaicompat -import "github.com/MAX-API-Next/MAX-API/setting/model_setting" +import ( + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/constant" + "github.com/MAX-API-Next/MAX-API/setting/model_setting" +) + +// The global bridge parses upstream OpenAI Responses, so a native Messages or +// GenerateContent conversion cannot use it. Keep this aligned with actual +// adaptor request/URL capabilities, not merely channel selection settings. +func SupportsOpenAIResponsesUpstream(channelType int) bool { + apiType, _ := common.ChannelType2APIType(channelType) + switch apiType { + case constant.APITypeOpenAI, constant.APITypeOpenRouter, constant.APITypeXinference, + constant.APITypeAli, constant.APITypeZhipuV4, constant.APITypePerplexity, + constant.APITypeCloudflare, constant.APITypeDeepSeek, constant.APITypeVolcEngine, + constant.APITypeXai, constant.APITypeCodex, constant.APITypeAdvancedCustom: + return true + default: + return false + } +} func ShouldChatCompletionsUseResponsesPolicy(policy model_setting.ChatCompletionsToResponsesPolicy, channelID int, channelType int, model string) bool { - if !policy.IsChannelEnabled(channelID, channelType) { + if !SupportsOpenAIResponsesUpstream(channelType) || !policy.IsChannelEnabled(channelID, channelType) { return false } return matchAnyRegex(policy.ModelPatterns, model) From e71ed04e76c3cf3e7eb81c994d268abb34753d0e Mon Sep 17 00:00:00 2001 From: CSCITech Date: Mon, 12 Oct 2026 13:48:43 +0800 Subject: [PATCH 12/24] test: assert stream finalization and upstream cancellation --- relay/channel/claude/relay_claude_test.go | 2 +- relay/file_search_stream_settlement_test.go | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/relay/channel/claude/relay_claude_test.go b/relay/channel/claude/relay_claude_test.go index 450cd3db..3a249216 100644 --- a/relay/channel/claude/relay_claude_test.go +++ b/relay/channel/claude/relay_claude_test.go @@ -241,7 +241,7 @@ func TestHandleStreamResponseDataConvertsResponsesFormat(t *testing.T) { for _, data := range streamData { require.Nil(t, HandleStreamResponseData(c, info, claudeInfo, data)) } - HandleStreamFinalResponse(c, info, claudeInfo) + require.Nil(t, HandleStreamFinalResponse(c, info, claudeInfo)) body := recorder.Body.String() assert.Contains(t, body, "response.created") diff --git a/relay/file_search_stream_settlement_test.go b/relay/file_search_stream_settlement_test.go index 21948dc6..d58b49a8 100644 --- a/relay/file_search_stream_settlement_test.go +++ b/relay/file_search_stream_settlement_test.go @@ -84,6 +84,7 @@ func TestFileSearchInterruptedStreamRetainsOriginalFunding(t *testing.T) { frames = append(frames, `{broken`) } var calls atomic.Int32 + upstreamCanceled := make(chan bool, 1) upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { calls.Add(1) w.Header().Set("Content-Type", "text/event-stream") @@ -95,7 +96,9 @@ func TestFileSearchInterruptedStreamRetainsOriginalFunding(t *testing.T) { if ending == "cancel" { select { case <-r.Context().Done(): + upstreamCanceled <- true case <-time.After(3 * time.Second): + upstreamCanceled <- false } } })) @@ -149,6 +152,12 @@ func TestFileSearchInterruptedStreamRetainsOriginalFunding(t *testing.T) { if cancelWriter != nil { require.True(t, cancelWriter.triggered.Load()) require.ErrorIs(t, c.Request.Context().Err(), context.Canceled) + select { + case canceled := <-upstreamCanceled: + require.True(t, canceled, "client cancellation must reach the upstream request") + case <-time.After(3 * time.Second): + t.Fatal("upstream handler did not finish within its cancellation guard") + } } assertStreamFundingBalances(t, db, funding, requestID, finalQuota, true, prompt, completion) var log model.Log From 49ffa9d52c86a1b48b56ef86edb4dc7b09418ddc Mon Sep 17 00:00:00 2001 From: CSCITech Date: Mon, 12 Oct 2026 15:07:42 +0800 Subject: [PATCH 13/24] fix: make Playground wallet reservations durable --- service/billing_playground_atomic_test.go | 263 +++++++++++++++++++++ service/billing_playground_process_test.go | 95 ++++++++ service/billing_process_recovery_test.go | 75 ++++-- service/billing_session.go | 27 ++- 4 files changed, 434 insertions(+), 26 deletions(-) create mode 100644 service/billing_playground_atomic_test.go create mode 100644 service/billing_playground_process_test.go diff --git a/service/billing_playground_atomic_test.go b/service/billing_playground_atomic_test.go new file mode 100644 index 00000000..79e44a80 --- /dev/null +++ b/service/billing_playground_atomic_test.go @@ -0,0 +1,263 @@ +package service + +import ( + "errors" + "fmt" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/MAX-API-Next/MAX-API/dto" + "github.com/MAX-API-Next/MAX-API/model" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func playgroundAtomicContext(requestID string, playground bool) (*gin.Context, *relaycommon.RelayInfo) { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + info := &relaycommon.RelayInfo{RequestId: requestID, UserId: 991, OriginModelName: "playground-atomic-model", + IsPlayground: playground, ForcePreConsume: true, UserSetting: dto.UserSetting{BillingPreference: "wallet_only"}} + if !playground { + info.TokenId, info.TokenKey = 992, "synthetic-playground-atomic-token" + } + return c, info +} + +func seedPlaygroundAtomic(t *testing.T) { + t.Helper() + truncate(t) + seedUser(t, 991, 100) + seedToken(t, 992, 991, "synthetic-playground-atomic-token", 100) +} + +func TestPlaygroundWalletLookupRaceCannotDebitNormalReservation(t *testing.T) { + for _, quota := range []int{0, 10} { + t.Run(fmt.Sprintf("quota_%d", quota), func(t *testing.T) { + seedPlaygroundAtomic(t) + lookedUp, resume := make(chan struct{}), make(chan struct{}) + var paused atomic.Bool + const callback = "test:playground-after-empty-source-read" + require.NoError(t, model.DB.Callback().Query().After("gorm:query").Register(callback, func(tx *gorm.DB) { + if tx.Statement.Table == "subscription_pre_consume_records" && paused.CompareAndSwap(false, true) { + close(lookedUp) + <-resume + } + })) + t.Cleanup(func() { _ = model.DB.Callback().Query().Remove(callback) }) + defer close(resume) + pgContext, pgInfo := playgroundAtomicContext("playground-normal-race", true) + pg := &BillingSession{relayInfo: pgInfo, funding: &WalletFunding{userId: 991}} + result := make(chan *types.MaxAPIError, 1) + go func() { result <- pg.preConsume(pgContext, quota) }() + select { + case <-lookedUp: + case <-time.After(5 * time.Second): + t.Fatal("Playground did not finish its empty source lookup") + } + paidContext, paidInfo := playgroundAtomicContext(pgInfo.RequestId, false) + paid := &BillingSession{relayInfo: paidInfo, funding: &WalletFunding{userId: 991}} + require.Nil(t, paid.preConsume(paidContext, quota)) + resume <- struct{}{} + select { + case apiErr := <-result: + t.Logf("wallet=%d token=%d original reservation=%d", getUserQuota(t, 991), getTokenRemainQuota(t, 992), quota) + require.EqualValues(t, 100-quota, getUserQuota(t, 991), "the Playground loser must not debit the paid reservation twice") + require.NotNil(t, apiErr) + require.ErrorIs(t, apiErr, model.ErrBillingSettlementOperationConflict) + require.True(t, types.IsSkipRetryError(apiErr)) + case <-time.After(5 * time.Second): + t.Fatal("Playground pre-consume did not return") + } + require.EqualValues(t, 100-quota, getUserQuota(t, 991)) + require.EqualValues(t, 100-quota, getTokenRemainQuota(t, 992)) + var record model.BillingSettlement + require.NoError(t, model.DB.Where("operation_key = ?", "request:"+pgInfo.RequestId+":pre-consume").First(&record).Error) + require.EqualValues(t, quota, record.AppliedFundingDelta) + require.EqualValues(t, quota, record.AppliedTokenDelta) + model.ProcessPendingBillingSettlementsOnce() + require.EqualValues(t, 100-quota, getUserQuota(t, 991)) + }) + } +} + +func TestPlaygroundWalletConcurrentEmptyLookupsChargeOnce(t *testing.T) { + seedPlaygroundAtomic(t) + lookedUp, resume := make(chan struct{}, 2), make(chan struct{}) + var lookups atomic.Int32 + const callback = "test:two-empty-playground-source-reads" + require.NoError(t, model.DB.Callback().Query().After("gorm:query").Register(callback, func(tx *gorm.DB) { + if tx.Statement.Table == "subscription_pre_consume_records" && lookups.Add(1) <= 2 { + lookedUp <- struct{}{} + <-resume + } + })) + t.Cleanup(func() { _ = model.DB.Callback().Query().Remove(callback) }) + defer close(resume) + results := make(chan *types.MaxAPIError, 2) + for range 2 { + c, info := playgroundAtomicContext("two-playground-race", true) + session := &BillingSession{relayInfo: info, funding: &WalletFunding{userId: 991}} + go func() { results <- session.preConsume(c, 10) }() + } + for range 2 { + select { + case <-lookedUp: + case <-time.After(5 * time.Second): + t.Fatal("both empty source lookups did not complete") + } + } + resume <- struct{}{} + resume <- struct{}{} + for range 2 { + select { + case apiErr := <-results: + require.Nil(t, apiErr) + case <-time.After(5 * time.Second): + t.Fatal("Playground debit did not complete") + } + } + require.EqualValues(t, 90, getUserQuota(t, 991)) + require.EqualValues(t, 100, getTokenRemainQuota(t, 992)) + var count int64 + require.NoError(t, model.DB.Model(&model.BillingSettlement{}).Where("operation_key = ?", "request:two-playground-race:pre-consume").Count(&count).Error) + require.EqualValues(t, 1, count) + model.ProcessPendingBillingSettlementsOnce() + require.EqualValues(t, 90, getUserQuota(t, 991)) +} + +func TestPlaygroundWalletDurableReservationPreservesLifecycle(t *testing.T) { + for _, actual := range []int{0, 6, 10, 14} { + t.Run(fmt.Sprintf("actual_%d", actual), func(t *testing.T) { + seedPlaygroundAtomic(t) + c, info := playgroundAtomicContext(fmt.Sprintf("playground-lifecycle-%d", actual), true) + session, apiErr := NewBillingSession(c, info, 10) + require.Nil(t, apiErr) + require.Equal(t, 10, session.GetPreConsumedQuota()) + require.Zero(t, session.tokenConsumed) + require.EqualValues(t, 90, getUserQuota(t, 991)) + require.EqualValues(t, 100, getTokenRemainQuota(t, 992)) + var initial model.BillingSettlement + require.NoError(t, model.DB.Where("operation_key = ?", "request:"+info.RequestId+":pre-consume").First(&initial).Error) + require.Equal(t, model.BillingSettlementStatusApplied, initial.Status) + require.EqualValues(t, 10, initial.AppliedFundingDelta) + require.Zero(t, initial.TokenDelta) + require.Zero(t, initial.AppliedTokenDelta) + paidContext, paidInfo := playgroundAtomicContext(info.RequestId, false) + _, conflict := NewBillingSession(paidContext, paidInfo, 10) + require.NotNil(t, conflict) + require.ErrorIs(t, conflict, model.ErrBillingSettlementOperationConflict) + effect := &model.BillingSettlementEffect{LogType: model.LogTypeConsume, ModelName: info.OriginModelName, + RequestID: info.RequestId, Quota: int64(actual), QuotaIsActual: true, UpdateUsage: true, PromptTokens: 1, CompletionTokens: 1} + require.NoError(t, session.SettleWithEffect(actual, effect)) + require.NoError(t, session.SettleWithEffect(actual, effect)) + for range 2 { + model.ProcessPendingBillingSettlementsOnce() + } + require.EqualValues(t, 100-actual, getUserQuota(t, 991)) + require.EqualValues(t, 100, getTokenRemainQuota(t, 992)) + var final model.BillingSettlement + require.NoError(t, model.DB.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(info.RequestId)).First(&final).Error) + require.EqualValues(t, actual-10, final.AppliedFundingDelta) + require.Zero(t, final.AppliedTokenDelta) + var logs int64 + require.NoError(t, model.DB.Model(&model.Log{}).Where("type = ?", model.LogTypeConsume).Count(&logs).Error) + require.EqualValues(t, 1, logs) + var user model.User + require.NoError(t, model.DB.First(&user, 991).Error) + require.EqualValues(t, actual, user.UsedQuota) + require.EqualValues(t, 1, user.RequestCount) + }) + } +} + +func TestPlaygroundWalletDurableReservationRefundAndRollback(t *testing.T) { + t.Run("refund", func(t *testing.T) { + seedPlaygroundAtomic(t) + c, info := playgroundAtomicContext("playground-refund", true) + session, apiErr := NewBillingSession(c, info, 10) + require.Nil(t, apiErr) + session.Refund(c) + session.Refund(c) + model.ProcessPendingBillingSettlementsOnce() + require.EqualValues(t, 100, getUserQuota(t, 991)) + require.EqualValues(t, 100, getTokenRemainQuota(t, 992)) + }) + t.Run("funding_transaction_rollback", func(t *testing.T) { + seedPlaygroundAtomic(t) + const callback = "test:playground-user-update-failure" + require.NoError(t, model.DB.Callback().Update().After("gorm:update").Register(callback, func(tx *gorm.DB) { + if tx.Statement.Table == "users" { + tx.AddError(errors.New("synthetic Playground funding transaction failure")) + } + })) + c, info := playgroundAtomicContext("playground-rollback", true) + _, apiErr := NewBillingSession(c, info, 10) + require.NotNil(t, apiErr) + require.NoError(t, model.DB.Callback().Update().Remove(callback)) + require.EqualValues(t, 100, getUserQuota(t, 991)) + require.EqualValues(t, 100, getTokenRemainQuota(t, 992)) + var count int64 + require.NoError(t, model.DB.Model(&model.BillingPreConsumeSelection{}).Count(&count).Error) + require.Zero(t, count) + var record model.BillingSettlement + require.NoError(t, model.DB.Where("operation_key = ?", "request:"+info.RequestId+":pre-consume").First(&record).Error) + require.NotEqual(t, model.BillingSettlementStatusApplied, record.Status) + require.Zero(t, record.AppliedFundingDelta) + // A failed paid admission stays inspectable and manual. Neither a fresh + // runner nor a repeated operation may charge a request that was rejected. + require.Equal(t, model.BillingSettlementStatusManual, record.Status) + applied, _, err := model.ApplyBillingSettlementOnce(model.BillingSettlementInput{ + OperationKey: record.OperationKey, Source: record.Source, UserID: record.UserID, TokenID: record.TokenID, + FundingDelta: record.FundingDelta, TokenDelta: record.TokenDelta, ManualOnFailure: record.ManualOnFailure, + PreConsumeRequestID: record.PreConsumeRequestID, PreConsumeModelName: record.PreConsumeModelName, + PreConsumeRequestedQuota: record.PreConsumeRequestedQuota, PreConsumeEffectiveQuota: record.PreConsumeEffectiveQuota}) + require.ErrorIs(t, err, model.ErrBillingSettlementManualReview) + require.Zero(t, applied) + model.ProcessPendingBillingSettlementsOnce() + require.EqualValues(t, 100, getUserQuota(t, 991)) + require.EqualValues(t, 100, getTokenRemainQuota(t, 992)) + }) +} + +func TestPlaygroundWalletReserveKeepsTokenZeroAndExactFunding(t *testing.T) { + for _, initial := range []int{0, 10} { + for _, ending := range []int{-1, 0, 6, 14, 18} { + t.Run(fmt.Sprintf("initial_%d/ending_%d", initial, ending), func(t *testing.T) { + seedPlaygroundAtomic(t) + c, info := playgroundAtomicContext("playground-extra-reserve", true) + session, apiErr := NewBillingSession(c, info, initial) + require.Nil(t, apiErr) + require.NoError(t, session.Reserve(14)) + require.NoError(t, session.Reserve(14)) + require.Equal(t, 14, session.GetPreConsumedQuota()) + require.Zero(t, session.tokenConsumed) + require.EqualValues(t, 86, getUserQuota(t, 991)) + require.EqualValues(t, 100, getTokenRemainQuota(t, 992)) + want := ending + if ending == -1 { + require.True(t, session.NeedsRefund()) + session.Refund(c) + session.Refund(c) + want = 0 + } else { + effect := &model.BillingSettlementEffect{LogType: model.LogTypeConsume, ModelName: info.OriginModelName, + RequestID: info.RequestId, Quota: int64(ending), QuotaIsActual: true, UpdateUsage: true, PromptTokens: 1, CompletionTokens: 1} + require.NoError(t, session.SettleWithEffect(ending, effect)) + require.NoError(t, session.SettleWithEffect(ending, effect)) + } + model.ProcessPendingBillingSettlementsOnce() + model.ProcessPendingBillingSettlementsOnce() + require.EqualValues(t, 100-want, getUserQuota(t, 991)) + require.EqualValues(t, 100, getTokenRemainQuota(t, 992)) + var final model.BillingSettlement + require.NoError(t, model.DB.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(info.RequestId)).First(&final).Error) + require.EqualValues(t, want-14, final.AppliedFundingDelta) + require.Zero(t, final.AppliedTokenDelta) + }) + } + } +} diff --git a/service/billing_playground_process_test.go b/service/billing_playground_process_test.go new file mode 100644 index 00000000..28e2fc33 --- /dev/null +++ b/service/billing_playground_process_test.go @@ -0,0 +1,95 @@ +package service + +import ( + "fmt" + "testing" + + "github.com/MAX-API-Next/MAX-API/model" + "github.com/stretchr/testify/require" +) + +func TestBillingPlaygroundProcessCrashRecovery(t *testing.T) { + funding := billingProcessFundingCases()[0] + phases := []string{"intent_committed", "funding_before_commit", "funding_committed", "receipt_before_log", + "log_before_commit", "log_committed", "usage_before_commit", "complete"} + for _, quota := range []int{130, 100, 6, 0} { + for _, phase := range phases { + t.Run(fmt.Sprintf("quota_%d/%s", quota, phase), func(t *testing.T) { + dir, db := seedBillingProcessDB(t, funding) + cfg := billingProcessConfig{Role: "work", Phase: phase, Preference: funding.preference, Source: funding.source, Quota: quota, Playground: true} + killBillingTestProcess(t, startBillingTestProcess(t, dir, cfg), phase) + assertBillingProcessLedger(t, db, funding, quota, phase, false, true) + original := billingProcessRecord(t, db) + cfg.Role, cfg.Phase = "recover", "" + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + assertBillingProcessLedger(t, db, funding, quota, "complete", false, true) + require.Equal(t, original.EffectPayload, billingProcessRecord(t, db).EffectPayload) + snapshot := billingProcessSnapshot(t, db) + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + require.Equal(t, snapshot, billingProcessSnapshot(t, db), "fresh PG recovery cannot repeat funding or effects") + }) + } + } +} + +func TestBillingPlaygroundProcessRefundRecovery(t *testing.T) { + funding := billingProcessFundingCases()[0] + for _, phase := range []string{"intent_committed", "funding_before_commit", "complete"} { + t.Run(phase, func(t *testing.T) { + dir, db := seedBillingProcessDB(t, funding) + cfg := billingProcessConfig{Role: "work", Phase: phase, Preference: funding.preference, Source: funding.source, Refund: true, Playground: true} + killBillingTestProcess(t, startBillingTestProcess(t, dir, cfg), phase) + assertBillingProcessLedger(t, db, funding, 0, phase, true, true) + cfg.Role, cfg.Phase = "recover", "" + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + assertBillingProcessLedger(t, db, funding, 0, "complete", true, true) + snapshot := billingProcessSnapshot(t, db) + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + require.Equal(t, snapshot, billingProcessSnapshot(t, db)) + }) + } +} + +func TestBillingPlaygroundProcessAdmissionRecovery(t *testing.T) { + funding := billingProcessFundingCases()[0] + for _, phase := range []string{"intent_committed", "funding_before_commit", "complete"} { + t.Run(phase, func(t *testing.T) { + dir, db := seedBillingProcessDB(t, funding) + cfg := billingProcessConfig{Role: "work", Phase: phase, Preference: funding.preference, Source: funding.source, Playground: true, AdmissionOnly: true} + killBillingTestProcess(t, startBillingTestProcess(t, dir, cfg), phase) + var initial model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", "request:process-recovery-request:pre-consume").First(&initial).Error) + require.EqualValues(t, 100, initial.FundingDelta) + require.Zero(t, initial.TokenDelta) + var user model.User + var token model.Token + require.NoError(t, db.First(&user, 951).Error) + require.NoError(t, db.First(&token, 952).Error) + want := int64(1000) + if phase == "complete" { + want = 900 + } + require.Equal(t, want, user.Quota) + require.EqualValues(t, 1000, token.RemainQuota) + cfg.Role, cfg.Phase = "recover", "" + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + require.NoError(t, db.First(&user, 951).Error) + require.EqualValues(t, 900, user.Quota) + // The original admitted reservation is reconstructable after a lost + // process; a refund uses the durable result, never a guessed amount. + cfg.Role = "refund_admission" + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + require.NoError(t, db.First(&user, 951).Error) + require.EqualValues(t, 1000, user.Quota) + require.NoError(t, db.First(&token, 952).Error) + require.EqualValues(t, 1000, token.RemainQuota) + var refund model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey("process-recovery-request")).First(&refund).Error) + require.EqualValues(t, -100, refund.AppliedFundingDelta) + require.Zero(t, refund.AppliedTokenDelta) + snapshot := billingProcessSnapshot(t, db) + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + require.Equal(t, snapshot, billingProcessSnapshot(t, db)) + }) + } +} diff --git a/service/billing_process_recovery_test.go b/service/billing_process_recovery_test.go index ddb55ef7..c4be11d0 100644 --- a/service/billing_process_recovery_test.go +++ b/service/billing_process_recovery_test.go @@ -47,15 +47,17 @@ func billingProcessFundingCases() []billingProcessFunding { } type billingProcessConfig struct { - Database string - Ready string - Release string - Role string - Phase string - Preference string - Source string - Quota int - Refund bool + Database string + Ready string + Release string + Role string + Phase string + Preference string + Source string + Quota int + Refund bool + Playground bool + AdmissionOnly bool } type billingProcessBarrier struct { @@ -89,10 +91,19 @@ func TestBillingProcessHelper(t *testing.T) { if cfg.Role == "work" { c, info := billingProcessContext(cfg) + if cfg.AdmissionOnly { + installBillingProcessBarrier(t, db, cfg) + } session, apiErr := NewBillingSession(c, info, 100) require.Nil(t, apiErr) require.Equal(t, cfg.Source, info.BillingSource) require.Equal(t, 100, session.GetPreConsumedQuota()) + if cfg.AdmissionOnly { + if cfg.Phase == "complete" { + pauseBillingProcess(t, cfg, billingProcessBarrier{Phase: cfg.Phase}) + } + return + } installBillingProcessBarrier(t, db, cfg) if cfg.Refund { input, err := session.PrepareRefundSettlement() @@ -104,7 +115,7 @@ func TestBillingProcessHelper(t *testing.T) { prompt, completion := billingProcessUsage(cfg.Quota) effect := newConsumeBillingSettlementEffect(info, model.RecordConsumeLogParams{ Content: "synthetic process recovery", ChannelId: 953, - ModelName: "process-recovery-model", TokenId: 952, TokenName: "synthetic-process-token", + ModelName: "process-recovery-model", TokenId: info.TokenId, TokenName: "synthetic-process-token", Group: "default", Quota: cfg.Quota, PromptTokens: prompt, CompletionTokens: completion, }, info.RequestId, "", true) require.NoError(t, session.SettleWithEffect(cfg.Quota, effect)) @@ -149,6 +160,9 @@ func TestBillingProcessHelper(t *testing.T) { funding = &SubscriptionFunding{requestId: info.RequestId, userId: 951, subscriptionId: 955, preConsumed: 100} } session := &BillingSession{relayInfo: info, funding: funding, preConsumedQuota: 100, tokenConsumed: 100} + if cfg.Playground { + session.tokenConsumed = 0 + } session.Refund(c) // The current durable Refund completes synchronously, including conflicts. session.mu.Lock() @@ -161,7 +175,21 @@ func TestBillingProcessHelper(t *testing.T) { // effect-only repair. It reconstructs inputs from the first persisted intent. model.ProcessPendingBillingSettlementsOnce() model.ProcessPendingBillingSettlementsOnce() - record := billingProcessRecord(t, db) + if cfg.Role == "refund_admission" { + c, info := billingProcessContext(cfg) + var initial model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", "request:process-recovery-request:pre-consume").First(&initial).Error) + require.Equal(t, model.BillingSettlementStatusApplied, initial.Status) + session := &BillingSession{relayInfo: info, funding: &WalletFunding{userId: 951, consumed: int(initial.AppliedFundingDelta)}, preConsumedQuota: int(initial.AppliedFundingDelta)} + session.Refund(c) + session.Refund(c) + } + var record model.BillingSettlement + if cfg.AdmissionOnly { + require.NoError(t, db.Where("operation_key = ?", "request:process-recovery-request:pre-consume").First(&record).Error) + } else { + record = billingProcessRecord(t, db) + } t.Logf("recovery pid=%d status=%s effect=%s attempts=%d next_attempt=%d", os.Getpid(), record.Status, record.EffectStatus, record.Attempts, record.NextAttempt) if record.Status == model.BillingSettlementStatusApplied { var effect *model.BillingSettlementEffect @@ -194,6 +222,9 @@ func billingProcessContext(cfg billingProcessConfig) (*gin.Context, *relaycommon info := &relaycommon.RelayInfo{RequestId: "process-recovery-request", UserId: 951, TokenId: 952, TokenKey: "synthetic-process-recovery-key", OriginModelName: "process-recovery-model", ForcePreConsume: true, UserSetting: dto.UserSetting{BillingPreference: cfg.Preference}} + if cfg.Playground { + info.IsPlayground, info.TokenId, info.TokenKey = true, 0, "" + } return c, info } @@ -258,7 +289,11 @@ func installBillingProcessBarrier(t *testing.T, db *gorm.DB, cfg billingProcessC create := func(tx *gorm.DB) { switch cfg.Phase { case "intent_committed": - if record, ok := tx.Statement.Dest.(*model.BillingSettlement); ok && record.OperationKey == model.BillingRequestFinalizeOperationKey("process-recovery-request") { + target := model.BillingRequestFinalizeOperationKey("process-recovery-request") + if cfg.AdmissionOnly { + target = "request:process-recovery-request:pre-consume" + } + if record, ok := tx.Statement.Dest.(*model.BillingSettlement); ok && record.OperationKey == target { stop(tx) } case "funding_committed", "receipt_before_log": @@ -426,7 +461,7 @@ func billingProcessRecord(t *testing.T, db *gorm.DB) model.BillingSettlement { return record } -func assertBillingProcessLedger(t *testing.T, db *gorm.DB, funding billingProcessFunding, quota int, phase string, refund bool) { +func assertBillingProcessLedger(t *testing.T, db *gorm.DB, funding billingProcessFunding, quota int, phase string, refund bool, playground ...bool) { t.Helper() beforeFunding := phase == "intent_committed" || phase == "funding_before_commit" fundedQuota := quota @@ -452,8 +487,12 @@ func assertBillingProcessLedger(t *testing.T, db *gorm.DB, funding billingProces } require.EqualValues(t, wallet, user.Quota) require.EqualValues(t, subUsed, sub.AmountUsed) - require.EqualValues(t, 1000-fundedQuota, token.RemainQuota) - require.EqualValues(t, fundedQuota, token.UsedQuota) + tokenFunded, tokenID, tokenDelta := fundedQuota, 952, quota-100 + if len(playground) > 0 && playground[0] { + tokenFunded, tokenID, tokenDelta = 0, 0, 0 + } + require.EqualValues(t, 1000-tokenFunded, token.RemainQuota) + require.EqualValues(t, tokenFunded, token.UsedQuota) counted := phase == "complete" && !refund wantUsed, wantRequests := 0, 0 if counted { @@ -483,9 +522,9 @@ func assertBillingProcessLedger(t *testing.T, db *gorm.DB, funding billingProces record := billingProcessRecord(t, db) require.Equal(t, funding.source, record.Source) require.EqualValues(t, 951, record.UserID) - require.EqualValues(t, 952, record.TokenID) + require.EqualValues(t, tokenID, record.TokenID) require.EqualValues(t, quota-100, record.FundingDelta) - require.EqualValues(t, quota-100, record.TokenDelta) + require.EqualValues(t, tokenDelta, record.TokenDelta) if beforeFunding { require.Equal(t, model.BillingSettlementStatusPending, record.Status) require.Zero(t, record.AppliedFundingDelta) @@ -494,7 +533,7 @@ func assertBillingProcessLedger(t *testing.T, db *gorm.DB, funding billingProces } else { require.Equal(t, model.BillingSettlementStatusApplied, record.Status) require.EqualValues(t, quota-100, record.AppliedFundingDelta) - require.EqualValues(t, quota-100, record.AppliedTokenDelta) + require.EqualValues(t, tokenDelta, record.AppliedTokenDelta) if !refund { if counted { require.Equal(t, model.BillingSettlementEffectApplied, record.EffectStatus) diff --git a/service/billing_session.go b/service/billing_session.go index b7d59cbf..a11dc06e 100644 --- a/service/billing_session.go +++ b/service/billing_session.go @@ -655,6 +655,9 @@ func (s *BillingSession) needsRefundLocked() bool { if s.tokenConsumed > 0 { return true } + if wallet, ok := s.funding.(*WalletFunding); ok && wallet.consumed > 0 { + return true + } // 订阅可能在 tokenConsumed=0 时仍预扣了额度 if sub, ok := s.funding.(*SubscriptionFunding); ok && sub.preConsumed > 0 { return true @@ -723,7 +726,9 @@ func (s *BillingSession) Reserve(targetQuota int) error { } s.preConsumedQuota += delta - s.tokenConsumed += delta + if !s.relayInfo.IsPlayground { + s.tokenConsumed += delta + } s.extraReserved += delta s.syncRelayInfo() return nil @@ -737,8 +742,8 @@ func (s *BillingSession) Reserve(targetQuota int) error { // 任一步骤失败时原子回滚已完成的步骤。 func (s *BillingSession) preConsume(c *gin.Context, quota int) *types.MaxAPIError { if _, wallet := s.funding.(*WalletFunding); wallet && s.relayInfo.IsPlayground { - // Playground wallets use legacy hooks without an operation key. Never - // let that path reinterpret an already persisted paid reservation. + // Reject persisted identities and lookup failures before wallet admission. + // The atomic source claim below also closes the empty-lookup race. _, persisted, err := model.ResolveBillingPreConsumeSource(s.relayInfo.RequestId) if err != nil { return mapAtomicPreConsumeError(err) @@ -760,7 +765,7 @@ func (s *BillingSession) preConsume(c *gin.Context, quota int) *types.MaxAPIErro } atomicPreConsumed := false - if !s.relayInfo.IsPlayground { + if _, wallet := s.funding.(*WalletFunding); !s.relayInfo.IsPlayground || wallet { switch funding := s.funding.(type) { case *WalletFunding: if s.relayInfo.RequestId == "" { @@ -769,6 +774,12 @@ func (s *BillingSession) preConsume(c *gin.Context, quota int) *types.MaxAPIErro } break } + tokenDelta := int64(effectiveQuota) + if s.relayInfo.IsPlayground { + // Playground has no paid API token, but its wallet reservation + // must claim the same durable identity before any debit. + tokenDelta = 0 + } applied, _, err := model.ApplyBillingSettlementOnce(model.BillingSettlementInput{ OperationKey: "request:" + s.relayInfo.RequestId + ":pre-consume", Source: model.BillingSettlementSourceWallet, @@ -776,7 +787,7 @@ func (s *BillingSession) preConsume(c *gin.Context, quota int) *types.MaxAPIErro TokenID: s.relayInfo.TokenId, TokenKey: s.relayInfo.TokenKey, FundingDelta: int64(effectiveQuota), - TokenDelta: int64(effectiveQuota), + TokenDelta: tokenDelta, ManualOnFailure: true, PreConsumeRequestID: s.relayInfo.RequestId, PreConsumeModelName: s.relayInfo.OriginModelName, @@ -790,7 +801,7 @@ func (s *BillingSession) preConsume(c *gin.Context, quota int) *types.MaxAPIErro return mapAtomicPreConsumeError(fmt.Errorf("wallet pre-consume applied unexpected amount: requested=%d applied=%d", effectiveQuota, applied)) } funding.consumed = int(applied) - s.tokenConsumed = int(applied) + s.tokenConsumed = int(tokenDelta) atomicPreConsumed = true case *SubscriptionFunding: if effectiveQuota <= 0 { @@ -814,8 +825,8 @@ func (s *BillingSession) preConsume(c *gin.Context, quota int) *types.MaxAPIErro } } - // Custom funding sources and playground requests retain their existing - // hooks; normal wallet/subscription requests use the atomic branches above. + // Custom sources and Playground subscriptions retain their existing hooks; + // wallets and normal subscriptions use the durable atomic branches above. if !atomicPreConsumed && effectiveQuota > 0 { if err := PreConsumeTokenQuota(s.relayInfo, effectiveQuota); err != nil { return types.NewErrorWithStatusCode(err, types.ErrorCodePreConsumeTokenQuotaFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog()) From a75a58b1aa52d29e8c37983da340e719bd2d15b2 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Mon, 12 Oct 2026 23:50:37 +0800 Subject: [PATCH 14/24] test: verify billing recovery with a separate SQLite log database --- service/billing_process_recovery_test.go | 109 ++++++--- service/billing_separate_log_process_test.go | 224 +++++++++++++++++++ 2 files changed, 304 insertions(+), 29 deletions(-) create mode 100644 service/billing_separate_log_process_test.go diff --git a/service/billing_process_recovery_test.go b/service/billing_process_recovery_test.go index c4be11d0..a2e09871 100644 --- a/service/billing_process_recovery_test.go +++ b/service/billing_process_recovery_test.go @@ -47,23 +47,26 @@ func billingProcessFundingCases() []billingProcessFunding { } type billingProcessConfig struct { - Database string - Ready string - Release string - Role string - Phase string - Preference string - Source string - Quota int - Refund bool - Playground bool - AdmissionOnly bool + Database string + SeparateLogDB bool + LogUnavailable bool + Ready string + Release string + Role string + Phase string + Preference string + Source string + Quota int + Refund bool + Playground bool + AdmissionOnly bool } type billingProcessBarrier struct { PID int Phase string Table string + Database string InsideTransaction bool } @@ -81,7 +84,17 @@ func TestBillingProcessHelper(t *testing.T) { require.NoError(t, common.Unmarshal(data, &cfg)) require.Equal(t, filepath.Join(filepath.Dir(configPath), "ledger.sqlite"), cfg.Database) db := openBillingProcessDB(t, cfg.Database) - model.DB, model.LOG_DB = db, db + logDB := db + if cfg.SeparateLogDB { + logDB = openBillingProcessDB(t, filepath.Join(filepath.Dir(cfg.Database), "logs.sqlite")) + } + model.DB, model.LOG_DB = db, logDB + if cfg.LogUnavailable { + require.True(t, cfg.SeparateLogDB, "only the independent log connection may be closed") + logSQL, err := logDB.DB() + require.NoError(t, err) + require.NoError(t, logSQL.Close()) + } common.UsingSQLite, common.RedisEnabled, common.BatchUpdateEnabled = true, false, false common.LogConsumeEnabled, common.DataExportEnabled = true, false gin.SetMode(gin.TestMode) @@ -92,7 +105,7 @@ func TestBillingProcessHelper(t *testing.T) { if cfg.Role == "work" { c, info := billingProcessContext(cfg) if cfg.AdmissionOnly { - installBillingProcessBarrier(t, db, cfg) + installBillingProcessBarrier(t, db, logDB, cfg) } session, apiErr := NewBillingSession(c, info, 100) require.Nil(t, apiErr) @@ -104,7 +117,7 @@ func TestBillingProcessHelper(t *testing.T) { } return } - installBillingProcessBarrier(t, db, cfg) + installBillingProcessBarrier(t, db, logDB, cfg) if cfg.Refund { input, err := session.PrepareRefundSettlement() require.NoError(t, err) @@ -150,7 +163,7 @@ func TestBillingProcessHelper(t *testing.T) { } } } - installBillingProcessBarrier(t, db, cfg) + installBillingProcessBarrier(t, db, logDB, cfg) if cfg.Role == "refund_competitor" { c, info := billingProcessContext(cfg) // The original session's refund intent competes with the already durable @@ -250,12 +263,29 @@ func openBillingProcessDB(t *testing.T, path string) *gorm.DB { } func seedBillingProcessDB(t *testing.T, funding billingProcessFunding) (string, *gorm.DB) { + t.Helper() + dir, db, _ := seedBillingProcessDatabases(t, funding, false) + return dir, db +} + +func seedBillingProcessDatabases(t *testing.T, funding billingProcessFunding, separateLog bool) (string, *gorm.DB, *gorm.DB) { t.Helper() dir := t.TempDir() db := openBillingProcessDB(t, filepath.Join(dir, "ledger.sqlite")) - require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}, &model.Channel{}, &model.Log{}, &model.BillingLogReceipt{}, + logDB := db + if separateLog { + logDB = openBillingProcessDB(t, filepath.Join(dir, "logs.sqlite")) + } + require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}, &model.Channel{}, &model.BillingSettlement{}, &model.BillingPreConsumeSelection{}, &model.CacheInvalidationTask{}, &model.UserSubscription{}, &model.SubscriptionPreConsumeRecord{}, &model.SubscriptionPlan{})) + require.NoError(t, logDB.AutoMigrate(&model.Log{}, &model.BillingLogReceipt{})) + if separateLog { + require.False(t, db.Migrator().HasTable(&model.Log{}), "the business DB must not silently receive consume logs") + require.False(t, db.Migrator().HasTable(&model.BillingLogReceipt{})) + require.False(t, logDB.Migrator().HasTable(&model.User{})) + require.False(t, logDB.Migrator().HasTable(&model.BillingSettlement{})) + } require.NoError(t, db.Create(&model.User{Id: 951, Username: "synthetic-process-recovery", Quota: funding.wallet}).Error) require.NoError(t, db.Create(&model.Token{Id: 952, UserId: 951, Key: "synthetic-process-recovery-key", Name: "synthetic-process-token", RemainQuota: 1000, Status: common.TokenStatusEnabled}).Error) @@ -266,24 +296,31 @@ func seedBillingProcessDB(t *testing.T, funding billingProcessFunding) (string, require.NoError(t, db.Create(&model.UserSubscription{Id: 955, UserId: 951, PlanId: 954, Status: "active", AmountTotal: 1000, AmountUsed: funding.subscriptionUsed, AllowWalletOverflow: funding.allowOverflow, StartTime: now - 3600, EndTime: now + 86400, LastResetTime: now - 3600}).Error) - return dir, db + return dir, db, logDB } // Callbacks exist only in the test child. The marker is written after the SQL // at that boundary succeeded, and the process remains blocked until killed. -func installBillingProcessBarrier(t *testing.T, db *gorm.DB, cfg billingProcessConfig) { +func installBillingProcessBarrier(t *testing.T, db, logDB *gorm.DB, cfg billingProcessConfig) { t.Helper() if cfg.Phase == "" || cfg.Phase == "complete" { return } var once sync.Once + barrierDB, database := db, "ledger" + if cfg.Phase == "funding_committed" || cfg.Phase == "receipt_before_log" || cfg.Phase == "log_before_commit" { + barrierDB = logDB + if cfg.SeparateLogDB { + database = "logs" + } + } stop := func(tx *gorm.DB) { if tx.Error != nil { return } once.Do(func() { _, inTx := tx.Statement.ConnPool.(*sql.Tx) - pauseBillingProcess(t, cfg, billingProcessBarrier{Phase: cfg.Phase, Table: tx.Statement.Table, InsideTransaction: inTx}) + pauseBillingProcess(t, cfg, billingProcessBarrier{Phase: cfg.Phase, Table: tx.Statement.Table, Database: database, InsideTransaction: inTx}) }) } create := func(tx *gorm.DB) { @@ -322,15 +359,15 @@ func installBillingProcessBarrier(t *testing.T, db *gorm.DB, cfg billingProcessC } const name = "test:billing-process-barrier" if cfg.Phase == "intent_committed" { - require.NoError(t, db.Callback().Create().After("gorm:commit_or_rollback_transaction").Register(name, create)) + require.NoError(t, barrierDB.Callback().Create().After("gorm:commit_or_rollback_transaction").Register(name, create)) } else if cfg.Phase == "funding_committed" { - require.NoError(t, db.Callback().Create().Before("gorm:create").Register(name, create)) + require.NoError(t, barrierDB.Callback().Create().Before("gorm:create").Register(name, create)) } else if cfg.Phase == "receipt_before_log" || cfg.Phase == "log_before_commit" { - require.NoError(t, db.Callback().Create().After("gorm:create").Register(name, create)) + require.NoError(t, barrierDB.Callback().Create().After("gorm:create").Register(name, create)) } else if cfg.Phase == "log_committed" { - require.NoError(t, db.Callback().Update().Before("gorm:update").Register(name, update)) + require.NoError(t, barrierDB.Callback().Update().Before("gorm:update").Register(name, update)) } else if cfg.Phase == "funding_before_commit" || cfg.Phase == "usage_before_commit" { - require.NoError(t, db.Callback().Update().After("gorm:update").Register(name, update)) + require.NoError(t, barrierDB.Callback().Update().After("gorm:update").Register(name, update)) } else { t.Fatalf("unknown interruption phase %q", cfg.Phase) } @@ -434,7 +471,7 @@ func billingProcessMarkerNotReady(err error) bool { return errors.Is(err, os.ErrNotExist) || runtime.GOOS == "windows" && errors.Is(err, syscall.Errno(32)) } -func killBillingTestProcess(t *testing.T, p *billingTestProcess, phase string) { +func killBillingTestProcess(t *testing.T, p *billingTestProcess, phase string) billingProcessBarrier { t.Helper() event := awaitBillingProcessBarrier(t, p) require.Equal(t, phase, event.Phase) @@ -446,7 +483,8 @@ func killBillingTestProcess(t *testing.T, p *billingTestProcess, phase string) { var exitErr *exec.ExitError require.ErrorAs(t, err, &exitErr, "an OS-forced termination must not be a normal test exit") require.False(t, exitErr.Success()) - t.Logf("forced termination pid=%d phase=%s table=%s inside_transaction=%t", event.PID, phase, event.Table, event.InsideTransaction) + t.Logf("forced termination pid=%d phase=%s database=%s table=%s inside_transaction=%t", event.PID, phase, event.Database, event.Table, event.InsideTransaction) + return event } func finishBillingTestProcess(t *testing.T, p *billingTestProcess) { @@ -462,6 +500,11 @@ func billingProcessRecord(t *testing.T, db *gorm.DB) model.BillingSettlement { } func assertBillingProcessLedger(t *testing.T, db *gorm.DB, funding billingProcessFunding, quota int, phase string, refund bool, playground ...bool) { + t.Helper() + assertBillingProcessDatabases(t, db, db, funding, quota, phase, refund, playground...) +} + +func assertBillingProcessDatabases(t *testing.T, db, logDB *gorm.DB, funding billingProcessFunding, quota int, phase string, refund bool, playground ...bool) { t.Helper() beforeFunding := phase == "intent_committed" || phase == "funding_before_commit" fundedQuota := quota @@ -503,8 +546,8 @@ func assertBillingProcessLedger(t *testing.T, db *gorm.DB, funding billingProces require.EqualValues(t, wantUsed, channel.UsedQuota) var logs []model.Log var receipts []model.BillingLogReceipt - require.NoError(t, db.Find(&logs).Error) - require.NoError(t, db.Find(&receipts).Error) + require.NoError(t, logDB.Find(&logs).Error) + require.NoError(t, logDB.Find(&receipts).Error) wantLogs := 0 if !refund && (phase == "log_committed" || phase == "usage_before_commit" || phase == "complete") { wantLogs = 1 @@ -581,6 +624,11 @@ func assertBillingProcessLedger(t *testing.T, db *gorm.DB, funding billingProces } func billingProcessSnapshot(t *testing.T, db *gorm.DB) [32]byte { + t.Helper() + return billingProcessDatabasesSnapshot(t, db, db) +} + +func billingProcessDatabasesSnapshot(t *testing.T, db, logDB *gorm.DB) [32]byte { t.Helper() var snapshot struct { Users []model.User @@ -595,9 +643,12 @@ func billingProcessSnapshot(t *testing.T, db *gorm.DB) [32]byte { Outbox []model.CacheInvalidationTask } for _, target := range []any{&snapshot.Users, &snapshot.Tokens, &snapshot.Channels, &snapshot.Subs, &snapshot.Reserved, - &snapshot.Selections, &snapshot.Settlements, &snapshot.Logs, &snapshot.Receipts, &snapshot.Outbox} { + &snapshot.Selections, &snapshot.Settlements, &snapshot.Outbox} { require.NoError(t, db.Order("id").Find(target).Error) } + for _, target := range []any{&snapshot.Logs, &snapshot.Receipts} { + require.NoError(t, logDB.Order("id").Find(target).Error) + } data, err := common.Marshal(snapshot) require.NoError(t, err) return sha256.Sum256(data) diff --git a/service/billing_separate_log_process_test.go b/service/billing_separate_log_process_test.go new file mode 100644 index 00000000..ddb6d384 --- /dev/null +++ b/service/billing_separate_log_process_test.go @@ -0,0 +1,224 @@ +package service + +import ( + "fmt" + "os" + "path/filepath" + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/model" + "github.com/stretchr/testify/require" +) + +type billingSeparateLogCase struct { + name string + funding billingProcessFunding + playground bool +} + +func billingSeparateLogCases() []billingSeparateLogCase { + var cases []billingSeparateLogCase + for _, funding := range billingProcessFundingCases() { + cases = append(cases, billingSeparateLogCase{name: funding.name, funding: funding}) + } + return append(cases, billingSeparateLogCase{name: "playground_wallet", funding: billingProcessFundingCases()[0], playground: true}) +} + +func (c billingSeparateLogCase) config(quota int, phase string) billingProcessConfig { + return billingProcessConfig{Role: "work", Phase: phase, Preference: c.funding.preference, Source: c.funding.source, + Quota: quota, Playground: c.playground, SeparateLogDB: true} +} + +func TestBillingSeparateLogProcessCrashRecovery(t *testing.T) { + // The first three barriers belong to the independent log transaction; the + // last two belong to the business transaction that applies usage counters. + phases := []string{"funding_committed", "receipt_before_log", "log_before_commit", "log_committed", "usage_before_commit"} + for _, c := range billingSeparateLogCases() { + for _, quota := range []int{130, 100, 6, 0} { + for _, phase := range phases { + t.Run(fmt.Sprintf("%s/quota_%d/%s", c.name, quota, phase), func(t *testing.T) { + dir, db, logDB := seedBillingProcessDatabases(t, c.funding, true) + cfg := c.config(quota, phase) + event := killBillingTestProcess(t, startBillingTestProcess(t, dir, cfg), phase) + wantDatabase := "logs" + if phase == "log_committed" || phase == "usage_before_commit" { + wantDatabase = "ledger" + } + require.Equal(t, wantDatabase, event.Database, "interrupt the transaction in the intended database") + assertBillingProcessDatabases(t, db, logDB, c.funding, quota, phase, false, c.playground) + original := billingProcessRecord(t, db) + cfg.Role, cfg.Phase = "recover", "" + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + assertBillingProcessDatabases(t, db, logDB, c.funding, quota, "complete", false, c.playground) + require.Equal(t, original.EffectPayload, billingProcessRecord(t, db).EffectPayload) + snapshot := billingProcessDatabasesSnapshot(t, db, logDB) + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + require.Equal(t, snapshot, billingProcessDatabasesSnapshot(t, db, logDB), "neither database may change on duplicate recovery") + }) + } + } + } +} + +func TestBillingSeparateLogProcessRepeatedCrashRecovery(t *testing.T) { + for _, c := range billingSeparateLogCases()[:2] { + for _, quota := range []int{130, 0} { + t.Run(fmt.Sprintf("%s/quota_%d", c.name, quota), func(t *testing.T) { + dir, db, logDB := seedBillingProcessDatabases(t, c.funding, true) + cfg := c.config(quota, "funding_committed") + killBillingTestProcess(t, startBillingTestProcess(t, dir, cfg), cfg.Phase) + original := billingProcessRecord(t, db) + cfg.Role = "recover" + for _, phase := range []string{"log_committed", "usage_before_commit"} { + cfg.Phase = phase + killBillingTestProcess(t, startBillingTestProcess(t, dir, cfg), phase) + assertBillingProcessDatabases(t, db, logDB, c.funding, quota, phase, false) + } + cfg.Phase = "" + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + assertBillingProcessDatabases(t, db, logDB, c.funding, quota, "complete", false) + require.Equal(t, original.EffectPayload, billingProcessRecord(t, db).EffectPayload) + snapshot := billingProcessDatabasesSnapshot(t, db, logDB) + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + require.Equal(t, snapshot, billingProcessDatabasesSnapshot(t, db, logDB)) + }) + } + } +} + +func TestBillingSeparateLogProcessUnavailableRecovery(t *testing.T) { + for _, c := range billingSeparateLogCases() { + for _, quota := range []int{130, 0} { + t.Run(fmt.Sprintf("%s/quota_%d", c.name, quota), func(t *testing.T) { + dir, db, logDB := seedBillingProcessDatabases(t, c.funding, true) + cfg := c.config(quota, "funding_committed") + killBillingTestProcess(t, startBillingTestProcess(t, dir, cfg), cfg.Phase) + original := billingProcessRecord(t, db) + snapshot := billingProcessDatabasesSnapshot(t, db, logDB) + cfg.Role, cfg.Phase, cfg.LogUnavailable = "recover", "", true + // Close only the child's independent sql.DB. The original funding + // is already applied, and failed effects must stay durably pending. + for attempt := 0; attempt < 2; attempt++ { + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + assertBillingProcessDatabases(t, db, logDB, c.funding, quota, "funding_committed", false, c.playground) + require.Equal(t, snapshot, billingProcessDatabasesSnapshot(t, db, logDB), "a log outage cannot change either ledger or drop the intent") + } + cfg.LogUnavailable = false + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + assertBillingProcessDatabases(t, db, logDB, c.funding, quota, "complete", false, c.playground) + require.Equal(t, original.EffectPayload, billingProcessRecord(t, db).EffectPayload) + snapshot = billingProcessDatabasesSnapshot(t, db, logDB) + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + require.Equal(t, snapshot, billingProcessDatabasesSnapshot(t, db, logDB)) + }) + } + } +} + +func TestBillingSeparateLogProcessConcurrentRecoveryAndRefund(t *testing.T) { + cases := billingSeparateLogCases() + for _, c := range []billingSeparateLogCase{cases[0], cases[1], cases[4]} { + for _, quota := range []int{130, 0} { + t.Run(fmt.Sprintf("%s/quota_%d", c.name, quota), func(t *testing.T) { + dir, db, logDB := seedBillingProcessDatabases(t, c.funding, true) + cfg := c.config(quota, "funding_committed") + killBillingTestProcess(t, startBillingTestProcess(t, dir, cfg), cfg.Phase) + original := billingProcessRecord(t, db) + cfg.Phase, cfg.Release = "", filepath.Join(dir, "release") + var workers []*billingTestProcess + for _, role := range []string{"recover", "recover", "refund_competitor"} { + cfg.Role = role + p := startBillingTestProcess(t, dir, cfg) + workers = append(workers, p) + awaitBillingProcessBarrier(t, p) + } + require.NoError(t, os.WriteFile(cfg.Release, []byte("release"), 0600)) + for _, p := range workers { + finishBillingTestProcess(t, p) + } + cfg.Role, cfg.Release = "recover", "" + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + assertBillingProcessDatabases(t, db, logDB, c.funding, quota, "complete", false, c.playground) + require.Equal(t, original.EffectPayload, billingProcessRecord(t, db).EffectPayload, "the competing refund must not replace finalized usage") + snapshot := billingProcessDatabasesSnapshot(t, db, logDB) + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + require.Equal(t, snapshot, billingProcessDatabasesSnapshot(t, db, logDB)) + }) + } + } +} + +func TestBillingSeparateLogProcessRefundRecovery(t *testing.T) { + for _, c := range billingSeparateLogCases() { + for _, phase := range []string{"intent_committed", "funding_before_commit", "complete"} { + t.Run(c.name+"/"+phase, func(t *testing.T) { + dir, db, logDB := seedBillingProcessDatabases(t, c.funding, true) + cfg := c.config(0, phase) + cfg.Refund = true + killBillingTestProcess(t, startBillingTestProcess(t, dir, cfg), phase) + assertBillingProcessDatabases(t, db, logDB, c.funding, 0, phase, true, c.playground) + cfg.Role, cfg.Phase = "recover", "" + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + assertBillingProcessDatabases(t, db, logDB, c.funding, 0, "complete", true, c.playground) + snapshot := billingProcessDatabasesSnapshot(t, db, logDB) + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + require.Equal(t, snapshot, billingProcessDatabasesSnapshot(t, db, logDB)) + }) + } + } +} + +func TestBillingSeparateLogProcessSubscriptionPeriodRecovery(t *testing.T) { + c := billingSeparateLogCases()[1] + for _, quota := range []int{130, 0} { + t.Run(fmt.Sprintf("quota_%d", quota), func(t *testing.T) { + dir, db, logDB := seedBillingProcessDatabases(t, c.funding, true) + cfg := c.config(quota, "funding_committed") + killBillingTestProcess(t, startBillingTestProcess(t, dir, cfg), cfg.Phase) + original := billingProcessRecord(t, db) + cfg.Role, cfg.Phase = "reset", "" + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + cfg.Role = "recover" + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + var sub model.UserSubscription + var user model.User + var token model.Token + var channel model.Channel + for _, target := range []struct { + row any + id int + }{{&sub, 955}, {&user, 951}, {&token, 952}, {&channel, 953}} { + require.NoError(t, db.First(target.row, target.id).Error) + } + require.Zero(t, sub.AmountUsed, "effect recovery must not charge the new subscription period") + require.EqualValues(t, 1000, user.Quota) + require.EqualValues(t, quota, token.UsedQuota) + require.EqualValues(t, 1000-quota, token.RemainQuota) + require.EqualValues(t, quota, user.UsedQuota) + require.EqualValues(t, 1, user.RequestCount) + require.EqualValues(t, quota, channel.UsedQuota) + record := billingProcessRecord(t, db) + require.Equal(t, original.EffectPayload, record.EffectPayload) + require.Equal(t, model.BillingSettlementStatusApplied, record.Status) + require.Equal(t, model.BillingSettlementEffectApplied, record.EffectStatus) + require.EqualValues(t, quota-100, record.AppliedFundingDelta) + require.EqualValues(t, quota-100, record.AppliedTokenDelta) + var logs []model.Log + var receipts []model.BillingLogReceipt + require.NoError(t, logDB.Find(&logs).Error) + require.NoError(t, logDB.Find(&receipts).Error) + require.Len(t, logs, 1) + require.Len(t, receipts, 1) + require.EqualValues(t, quota, logs[0].Quota) + require.Equal(t, record.OperationKey, receipts[0].OperationKey) + var other map[string]any + require.NoError(t, common.UnmarshalJsonStr(logs[0].Other, &other)) + require.EqualValues(t, quota, other["subscription_used"]) + require.EqualValues(t, 1000-quota, other["subscription_remain"]) + snapshot := billingProcessDatabasesSnapshot(t, db, logDB) + finishBillingTestProcess(t, startBillingTestProcess(t, dir, cfg)) + require.Equal(t, snapshot, billingProcessDatabasesSnapshot(t, db, logDB)) + }) + } +} From 690f9711e99bb3b352e8c496200df2f8a04bc937 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Tue, 13 Oct 2026 01:57:34 +0800 Subject: [PATCH 15/24] fix: finalize normal Responses stream endings consistently --- relay/channel/claude/relay-claude.go | 6 +-- relay/channel/gemini/relay_responses.go | 2 +- relay/channel/openai/empty_completion_test.go | 40 +++++++++++++++++++ relay/channel/openai/relay_responses.go | 6 +-- relay/channel/openai/responses_via_chat.go | 2 +- ...sponses_custom_tool_eof_settlement_test.go | 13 ++++-- .../chat_to_responses_response.go | 2 +- .../openaicompat/custom_tool_response_test.go | 36 +++++++++++++++++ 8 files changed, 95 insertions(+), 12 deletions(-) diff --git a/relay/channel/claude/relay-claude.go b/relay/channel/claude/relay-claude.go index 55497651..0bb20f73 100644 --- a/relay/channel/claude/relay-claude.go +++ b/relay/channel/claude/relay-claude.go @@ -1088,10 +1088,10 @@ func ClaudeStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon. } if claudeInfo.ResponsesState != nil { claudeInfo.ResponsesState.MarkIncomplete(reason) - // EOF means the client can still receive a terminal Responses - // event. For cancellation or transport failure, do not write + // A normal ending still permits a terminal Responses event. + // For cancellation or transport failure, do not write // after the client has gone away. - if streamStatus == nil || streamStatus.EndReason == relaycommon.StreamEndReasonEOF { + if !streamStatus.IsAbnormalEnd() { if finalErr := HandleStreamFinalResponse(c, info, claudeInfo); finalErr != nil { if !claudeInfo.responsesOutputDelivered { return nil, finalErr diff --git a/relay/channel/gemini/relay_responses.go b/relay/channel/gemini/relay_responses.go index 36c9b741..f2c26bdf 100644 --- a/relay/channel/gemini/relay_responses.go +++ b/relay/channel/gemini/relay_responses.go @@ -285,7 +285,7 @@ func GeminiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, r reason = string(info.StreamStatus.EndReason) } state.MarkIncomplete(reason) - if info.StreamStatus == nil || info.StreamStatus.EndReason == relaycommon.StreamEndReasonEOF { + if !info.StreamStatus.IsAbnormalEnd() { if !finalizeStream() { finalErr := types.NewError(streamErr, streamErr.GetErrorCode(), types.ErrOptionWithSkipRetry()) if !outputDelivered { diff --git a/relay/channel/openai/empty_completion_test.go b/relay/channel/openai/empty_completion_test.go index 3db7fbfc..cb5694f3 100644 --- a/relay/channel/openai/empty_completion_test.go +++ b/relay/channel/openai/empty_completion_test.go @@ -8,13 +8,53 @@ import ( "testing" "github.com/MAX-API-Next/MAX-API/common" + "github.com/MAX-API-Next/MAX-API/constant" "github.com/MAX-API-Next/MAX-API/dto" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "github.com/MAX-API-Next/MAX-API/service" "github.com/MAX-API-Next/MAX-API/types" "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" ) +func TestOaiResponsesRetrySuccessRequiresProtocolTerminal(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitTokenEncoders() + for _, ending := range []string{"eof", "done", "completed", "incomplete"} { + t.Run(ending, func(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + retry := map[string]any{"empty_retry": true, "empty_retry_result": "retrying"} + c.Set(string(constant.ContextKeyEmptyCompletionInfo), retry) + info := &relaycommon.RelayInfo{RelayFormat: types.RelayFormatOpenAIResponses, OriginModelName: "gpt-test", DisablePing: true, + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-test"}} + body := "data: " + `{"type":"response.output_text.delta","delta":"visible"}` + "\n\n" + terminal := ending == "completed" || ending == "incomplete" + if terminal { + body += "data: " + `{"type":"response.` + ending + `","response":{"id":"resp_retry","status":"` + ending + `","output":[],"usage":{"input_tokens":4,"output_tokens":2,"total_tokens":6}}}` + "\n\n" + } else if ending == "done" { + body += "data: [DONE]\n\n" + } + usage, apiErr := OaiResponsesStreamHandler(c, info, &http.Response{StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(body))}) + require.NotNil(t, usage) + require.Contains(t, recorder.Body.String(), "visible") + if terminal { + require.Nil(t, apiErr) + require.Equal(t, "success", retry["empty_retry_result"]) + require.Equal(t, 6, usage.TotalTokens) + } else { + require.NotNil(t, apiErr) + require.True(t, types.IsSkipRetryError(apiErr)) + require.Equal(t, "retrying", retry["empty_retry_result"]) + require.NotContains(t, retry, "final_usage") + } + }) + } +} + func TestIsEmptyResponsesCompletion(t *testing.T) { tests := []struct { name string diff --git a/relay/channel/openai/relay_responses.go b/relay/channel/openai/relay_responses.go index 13945ff1..4c00bd4c 100644 --- a/relay/channel/openai/relay_responses.go +++ b/relay/channel/openai/relay_responses.go @@ -571,9 +571,6 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp if service.ResponseAuditEnabled() { service.SetRelayResponseAuditContent(info, responseTextBuilder.String()) } - if hasVisiblePayload { - recordEmptyCompletionRetrySuccess(c, info, usage) - } if !terminalEventSeen { reason := "upstream_eof" if info.StreamStatus != nil && info.StreamStatus.IsAbnormalEnd() { @@ -587,6 +584,9 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp } return usage, openaicompat.NewResponsesStreamIncompleteError(reason) } + if hasVisiblePayload { + recordEmptyCompletionRetrySuccess(c, info, usage) + } return usage, nil } diff --git a/relay/channel/openai/responses_via_chat.go b/relay/channel/openai/responses_via_chat.go index 418c23bd..a0a64517 100644 --- a/relay/channel/openai/responses_via_chat.go +++ b/relay/channel/openai/responses_via_chat.go @@ -224,7 +224,7 @@ func OaiChatToResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo reason = string(info.StreamStatus.EndReason) } state.MarkIncomplete(reason) - if info.StreamStatus == nil || info.StreamStatus.EndReason == relaycommon.StreamEndReasonEOF { + if !info.StreamStatus.IsAbnormalEnd() { if !finalizeStream() { finalErr := types.NewError(streamErr, streamErr.GetErrorCode(), types.ErrOptionWithSkipRetry()) if !outputDelivered { diff --git a/relay/responses_custom_tool_eof_settlement_test.go b/relay/responses_custom_tool_eof_settlement_test.go index d1377385..0105a7bb 100644 --- a/relay/responses_custom_tool_eof_settlement_test.go +++ b/relay/responses_custom_tool_eof_settlement_test.go @@ -73,7 +73,7 @@ func TestResponsesCustomToolEOFFinalizationUsesDurableLedger(t *testing.T) { if provider == "chat" && strings.HasSuffix(usageKind, "input_only") { continue } - for _, ending := range []string{"eof", "first_input_write_failure", "second_input_write_failure", "finish_second_input_write_failure", "input_flush_failure", "terminal_write_failure", "no_delivery"} { + for _, ending := range []string{"eof", "done", "first_input_write_failure", "second_input_write_failure", "finish_second_input_write_failure", "input_flush_failure", "terminal_write_failure", "no_delivery"} { t.Run(strings.Join([]string{provider, funding, usageKind, ending}, "/"), func(t *testing.T) { db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) require.NoError(t, err) @@ -118,6 +118,9 @@ func TestResponsesCustomToolEOFFinalizationUsesDurableLedger(t *testing.T) { prompt = 0 } channelType, frames := customToolEOFFrames(t, provider, usageKind, inputs) + if ending == "done" { + frames = append(frames, "[DONE]") + } if ending == "finish_second_input_write_failure" && usageKind != "known" && usageKind != "zero" { switch provider { case "chat": @@ -159,7 +162,7 @@ func TestResponsesCustomToolEOFFinalizationUsesDurableLedger(t *testing.T) { info.SetEstimatePromptTokens(10) require.Nil(t, service.PreConsumeBilling(c, 100, info)) apiErr := ResponsesHelper(c, info) - if ending == "eof" && (usageKind == "known" || usageKind == "zero") { + if (ending == "eof" || ending == "done") && (usageKind == "known" || usageKind == "zero") { require.Nil(t, apiErr) } else { require.NotNil(t, apiErr) @@ -240,8 +243,12 @@ func TestResponsesCustomToolEOFFinalizationUsesDurableLedger(t *testing.T) { } } require.Equal(t, deliveredInput, gotInput.String(), "only fully written input deltas may participate in settlement") - if ending == "eof" { + if ending == "eof" || ending == "done" { require.Equal(t, 1, terminalCount) + if usageKind != "known" && usageKind != "zero" { + require.Contains(t, writer.Body.String(), `"type":"response.incomplete"`) + require.NotContains(t, writer.Body.String(), `"type":"response.completed"`) + } } else { require.Zero(t, terminalCount, "no terminal may be appended after the downstream fails") } diff --git a/service/openaicompat/chat_to_responses_response.go b/service/openaicompat/chat_to_responses_response.go index bb1cbe5c..f1b3f429 100644 --- a/service/openaicompat/chat_to_responses_response.go +++ b/service/openaicompat/chat_to_responses_response.go @@ -147,7 +147,7 @@ func ChatCompletionsResponseToResponsesResponseWithCustomTools(resp *dto.OpenAIT } for i, toolCall := range choice.Message.ParseToolCalls() { item := chatToolCallToResponsesOutput(toolCall, id, i, status) - if _, custom := customToolNames[toolCall.Function.Name]; custom && item.Type == responsesOutputTypeFunctionCall { + if _, custom := customToolNames[strings.TrimSpace(toolCall.Function.Name)]; custom && item.Type == responsesOutputTypeFunctionCall { item.Type = responsesOutputTypeCustomToolCall item.Input = customToolInputFromArguments(toolCall.Function.Arguments) item.Arguments = nil diff --git a/service/openaicompat/custom_tool_response_test.go b/service/openaicompat/custom_tool_response_test.go index 64a089dd..279ae4ce 100644 --- a/service/openaicompat/custom_tool_response_test.go +++ b/service/openaicompat/custom_tool_response_test.go @@ -60,3 +60,39 @@ func TestResponsesCustomToolPreservesRawInputAndOrdinaryFunctions(t *testing.T) }) } } + +func TestResponsesCustomToolWhitespaceNameMatchesStream(t *testing.T) { + for _, name := range []string{" apply_patch", "apply_patch\t\n", " \tapply_patch \n"} { + t.Run(name, func(t *testing.T) { + tool := dto.ToolCallResponse{ID: "call_custom", Type: "function", Function: dto.FunctionResponse{Name: name, Arguments: `{"input":"raw patch"}`}} + message := dto.Message{Role: "assistant"} + message.SetToolCalls([]dto.ToolCallResponse{tool}) + names := map[string]struct{}{"apply_patch": {}} + response, _, err := ChatCompletionsResponseToResponsesResponseWithCustomTools(&dto.OpenAITextResponse{ + Choices: []dto.OpenAITextResponseChoice{{Message: message, FinishReason: "tool_calls"}}, + }, "resp_custom", names) + require.NoError(t, err) + require.Len(t, response.Output, 1) + require.Equal(t, "custom_tool_call", response.Output[0].Type) + require.Equal(t, name, response.Output[0].Name) + require.Equal(t, "raw patch", response.Output[0].Input) + require.Nil(t, response.Output[0].Arguments) + require.Equal(t, "call_custom", response.Output[0].CallId) + + state := NewChatToResponsesStreamState("resp_custom", "test") + state.CustomToolNames = names + tool.SetIndex(0) + _, err = ChatCompletionsStreamChunkToResponsesEvents(&dto.ChatCompletionsStreamResponse{ + Choices: []dto.ChatCompletionsStreamResponseChoice{{Delta: dto.ChatCompletionsStreamResponseChoiceDelta{ToolCalls: []dto.ToolCallResponse{tool}}}}, + }, state) + require.NoError(t, err) + events := FinalizeChatCompletionsStreamToResponses(state) + require.NotEmpty(t, events) + stream := events[len(events)-1].Payload.Response + require.NotNil(t, stream) + require.Equal(t, response.Output[0].Type, stream.Output[0].Type) + require.Equal(t, "apply_patch", stream.Output[0].Name) + require.Equal(t, response.Output[0].Input, stream.Output[0].Input) + }) + } +} From d802f167fb5b1c07820013778cd4657bf7a7f949 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Tue, 13 Oct 2026 03:24:05 +0800 Subject: [PATCH 16/24] fix: preserve delivered Responses terminal outcomes and tool identity --- model/subscription_reset_billing_test.go | 7 ++ relay/channel/openai/relay_responses.go | 13 +++ .../openai/responses_terminal_repair_test.go | 53 ++++++++++++ relay/chat_stream_billing_test.go | 27 ++++-- ...s_native_missing_output_settlement_test.go | 82 ++++++++++++++++++- .../chat_to_responses_response.go | 4 +- .../openaicompat/custom_tool_response_test.go | 14 +++- 7 files changed, 185 insertions(+), 15 deletions(-) diff --git a/model/subscription_reset_billing_test.go b/model/subscription_reset_billing_test.go index 71f610f4..49e4de2a 100644 --- a/model/subscription_reset_billing_test.go +++ b/model/subscription_reset_billing_test.go @@ -279,6 +279,13 @@ func TestSubscriptionResetChecksAllReservationsBeyondOneBatch(t *testing.T) { var got UserSubscription require.NoError(t, DB.First(&got, sub.Id).Error) require.EqualValues(t, 100, got.AmountUsed) + + _, _, err = ApplyBillingSettlementOnce(resetBillingFinalize(sub, requestID)) + require.NoError(t, err) + _, err = AdminResetUserSubscriptionsByPlan(sub.UserId, plan.Id, false) + require.NoError(t, err, "all applied reservations across batches must pass validation") + require.NoError(t, DB.First(&got, sub.Id).Error) + require.Zero(t, got.AmountUsed) } func TestSubscriptionResetAllowsEffectOnlyRecovery(t *testing.T) { diff --git a/relay/channel/openai/relay_responses.go b/relay/channel/openai/relay_responses.go index 4c00bd4c..0b0086a7 100644 --- a/relay/channel/openai/relay_responses.go +++ b/relay/channel/openai/relay_responses.go @@ -375,6 +375,7 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp hasVisiblePayload := false emptyCompletionRecorded := false terminalEventSeen := false + terminalDelivered := false usageReported := false outputDelivered := false @@ -394,6 +395,13 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp streamErr = types.NewError(writeErr, types.ErrorCodeBadResponse) return false } + if delivered { + switch event.response.Type { + case "response.completed", "response.incomplete", "response.done", "response.failed", + "response.cancelled", "response.canceled", "error", "response.error": + terminalDelivered = true + } + } if delivered && responsesStreamHasVisibleOutput(&event.response) { info.SetFirstResultTime() } @@ -428,6 +436,11 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp sr.Stop(streamErr) return } + // A delivered terminal fixes this generation's outcome and usage. Late + // upstream frames must not alter billing or write to the closed stream. + if terminalDelivered { + return + } // 检查当前数据是否包含 completed 状态和 usage 信息 var streamResponse dto.ResponsesStreamResponse diff --git a/relay/channel/openai/responses_terminal_repair_test.go b/relay/channel/openai/responses_terminal_repair_test.go index bb789268..9b4acadc 100644 --- a/relay/channel/openai/responses_terminal_repair_test.go +++ b/relay/channel/openai/responses_terminal_repair_test.go @@ -243,3 +243,56 @@ func TestNativeResponsesLegalTerminalsPreserveReportedUsage(t *testing.T) { } } } + +func TestNativeResponsesDeliveredTerminalIgnoresLaterFrames(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitTokenEncoders() + oldEmptyRetry := common.EmptyCompletionRetryEnabled + common.EmptyCompletionRetryEnabled = false + t.Cleanup(func() { common.EmptyCompletionRetryEnabled = oldEmptyRetry }) + tails := []struct{ name, frame string }{ + {"duplicate", `{"type":"response.completed","response":{"status":"completed","usage":{"input_tokens":1000,"output_tokens":999,"total_tokens":1999}}}`}, + {"metadata", `{"type":"response.in_progress","response":{"usage":{"input_tokens":1000,"output_tokens":999,"total_tokens":1999}}}`}, + {"malformed", `{broken`}, + {"failure", `{"type":"response.failed","response":{"status":"failed","error":{"message":"late failure"}}}`}, + {"tool", `{"type":"response.output_item.done","item":{"type":"function_call","name":"late_tool","id":"late_tool","arguments":"{}"}}`}, + {"output", `{"type":"response.output_text.delta","delta":"late output"}`}, + } + for _, event := range []string{"response.completed", "response.incomplete", "response.done"} { + for _, usageKind := range []string{"missing", "known", "zero"} { + for _, tail := range tails { + t.Run(event+"/"+usageKind+"/"+tail.name, func(t *testing.T) { + c, info := newResponsesUsageTestContext() + recorder := httptest.NewRecorder() + written, _ := gin.CreateTestContext(recorder) + c.Writer = written.Writer + info.SetEstimatePromptTokens(10) + response := map[string]any{"id": "resp_terminal", "status": "completed"} + if event == "response.incomplete" { + response["status"] = "incomplete" + } + prompt, completion := 10, service.CountTextToken("accepted output", "gpt-test") + if usageKind != "missing" { + prompt, completion = 4, 2 + if usageKind == "zero" { + prompt, completion = 0, 0 + } + response["usage"] = map[string]int{"input_tokens": prompt, "output_tokens": completion, "total_tokens": prompt + completion} + } + terminal, err := common.Marshal(map[string]any{"type": event, "response": response}) + require.NoError(t, err) + body := "data: {\"type\":\"response.output_text.delta\",\"delta\":\"accepted output\"}\n\ndata: " + string(terminal) + "\n\ndata: " + tail.frame + "\n\n" + usage, apiErr := OaiResponsesStreamHandler(c, info, &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(body))}) + require.Nil(t, apiErr, "a delivered terminal fixes this generation's outcome") + require.NotNil(t, usage) + require.Equal(t, prompt, usage.PromptTokens) + require.Equal(t, completion, usage.CompletionTokens) + require.Equal(t, prompt+completion, usage.TotalTokens) + require.Equal(t, 2, strings.Count(recorder.Body.String(), "event: ")) + require.NotContains(t, recorder.Body.String(), "late") + require.NotContains(t, recorder.Body.String(), "1999") + }) + } + } + } +} diff --git a/relay/chat_stream_billing_test.go b/relay/chat_stream_billing_test.go index 2bb1d770..4a8c8199 100644 --- a/relay/chat_stream_billing_test.go +++ b/relay/chat_stream_billing_test.go @@ -305,16 +305,25 @@ func TestResponsesPartialBillingRecorderPreservesFailureMetricAndRefundBoundary( require.Zero(t, billing.refunds, "partial settlement must prevent a full refund") perfmetrics.RecordRelaySample(info, false, 0) - result, queryErr := perfmetrics.QuerySummaryAll(1, nil) - require.NoError(t, queryErr) - found := false - for _, summary := range result.Models { - if summary.ModelName == info.OriginModelName { - found = true - require.EqualValues(t, 1, summary.RequestCount) - require.Zero(t, summary.SuccessRate) + // Success samples use gopool.Go. Observe a bounded window for an erroneous + // enqueue, keeping all test queries on this goroutine before DB cleanup. + deadline := time.Now().Add(200 * time.Millisecond) + for { + result, queryErr := perfmetrics.QuerySummaryAll(1, nil) + require.NoError(t, queryErr) + found := false + for _, summary := range result.Models { + if summary.ModelName == info.OriginModelName { + found = true + require.EqualValues(t, 1, summary.RequestCount) + require.Zero(t, summary.SuccessRate) + break + } + } + require.True(t, found, "missing partial Responses failure metric") + if !time.Now().Before(deadline) { break } + time.Sleep(10 * time.Millisecond) } - require.True(t, found, "missing partial Responses failure metric") } diff --git a/relay/responses_native_missing_output_settlement_test.go b/relay/responses_native_missing_output_settlement_test.go index 2a23648d..f89e6dee 100644 --- a/relay/responses_native_missing_output_settlement_test.go +++ b/relay/responses_native_missing_output_settlement_test.go @@ -11,8 +11,10 @@ import ( "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/constant" + "github.com/MAX-API-Next/MAX-API/dto" "github.com/MAX-API-Next/MAX-API/model" "github.com/MAX-API-Next/MAX-API/pkg/billingexpr" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" "github.com/MAX-API-Next/MAX-API/service" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" @@ -82,10 +84,16 @@ func nativeMissingOutputCases() []nativeMissingOutputCase { type nativeMissingOutputWriter struct { *httptest.ResponseRecorder short, triggered bool + failOn string + flush bool } func (w *nativeMissingOutputWriter) Write(data []byte) (int, error) { - if strings.Contains(string(data), "unwritten-generated-output") { + failOn := w.failOn + if failOn == "" { + failOn = "unwritten-generated-output" + } + if !w.flush && strings.Contains(string(data), failOn) { w.triggered = true if w.short { return w.ResponseRecorder.Write(data[:3]) @@ -99,6 +107,47 @@ func (w *nativeMissingOutputWriter) WriteString(data string) (int, error) { return w.Write([]byte(data)) } +func (w *nativeMissingOutputWriter) FlushError() error { + if w.flush && strings.Contains(w.Body.String(), w.failOn) { + w.triggered = true + return io.ErrClosedPipe + } + w.ResponseRecorder.Flush() + return nil +} + +func TestNativeResponsesTerminalTailRetainsOriginalFunding(t *testing.T) { + initRequestReplayTests(t) + oldAudit := common.LogResponseContentEnabled + common.LogResponseContentEnabled = true + t.Cleanup(func() { common.LogResponseContentEnabled = oldAudit }) + setFileSearchLedgerPriceForTest(t, 0.02) + for _, event := range []string{"response.completed", "response.incomplete", "response.done"} { + for _, funding := range streamFundingCases() { + for _, usage := range []string{"missing", "known", "zero"} { + t.Run(event+"/"+funding.name+"/"+usage, func(t *testing.T) { + scenario := nativeMissingOutputCase{event, "accepted output", []map[string]any{ + {"type": "response.output_text.delta", "delta": "accepted output"}, + {"type": event, "response": map[string]any{"id": "resp_missing", "status": strings.TrimPrefix(event, "response.")}}, + }} + runNativeMissingOutputLedger(t, scenario, funding, usage, "after_terminal", false) + }) + } + } + } + for _, funding := range streamFundingCases() { + for _, failure := range []string{"terminal_write", "terminal_short", "terminal_flush"} { + t.Run(funding.name+"/"+failure, func(t *testing.T) { + scenario := nativeMissingOutputCase{"response.completed", "accepted output", []map[string]any{ + {"type": "response.output_text.delta", "delta": "accepted output"}, + {"type": "response.completed", "response": map[string]any{"id": "resp_missing", "status": "completed"}}, + }} + runNativeMissingOutputLedger(t, scenario, funding, "known", failure, false) + }) + } + } +} + func TestNativeResponsesMissingOutputRetainsOriginalFunding(t *testing.T) { initRequestReplayTests(t) oldAudit := common.LogResponseContentEnabled @@ -166,7 +215,7 @@ func runNativeMissingOutputLedger(t *testing.T, scenario nativeMissingOutputCase hasTerminal := false if !noDelivery { for _, frame := range scenario.frames { - if frame["type"] == "response.completed" { + if frame["type"] == "response.completed" || frame["type"] == "response.incomplete" || frame["type"] == "response.done" { hasTerminal = true if usage == "terminal_zero" { // A terminal snapshot supplies usage in that one terminal; @@ -190,6 +239,13 @@ func runNativeMissingOutputLedger(t *testing.T, scenario nativeMissingOutputCase frames = append(frames, `{"type":"response.failed","response":{"id":"resp_missing","status":"failed","error":{"code":"server_error","message":"synthetic failure"}}}`) case "write", "short": frames = append(frames, `{"type":"response.`+scenario.name+`.delta","item_id":"item_0","output_index":0,"delta":"unwritten-generated-output"}`) + case "after_terminal", "terminal_write", "terminal_short", "terminal_flush": + frames = append(frames, + `{"type":"response.completed","response":{"id":"late_response","status":"completed","usage":{"input_tokens":1000,"output_tokens":999,"total_tokens":1999}}}`, + `{"type":"response.output_item.done","item":{"type":"file_search_call","id":"late_tool","status":"completed"}}`, + `{"type":"response.output_text.delta","delta":"late output"}`, + `{"type":"response.failed","response":{"status":"failed","error":{"message":"late failure"}}}`, + `{broken`) } var calls atomic.Int32 upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -203,6 +259,12 @@ func runNativeMissingOutputLedger(t *testing.T, scenario nativeMissingOutputCase t.Cleanup(upstream.Close) requestID := "native-missing-" + fmt.Sprintf("%x", len(t.Name())) + "-" + scenario.name c, info, recorder := streamFundingContext(t, "responses_native", requestID, upstream.URL, constant.ChannelTypeOpenAI, funding) + if ending == "after_terminal" { + info.Request.(*dto.OpenAIResponsesRequest).Tools = []byte(`[{"type":"file_search","vector_store_ids":["vs_synthetic"]}]`) + info.ResponsesUsageInfo = &relaycommon.ResponsesUsageInfo{BuiltInTools: map[string]*relaycommon.BuildInToolInfo{ + dto.BuildInToolFileSearch: {ToolName: dto.BuildInToolFileSearch}, + }} + } if len(pricing) > 0 { if pricing[0] == "free" { info.PriceData.ModelRatio = 0 @@ -213,13 +275,17 @@ func runNativeMissingOutputLedger(t *testing.T, scenario nativeMissingOutputCase var writer *nativeMissingOutputWriter if ending == "write" || ending == "short" { writer = &nativeMissingOutputWriter{ResponseRecorder: recorder, short: ending == "short"} + } else if strings.HasPrefix(ending, "terminal_") && ending != "terminal_failed" { + writer = &nativeMissingOutputWriter{ResponseRecorder: recorder, failOn: "event: response.completed", short: ending == "terminal_short", flush: ending == "terminal_flush"} + } + if writer != nil { failed, _ := gin.CreateTestContext(writer) c.Writer = failed.Writer } require.Nil(t, service.PreConsumeBilling(c, 100, info)) apiErr := ResponsesHelper(c, info) legalTerminal := usage == "terminal_zero" || hasTerminal - if legalTerminal && ending == "eof" { + if legalTerminal && (ending == "eof" || ending == "after_terminal") { require.Nil(t, apiErr) } else { require.NotNil(t, apiErr) @@ -228,6 +294,14 @@ func runNativeMissingOutputLedger(t *testing.T, scenario nativeMissingOutputCase require.True(t, writer.triggered) require.NotContains(t, recorder.Body.String(), "unwritten-generated-output") } + if ending == "after_terminal" || strings.HasPrefix(ending, "terminal_") && ending != "terminal_failed" { + require.NotContains(t, recorder.Body.String(), "late") + require.NotContains(t, recorder.Body.String(), "1999") + } + if ending == "after_terminal" { + require.Zero(t, info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolFileSearch].CallCount, "post-terminal tools must not add fees") + require.Equal(t, 1, strings.Count(recorder.Body.String(), "event: "+scenario.name+"\n")) + } service.HandleFailedBilling(c, info, apiErr) service.HandleFailedBilling(c, info, apiErr) require.EqualValues(t, 1, calls.Load()) @@ -265,7 +339,7 @@ func runNativeMissingOutputLedger(t *testing.T, scenario nativeMissingOutputCase require.True(t, effect.QuotaIsActual) } assertStreamFundingBalances(t, db, funding, requestID, quota, !noDelivery, prompt, completion) - if ending == "eof" { + if ending == "eof" || ending == "after_terminal" { var audit strings.Builder for _, frame := range scenario.frames { if frame["type"] == "response.output_text.delta" { diff --git a/service/openaicompat/chat_to_responses_response.go b/service/openaicompat/chat_to_responses_response.go index f1b3f429..896133eb 100644 --- a/service/openaicompat/chat_to_responses_response.go +++ b/service/openaicompat/chat_to_responses_response.go @@ -147,8 +147,10 @@ func ChatCompletionsResponseToResponsesResponseWithCustomTools(resp *dto.OpenAIT } for i, toolCall := range choice.Message.ParseToolCalls() { item := chatToolCallToResponsesOutput(toolCall, id, i, status) - if _, custom := customToolNames[strings.TrimSpace(toolCall.Function.Name)]; custom && item.Type == responsesOutputTypeFunctionCall { + name := strings.TrimSpace(toolCall.Function.Name) + if _, custom := customToolNames[name]; custom && item.Type == responsesOutputTypeFunctionCall { item.Type = responsesOutputTypeCustomToolCall + item.Name = name item.Input = customToolInputFromArguments(toolCall.Function.Arguments) item.Arguments = nil } diff --git a/service/openaicompat/custom_tool_response_test.go b/service/openaicompat/custom_tool_response_test.go index 279ae4ce..8aece9de 100644 --- a/service/openaicompat/custom_tool_response_test.go +++ b/service/openaicompat/custom_tool_response_test.go @@ -74,7 +74,9 @@ func TestResponsesCustomToolWhitespaceNameMatchesStream(t *testing.T) { require.NoError(t, err) require.Len(t, response.Output, 1) require.Equal(t, "custom_tool_call", response.Output[0].Type) - require.Equal(t, name, response.Output[0].Name) + require.Equal(t, "apply_patch", response.Output[0].Name) + _, registered := names[response.Output[0].Name] + require.True(t, registered, "clients must resolve the returned custom name against the tool registry") require.Equal(t, "raw patch", response.Output[0].Input) require.Nil(t, response.Output[0].Arguments) require.Equal(t, "call_custom", response.Output[0].CallId) @@ -92,7 +94,17 @@ func TestResponsesCustomToolWhitespaceNameMatchesStream(t *testing.T) { require.NotNil(t, stream) require.Equal(t, response.Output[0].Type, stream.Output[0].Type) require.Equal(t, "apply_patch", stream.Output[0].Name) + require.Equal(t, response.Output[0].Name, stream.Output[0].Name) require.Equal(t, response.Output[0].Input, stream.Output[0].Input) + + tool.Function.Name = " lookup " + message.SetToolCalls([]dto.ToolCallResponse{tool}) + ordinary, _, err := ChatCompletionsResponseToResponsesResponseWithCustomTools(&dto.OpenAITextResponse{ + Choices: []dto.OpenAITextResponseChoice{{Message: message, FinishReason: "tool_calls"}}, + }, "resp_function", names) + require.NoError(t, err) + require.Equal(t, "function_call", ordinary.Output[0].Type) + require.Equal(t, " lookup ", ordinary.Output[0].Name, "custom restoration must preserve ordinary function behavior") }) } } From 1f82e5ad05aabfefbdcbd48a505e0a0090c9bdb0 Mon Sep 17 00:00:00 2001 From: CSCITech Date: Tue, 13 Oct 2026 11:49:38 +0800 Subject: [PATCH 17/24] test: honor short mode in billing process recovery matrices --- service/billing_playground_process_test.go | 9 +++++++++ service/billing_process_recovery_test.go | 18 ++++++++++++++++++ service/billing_separate_log_process_test.go | 18 ++++++++++++++++++ 3 files changed, 45 insertions(+) diff --git a/service/billing_playground_process_test.go b/service/billing_playground_process_test.go index 28e2fc33..b49a1719 100644 --- a/service/billing_playground_process_test.go +++ b/service/billing_playground_process_test.go @@ -9,6 +9,9 @@ import ( ) func TestBillingPlaygroundProcessCrashRecovery(t *testing.T) { + if testing.Short() { + t.Skip("OS-process recovery matrix; run without -short") + } funding := billingProcessFundingCases()[0] phases := []string{"intent_committed", "funding_before_commit", "funding_committed", "receipt_before_log", "log_before_commit", "log_committed", "usage_before_commit", "complete"} @@ -33,6 +36,9 @@ func TestBillingPlaygroundProcessCrashRecovery(t *testing.T) { } func TestBillingPlaygroundProcessRefundRecovery(t *testing.T) { + if testing.Short() { + t.Skip("OS-process recovery matrix; run without -short") + } funding := billingProcessFundingCases()[0] for _, phase := range []string{"intent_committed", "funding_before_commit", "complete"} { t.Run(phase, func(t *testing.T) { @@ -51,6 +57,9 @@ func TestBillingPlaygroundProcessRefundRecovery(t *testing.T) { } func TestBillingPlaygroundProcessAdmissionRecovery(t *testing.T) { + if testing.Short() { + t.Skip("OS-process recovery matrix; run without -short") + } funding := billingProcessFundingCases()[0] for _, phase := range []string{"intent_committed", "funding_before_commit", "complete"} { t.Run(phase, func(t *testing.T) { diff --git a/service/billing_process_recovery_test.go b/service/billing_process_recovery_test.go index a2e09871..df6b908a 100644 --- a/service/billing_process_recovery_test.go +++ b/service/billing_process_recovery_test.go @@ -655,6 +655,9 @@ func billingProcessDatabasesSnapshot(t *testing.T, db, logDB *gorm.DB) [32]byte } func TestBillingProcessCrashRecovery(t *testing.T) { + if testing.Short() { + t.Skip("OS-process recovery matrix; run without -short") + } phases := []string{"intent_committed", "funding_before_commit", "funding_committed", "receipt_before_log", "log_before_commit", "log_committed", "usage_before_commit", "complete"} for _, funding := range billingProcessFundingCases() { @@ -681,6 +684,9 @@ func TestBillingProcessCrashRecovery(t *testing.T) { } func TestBillingProcessRefundRecovery(t *testing.T) { + if testing.Short() { + t.Skip("OS-process recovery matrix; run without -short") + } for _, funding := range billingProcessFundingCases() { for _, phase := range []string{"intent_committed", "funding_before_commit", "complete"} { t.Run(funding.name+"/"+phase, func(t *testing.T) { @@ -700,6 +706,9 @@ func TestBillingProcessRefundRecovery(t *testing.T) { } func TestBillingProcessRepeatedCrashRecovery(t *testing.T) { + if testing.Short() { + t.Skip("OS-process recovery matrix; run without -short") + } for _, funding := range billingProcessFundingCases() { for _, quota := range []int{130, 0} { t.Run(fmt.Sprintf("%s/quota_%d", funding.name, quota), func(t *testing.T) { @@ -725,6 +734,9 @@ func TestBillingProcessRepeatedCrashRecovery(t *testing.T) { } func TestBillingProcessConcurrentRecoveryAndRefund(t *testing.T) { + if testing.Short() { + t.Skip("OS-process recovery matrix; run without -short") + } for _, funding := range billingProcessFundingCases() { for _, quota := range []int{130, 0} { t.Run(fmt.Sprintf("%s/quota_%d", funding.name, quota), func(t *testing.T) { @@ -777,6 +789,9 @@ func TestBillingProcessConcurrentRecoveryAndRefund(t *testing.T) { } func TestBillingProcessUnsupportedStateIsNotRetried(t *testing.T) { + if testing.Short() { + t.Skip("OS-process recovery matrix; run without -short") + } for _, funding := range billingProcessFundingCases()[:2] { t.Run(funding.name, func(t *testing.T) { dir, db := seedBillingProcessDB(t, funding) @@ -796,6 +811,9 @@ func TestBillingProcessUnsupportedStateIsNotRetried(t *testing.T) { } func TestBillingProcessSubscriptionPeriodRecovery(t *testing.T) { + if testing.Short() { + t.Skip("OS-process recovery matrix; run without -short") + } funding := billingProcessFundingCases()[1] for _, phase := range []string{"intent_committed", "funding_committed"} { t.Run(phase, func(t *testing.T) { diff --git a/service/billing_separate_log_process_test.go b/service/billing_separate_log_process_test.go index ddb6d384..78ce8cb5 100644 --- a/service/billing_separate_log_process_test.go +++ b/service/billing_separate_log_process_test.go @@ -31,6 +31,9 @@ func (c billingSeparateLogCase) config(quota int, phase string) billingProcessCo } func TestBillingSeparateLogProcessCrashRecovery(t *testing.T) { + if testing.Short() { + t.Skip("OS-process recovery matrix; run without -short") + } // The first three barriers belong to the independent log transaction; the // last two belong to the business transaction that applies usage counters. phases := []string{"funding_committed", "receipt_before_log", "log_before_commit", "log_committed", "usage_before_commit"} @@ -62,6 +65,9 @@ func TestBillingSeparateLogProcessCrashRecovery(t *testing.T) { } func TestBillingSeparateLogProcessRepeatedCrashRecovery(t *testing.T) { + if testing.Short() { + t.Skip("OS-process recovery matrix; run without -short") + } for _, c := range billingSeparateLogCases()[:2] { for _, quota := range []int{130, 0} { t.Run(fmt.Sprintf("%s/quota_%d", c.name, quota), func(t *testing.T) { @@ -88,6 +94,9 @@ func TestBillingSeparateLogProcessRepeatedCrashRecovery(t *testing.T) { } func TestBillingSeparateLogProcessUnavailableRecovery(t *testing.T) { + if testing.Short() { + t.Skip("OS-process recovery matrix; run without -short") + } for _, c := range billingSeparateLogCases() { for _, quota := range []int{130, 0} { t.Run(fmt.Sprintf("%s/quota_%d", c.name, quota), func(t *testing.T) { @@ -117,6 +126,9 @@ func TestBillingSeparateLogProcessUnavailableRecovery(t *testing.T) { } func TestBillingSeparateLogProcessConcurrentRecoveryAndRefund(t *testing.T) { + if testing.Short() { + t.Skip("OS-process recovery matrix; run without -short") + } cases := billingSeparateLogCases() for _, c := range []billingSeparateLogCase{cases[0], cases[1], cases[4]} { for _, quota := range []int{130, 0} { @@ -150,6 +162,9 @@ func TestBillingSeparateLogProcessConcurrentRecoveryAndRefund(t *testing.T) { } func TestBillingSeparateLogProcessRefundRecovery(t *testing.T) { + if testing.Short() { + t.Skip("OS-process recovery matrix; run without -short") + } for _, c := range billingSeparateLogCases() { for _, phase := range []string{"intent_committed", "funding_before_commit", "complete"} { t.Run(c.name+"/"+phase, func(t *testing.T) { @@ -170,6 +185,9 @@ func TestBillingSeparateLogProcessRefundRecovery(t *testing.T) { } func TestBillingSeparateLogProcessSubscriptionPeriodRecovery(t *testing.T) { + if testing.Short() { + t.Skip("OS-process recovery matrix; run without -short") + } c := billingSeparateLogCases()[1] for _, quota := range []int{130, 0} { t.Run(fmt.Sprintf("quota_%d", quota), func(t *testing.T) { From a0cce548bb70b6f83964bd4d38dacf9c2a8b877f Mon Sep 17 00:00:00 2001 From: CSCITech Date: Tue, 13 Oct 2026 13:25:26 +0800 Subject: [PATCH 18/24] test: clean up Playground callback after assertion failures --- service/billing_playground_atomic_test.go | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/service/billing_playground_atomic_test.go b/service/billing_playground_atomic_test.go index 79e44a80..16709d6b 100644 --- a/service/billing_playground_atomic_test.go +++ b/service/billing_playground_atomic_test.go @@ -194,10 +194,17 @@ func TestPlaygroundWalletDurableReservationRefundAndRollback(t *testing.T) { tx.AddError(errors.New("synthetic Playground funding transaction failure")) } })) + callbackRemoved := false + t.Cleanup(func() { + if !callbackRemoved { + _ = model.DB.Callback().Update().Remove(callback) + } + }) c, info := playgroundAtomicContext("playground-rollback", true) _, apiErr := NewBillingSession(c, info, 10) require.NotNil(t, apiErr) require.NoError(t, model.DB.Callback().Update().Remove(callback)) + callbackRemoved = true require.EqualValues(t, 100, getUserQuota(t, 991)) require.EqualValues(t, 100, getTokenRemainQuota(t, 992)) var count int64 From 5f7b8a5cec4890d53be72eb31c70dfca7576fe0a Mon Sep 17 00:00:00 2001 From: CSCITech Date: Tue, 13 Oct 2026 14:24:55 +0800 Subject: [PATCH 19/24] test: assert upstream cancellation in settlement fixtures --- relay/metadata_funding_settlement_test.go | 11 +++++++++++ relay/responses_first_result_settlement_test.go | 9 +++++++++ 2 files changed, 20 insertions(+) diff --git a/relay/metadata_funding_settlement_test.go b/relay/metadata_funding_settlement_test.go index bec67668..3ec8b8c8 100644 --- a/relay/metadata_funding_settlement_test.go +++ b/relay/metadata_funding_settlement_test.go @@ -119,6 +119,7 @@ func TestMetadataOnlyFailureRetainsOriginalFunding(t *testing.T) { frames = append(frames, `{broken`) } var calls atomic.Int32 + upstreamCanceled := make(chan bool, 1) upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { calls.Add(1) w.Header().Set("Content-Type", "text/event-stream") @@ -130,7 +131,9 @@ func TestMetadataOnlyFailureRetainsOriginalFunding(t *testing.T) { if ending == "timeout" || ending == "cancel" { select { case <-r.Context().Done(): + upstreamCanceled <- true case <-time.After(3 * time.Second): + upstreamCanceled <- false } } })) @@ -171,6 +174,14 @@ func TestMetadataOnlyFailureRetainsOriginalFunding(t *testing.T) { apiErr = ResponsesHelper(c, info) } require.NotNil(t, apiErr) + if ending == "cancel" || ending == "timeout" { + select { + case canceled := <-upstreamCanceled: + require.True(t, canceled, "cancellation or timeout must reach the upstream request") + case <-time.After(3 * time.Second): + t.Fatal("upstream handler did not finish within its cancellation guard") + } + } require.False(t, info.HasRecordedChannelFirstResult()) require.True(t, info.FirstResultTime.IsZero()) if bufferedMetadata { diff --git a/relay/responses_first_result_settlement_test.go b/relay/responses_first_result_settlement_test.go index 393ce508..132ceb33 100644 --- a/relay/responses_first_result_settlement_test.go +++ b/relay/responses_first_result_settlement_test.go @@ -125,6 +125,7 @@ func TestResponsesFirstResultTimeoutSettlesDeliveredWhitespace(t *testing.T) { StartTime: time.Now().Add(-time.Hour).Unix(), EndTime: time.Now().Add(time.Hour).Unix()}).Error) } var calls atomic.Int32 + upstreamCanceled := make(chan bool, 1) upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { calls.Add(1) w.Header().Set("Content-Type", "text/event-stream") @@ -132,7 +133,9 @@ func TestResponsesFirstResultTimeoutSettlesDeliveredWhitespace(t *testing.T) { w.(http.Flusher).Flush() select { case <-r.Context().Done(): + upstreamCanceled <- true case <-time.After(3 * time.Second): + upstreamCanceled <- false } })) t.Cleanup(upstream.Close) @@ -164,6 +167,12 @@ func TestResponsesFirstResultTimeoutSettlesDeliveredWhitespace(t *testing.T) { apiErr := ResponsesHelper(c, info) require.NotNil(t, apiErr) require.Equal(t, types.ErrorCodeChannelResponseTimeExceeded, apiErr.GetErrorCode()) + select { + case canceled := <-upstreamCanceled: + require.True(t, canceled, "first-result timeout must reach the upstream request") + case <-time.After(3 * time.Second): + t.Fatal("upstream handler did not finish within its cancellation guard") + } require.True(t, types.IsSkipRetryError(apiErr)) require.Equal(t, relaycommon.StreamEndReasonTimeout, info.StreamStatus.EndReason) require.False(t, info.HasRecordedChannelFirstResult()) From 0b3fd93fb5c3baaffa33e1e2a0dd87404b783c2e Mon Sep 17 00:00:00 2001 From: CSCITech Date: Tue, 13 Oct 2026 15:39:22 +0800 Subject: [PATCH 20/24] fix: preserve native chat for Zhipu coding-plan aliases --- relay/chat_completions_via_responses.go | 2 +- ...chat_responses_protocol_settlement_test.go | 128 ++++++++++++++++-- relay/compatible_handler.go | 3 +- service/openai_chat_responses_mode.go | 4 + service/openaicompat/policy.go | 14 ++ 5 files changed, 140 insertions(+), 11 deletions(-) diff --git a/relay/chat_completions_via_responses.go b/relay/chat_completions_via_responses.go index 0f4a376b..7426b4eb 100644 --- a/relay/chat_completions_via_responses.go +++ b/relay/chat_completions_via_responses.go @@ -71,7 +71,7 @@ func applySystemPromptIfNeeded(c *gin.Context, info *relaycommon.RelayInfo, requ } func chatCompletionsViaResponses(c *gin.Context, info *relaycommon.RelayInfo, adaptor channel.Adaptor, request *dto.GeneralOpenAIRequest) (*dto.Usage, *types.MaxAPIError) { - if info == nil || adaptor == nil || !service.SupportsOpenAIResponsesUpstream(info.ChannelType) { + if info == nil || adaptor == nil || !service.SupportsOpenAIResponsesUpstreamForBaseURL(info.ChannelType, info.ChannelBaseUrl) { return nil, types.NewError(fmt.Errorf("channel does not use OpenAI Responses upstream"), types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) } chatJSON, err := common.Marshal(request) diff --git a/relay/chat_responses_protocol_settlement_test.go b/relay/chat_responses_protocol_settlement_test.go index bb854d3d..5f16b79d 100644 --- a/relay/chat_responses_protocol_settlement_test.go +++ b/relay/chat_responses_protocol_settlement_test.go @@ -14,11 +14,14 @@ import ( "github.com/MAX-API-Next/MAX-API/constant" "github.com/MAX-API-Next/MAX-API/dto" "github.com/MAX-API-Next/MAX-API/model" + "github.com/MAX-API-Next/MAX-API/relay/channel" "github.com/MAX-API-Next/MAX-API/relay/channel/claude" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" relayconstant "github.com/MAX-API-Next/MAX-API/relay/constant" "github.com/MAX-API-Next/MAX-API/service" "github.com/MAX-API-Next/MAX-API/setting/model_setting" "github.com/MAX-API-Next/MAX-API/types" + "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) @@ -27,18 +30,22 @@ func TestGlobalChatResponsesProtocolRetainsOriginalFunding(t *testing.T) { providers := []struct { name, wire, converter string channelType int + baseAlias string }{ - {"claude", "claude", "", constant.ChannelTypeAnthropic}, - {"gemini", "gemini", "", constant.ChannelTypeGemini}, - {"openai", "responses", "", constant.ChannelTypeOpenAI}, - {"custom_none", "responses", dto.AdvancedCustomConverterNone, constant.ChannelTypeAdvancedCustom}, - {"custom_claude", "claude", dto.AdvancedCustomConverterOpenAIChatCompletionsToAnthropicMessages, constant.ChannelTypeAdvancedCustom}, - {"custom_gemini", "gemini", dto.AdvancedCustomConverterOpenAIChatCompletionsToGeminiGenerateContent, constant.ChannelTypeAdvancedCustom}, - {"custom_responses", "responses", dto.AdvancedCustomConverterOpenAIChatCompletionsToOpenAIResponses, constant.ChannelTypeAdvancedCustom}, + {"claude", "claude", "", constant.ChannelTypeAnthropic, ""}, + {"gemini", "gemini", "", constant.ChannelTypeGemini, ""}, + {"openai", "responses", "", constant.ChannelTypeOpenAI, ""}, + {"custom_none", "responses", dto.AdvancedCustomConverterNone, constant.ChannelTypeAdvancedCustom, ""}, + {"custom_claude", "claude", dto.AdvancedCustomConverterOpenAIChatCompletionsToAnthropicMessages, constant.ChannelTypeAdvancedCustom, ""}, + {"custom_gemini", "gemini", dto.AdvancedCustomConverterOpenAIChatCompletionsToGeminiGenerateContent, constant.ChannelTypeAdvancedCustom, ""}, + {"custom_responses", "responses", dto.AdvancedCustomConverterOpenAIChatCompletionsToOpenAIResponses, constant.ChannelTypeAdvancedCustom, ""}, + {"zhipu_standard", "responses", "", constant.ChannelTypeZhipu_v4, ""}, + {"zhipu_coding", "chat_native", "", constant.ChannelTypeZhipu_v4, "glm-coding-plan"}, + {"zhipu_coding_international", "chat_native", "", constant.ChannelTypeZhipu_v4, "glm-coding-plan-international"}, } for _, provider := range providers { modes := []bool{true} - if provider.name == "claude" || provider.name == "gemini" { + if provider.name == "claude" || provider.name == "gemini" || provider.baseAlias != "" { modes = []bool{false, true} } for _, enabled := range modes { @@ -72,6 +79,15 @@ func TestGlobalChatResponsesProtocolRetainsOriginalFunding(t *testing.T) { t.Cleanup(upstream.Close) requestID := "global-protocol-" + provider.name c, info, recorder := streamFundingContext(t, "chat_native", requestID, upstream.URL, provider.channelType, funding) + if provider.baseAlias != "" { + original := constant.ChannelSpecialBases[provider.baseAlias] + synthetic := original + synthetic.OpenAIBaseURL = upstream.URL + "/coding" + constant.ChannelSpecialBases[provider.baseAlias] = synthetic + t.Cleanup(func() { constant.ChannelSpecialBases[provider.baseAlias] = original }) + c.Set(string(constant.ContextKeyChannelBaseUrl), provider.baseAlias) + // The adaptor still resolves the real alias, but its HTTP target is local. + } info.IsStream = stream info.Request.(*dto.GeneralOpenAIRequest).Stream = common.GetPointer(stream) model_setting.GetGlobalSettings().ChatCompletionsToResponsesPolicy = model_setting.ChatCompletionsToResponsesPolicy{Enabled: enabled, AllChannels: true, ModelPatterns: []string{"^gpt-test$"}} @@ -102,8 +118,16 @@ func TestGlobalChatResponsesProtocolRetainsOriginalFunding(t *testing.T) { require.Contains(t, strings.ToLower(request.path), "generatecontent") require.Contains(t, request.body, "contents") require.NotContains(t, request.body, "input") + case "chat_native": + require.Equal(t, "/coding/chat/completions", request.path) + require.Contains(t, request.body, "messages") + require.NotContains(t, request.body, "input") default: - require.Equal(t, "/v1/responses", request.path) + path := "/v1/responses" + if provider.channelType == constant.ChannelTypeZhipu_v4 { + path = "/api/v1/responses" + } + require.Equal(t, path, request.path) require.Contains(t, request.body, "input") require.NotContains(t, request.body, "messages") } @@ -218,10 +242,85 @@ func TestGlobalResponsesBridgeRejectsNativeBeforeUpstreamAndRefunds(t *testing.T } } +func TestGlobalResponsesBridgeZhipuEndpointEligibility(t *testing.T) { + initRequestReplayTests(t) + bases := []string{"", "https://open.bigmodel.cn", "https://synthetic.invalid"} + for alias := range constant.ChannelSpecialBases { + bases = append(bases, alias) + } + for _, base := range bases { + t.Run(base, func(t *testing.T) { + c, info, _ := streamFundingContext(t, "chat_native", "zhipu-capability", base, constant.ChannelTypeZhipu_v4, streamFundingCases()[3]) + info.InitChannelMeta(c) + adaptor := GetAdaptor(info.ApiType) + adaptor.Init(info) + info.RelayMode = relayconstant.RelayModeResponses + _, err := adaptor.GetRequestURL(info) + info.RelayMode = relayconstant.RelayModeChatCompletions + model_setting.GetGlobalSettings().ChatCompletionsToResponsesPolicy = model_setting.ChatCompletionsToResponsesPolicy{Enabled: true, AllChannels: true, ModelPatterns: []string{"^gpt-test$"}} + require.Equal(t, err == nil, shouldUseGlobalResponsesBridge(c, info), "eligibility must match the actual adaptor URL capability") + }) + } +} + +func TestDirectResponsesBridgeRejectsZhipuCodingAliasesAndRefunds(t *testing.T) { + initRequestReplayTests(t) + for alias := range constant.ChannelSpecialBases { + for _, funding := range streamFundingCases() { + t.Run(alias+"/"+funding.name, func(t *testing.T) { + db := setupStreamFundingLedger(t, funding) + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { calls.Add(1) })) + t.Cleanup(upstream.Close) + original := constant.ChannelSpecialBases[alias] + synthetic := original + synthetic.OpenAIBaseURL = upstream.URL + constant.ChannelSpecialBases[alias] = synthetic + t.Cleanup(func() { constant.ChannelSpecialBases[alias] = original }) + requestID := "rejected-zhipu-coding" + c, info, recorder := streamFundingContext(t, "chat_native", requestID, alias, constant.ChannelTypeZhipu_v4, funding) + info.InitChannelMeta(c) + info.RequestURLPath = "/v1/chat/completions" + adaptor := &responsesConversionProbe{Adaptor: GetAdaptor(info.ApiType)} + adaptor.Init(info) + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + usage, apiErr := chatCompletionsViaResponses(c, info, adaptor, info.Request.(*dto.GeneralOpenAIRequest)) + require.Zero(t, adaptor.conversions, "unsupported endpoint must be rejected before converting the request") + require.Nil(t, usage) + require.NotNil(t, apiErr) + require.Equal(t, types.ErrorCodeConvertRequestFailed, apiErr.GetErrorCode()) + require.True(t, types.IsSkipRetryError(apiErr)) + require.Zero(t, calls.Load()) + require.Empty(t, recorder.Body.String()) + require.Equal(t, relayconstant.RelayModeChatCompletions, info.RelayMode) + require.Equal(t, "/v1/chat/completions", info.RequestURLPath) + service.HandleFailedBilling(c, info, apiErr) + service.HandleFailedBilling(c, info, apiErr) + assertStreamFundingBalances(t, db, funding, requestID, 0, false, 0, 0) + before := requestReplaySnapshot(t, db) + model.ProcessPendingBillingSettlementsOnce() + require.Equal(t, before, requestReplaySnapshot(t, db)) + }) + } + } +} + +type responsesConversionProbe struct { + channel.Adaptor + conversions int +} + +func (a *responsesConversionProbe) ConvertOpenAIResponsesRequest(c *gin.Context, info *relaycommon.RelayInfo, request dto.OpenAIResponsesRequest) (any, error) { + a.conversions++ + return a.Adaptor.ConvertOpenAIResponsesRequest(c, info, request) +} + func globalBridgeWireBody(t *testing.T, wire string, stream bool) string { t.Helper() if !stream { switch wire { + case "chat_native": + return `{"id":"chat_funding","object":"chat.completion","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"partial output"},"finish_reason":"stop"}],"usage":{"prompt_tokens":4,"completion_tokens":2,"total_tokens":6}}` case "claude": return `{"id":"msg_funding","type":"message","role":"assistant","model":"gpt-test","content":[{"type":"text","text":"partial output"}],"stop_reason":"end_turn","usage":{"input_tokens":4,"output_tokens":2}}` case "gemini": @@ -232,6 +331,17 @@ func globalBridgeWireBody(t *testing.T, wire string, stream bool) string { } var frames []string switch wire { + case "chat_native": + _, frames = streamFundingFrames(t, "chat_native", "known") + for i, frame := range frames { + var chunk map[string]any + require.NoError(t, common.UnmarshalJsonStr(frame, &chunk)) + chunk["object"] = "chat.completion.chunk" + data, err := common.Marshal(chunk) + require.NoError(t, err) + frames[i] = string(data) + } + frames = append(frames, "[DONE]") case "claude": _, frames = streamFundingFrames(t, "claude_to_responses", "known") frames = append(frames, `{"type":"message_stop"}`) diff --git a/relay/compatible_handler.go b/relay/compatible_handler.go index 482935c9..d4365d82 100644 --- a/relay/compatible_handler.go +++ b/relay/compatible_handler.go @@ -255,7 +255,8 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (maxAPIError *types } func shouldUseGlobalResponsesBridge(c *gin.Context, info *relaycommon.RelayInfo) bool { - if info == nil || !service.ShouldChatCompletionsUseResponsesGlobal(info.ChannelId, info.ChannelType, info.OriginModelName) { + if info == nil || !service.SupportsOpenAIResponsesUpstreamForBaseURL(info.ChannelType, info.ChannelBaseUrl) || + !service.ShouldChatCompletionsUseResponsesGlobal(info.ChannelId, info.ChannelType, info.OriginModelName) { return false } if info.ChannelType != constant.ChannelTypeAdvancedCustom { diff --git a/service/openai_chat_responses_mode.go b/service/openai_chat_responses_mode.go index 6e727023..dfc11c0f 100644 --- a/service/openai_chat_responses_mode.go +++ b/service/openai_chat_responses_mode.go @@ -16,3 +16,7 @@ func ShouldChatCompletionsUseResponsesGlobal(channelID int, channelType int, mod func SupportsOpenAIResponsesUpstream(channelType int) bool { return openaicompat.SupportsOpenAIResponsesUpstream(channelType) } + +func SupportsOpenAIResponsesUpstreamForBaseURL(channelType int, baseURL string) bool { + return openaicompat.SupportsOpenAIResponsesUpstreamForBaseURL(channelType, baseURL) +} diff --git a/service/openaicompat/policy.go b/service/openaicompat/policy.go index 1c8406ed..77951977 100644 --- a/service/openaicompat/policy.go +++ b/service/openaicompat/policy.go @@ -22,6 +22,20 @@ func SupportsOpenAIResponsesUpstream(channelType int) bool { } } +// Endpoint aliases can narrow a Responses-capable adaptor's support. Use the +// same alias registry as its URL resolver, while retaining standard endpoints. +func SupportsOpenAIResponsesUpstreamForBaseURL(channelType int, baseURL string) bool { + if !SupportsOpenAIResponsesUpstream(channelType) { + return false + } + apiType, _ := common.ChannelType2APIType(channelType) + if apiType == constant.APITypeZhipuV4 { + _, codingPlan := constant.ChannelSpecialBases[baseURL] + return !codingPlan + } + return true +} + func ShouldChatCompletionsUseResponsesPolicy(policy model_setting.ChatCompletionsToResponsesPolicy, channelID int, channelType int, model string) bool { if !SupportsOpenAIResponsesUpstream(channelType) || !policy.IsChannelEnabled(channelID, channelType) { return false From 195a7bad65330a599e3ccdeb090446309c52f2af Mon Sep 17 00:00:00 2001 From: CSCITech Date: Tue, 13 Oct 2026 17:11:09 +0800 Subject: [PATCH 21/24] fix: align search settlement and Responses bridge routing --- relay/channel/openai/empty_completion.go | 6 +- relay/channel/openai/empty_completion_test.go | 13 ++ relay/chat_completions_via_responses.go | 9 +- ...chat_responses_protocol_settlement_test.go | 163 +++++++++++------- relay/claude_handler.go | 1 + relay/file_search_stream_settlement_test.go | 162 +++++++++++------ relay/helper/common.go | 2 +- relay/metadata_funding_settlement_test.go | 9 +- service/openaicompat/policy.go | 2 +- service/text_quota.go | 37 +++- service/tool_usage_billing_test.go | 53 ++++++ 11 files changed, 333 insertions(+), 124 deletions(-) diff --git a/relay/channel/openai/empty_completion.go b/relay/channel/openai/empty_completion.go index bb7eec8e..525517e8 100644 --- a/relay/channel/openai/empty_completion.go +++ b/relay/channel/openai/empty_completion.go @@ -236,7 +236,7 @@ func responsesStreamHasVisibleOutput(streamResp *dto.ResponsesStreamResponse) bo case "response.custom_tool_call_input.delta": return strings.TrimSpace(streamResp.Delta) != "" case dto.ResponsesOutputTypeItemAdded, dto.ResponsesOutputTypeItemDone: - if streamResp.Item != nil && streamResp.Item.Type == dto.BuildInCallFileSearchCall { + if streamResp.Item != nil && (streamResp.Item.Type == dto.BuildInCallFileSearchCall || streamResp.Item.Type == dto.BuildInCallWebSearchCall) { return helper.ResponsesStreamEventHasOutput(*streamResp) } return responsesOutputHasVisiblePayload(streamResp.Item) @@ -252,8 +252,10 @@ func responsesOutputHasVisiblePayload(output *dto.ResponsesOutput) bool { return false } switch output.Type { - case "function_call", dto.BuildInCallCustomToolCall, dto.ResponsesOutputTypeImageGenerationCall, dto.BuildInCallWebSearchCall: + case "function_call", dto.BuildInCallCustomToolCall, dto.ResponsesOutputTypeImageGenerationCall: return true + case dto.BuildInCallWebSearchCall: + return isBillableResponsesToolStatus(output.Status) case "message": if output.Role != "" && output.Role != "assistant" { return false diff --git a/relay/channel/openai/empty_completion_test.go b/relay/channel/openai/empty_completion_test.go index cb5694f3..f0b6ce2f 100644 --- a/relay/channel/openai/empty_completion_test.go +++ b/relay/channel/openai/empty_completion_test.go @@ -131,6 +131,19 @@ func TestIsEmptyResponsesCompletion(t *testing.T) { } } +func TestResponsesWebSearchLifecycleRequiresCompletedItem(t *testing.T) { + for _, status := range []string{"in_progress", "searching", "failed", "incomplete", "future_status", "completed", ""} { + for _, eventType := range []string{dto.ResponsesOutputTypeItemAdded, dto.ResponsesOutputTypeItemDone} { + t.Run(eventType+"/"+status, func(t *testing.T) { + item := &dto.ResponsesOutput{Type: dto.BuildInCallWebSearchCall, Status: status} + expected := eventType == dto.ResponsesOutputTypeItemDone && (status == "completed" || status == "") + require.Equal(t, expected, responsesStreamHasVisibleOutput(&dto.ResponsesStreamResponse{Type: eventType, Item: item})) + require.Equal(t, status == "completed" || status == "", responsesHasVisibleOutput(&dto.OpenAIResponsesResponse{Output: []dto.ResponsesOutput{*item}})) + }) + } + } +} + func TestShouldRetryEmptyCompletionOnlyOnce(t *testing.T) { gin.SetMode(gin.TestMode) c, _ := gin.CreateTestContext(httptest.NewRecorder()) diff --git a/relay/chat_completions_via_responses.go b/relay/chat_completions_via_responses.go index 7426b4eb..0e56a970 100644 --- a/relay/chat_completions_via_responses.go +++ b/relay/chat_completions_via_responses.go @@ -153,7 +153,14 @@ func chatCompletionsViaResponses(c *gin.Context, info *relaycommon.RelayInfo, ad var requestBody io.Reader = body var httpResp *http.Response - resp, err := adaptor.DoRequest(c, info, requestBody) + // URL/header selection describes the converted upstream protocol. Restore + // the client format before converting the reply or recording settlement. + resp, err := func() (any, error) { + savedFormat := info.RelayFormat + info.RelayFormat = types.RelayFormatOpenAIResponses + defer func() { info.RelayFormat = savedFormat }() + return adaptor.DoRequest(c, info, requestBody) + }() if err != nil { return nil, types.NewOpenAIError(err, types.ErrorCodeDoRequestFailed, http.StatusInternalServerError) } diff --git a/relay/chat_responses_protocol_settlement_test.go b/relay/chat_responses_protocol_settlement_test.go index 5f16b79d..89de9e81 100644 --- a/relay/chat_responses_protocol_settlement_test.go +++ b/relay/chat_responses_protocol_settlement_test.go @@ -5,6 +5,7 @@ import ( "io" "net/http" "net/http/httptest" + "net/url" "strings" "sync/atomic" "testing" @@ -31,17 +32,27 @@ func TestGlobalChatResponsesProtocolRetainsOriginalFunding(t *testing.T) { name, wire, converter string channelType int baseAlias string + claudeClient bool }{ - {"claude", "claude", "", constant.ChannelTypeAnthropic, ""}, - {"gemini", "gemini", "", constant.ChannelTypeGemini, ""}, - {"openai", "responses", "", constant.ChannelTypeOpenAI, ""}, - {"custom_none", "responses", dto.AdvancedCustomConverterNone, constant.ChannelTypeAdvancedCustom, ""}, - {"custom_claude", "claude", dto.AdvancedCustomConverterOpenAIChatCompletionsToAnthropicMessages, constant.ChannelTypeAdvancedCustom, ""}, - {"custom_gemini", "gemini", dto.AdvancedCustomConverterOpenAIChatCompletionsToGeminiGenerateContent, constant.ChannelTypeAdvancedCustom, ""}, - {"custom_responses", "responses", dto.AdvancedCustomConverterOpenAIChatCompletionsToOpenAIResponses, constant.ChannelTypeAdvancedCustom, ""}, - {"zhipu_standard", "responses", "", constant.ChannelTypeZhipu_v4, ""}, - {"zhipu_coding", "chat_native", "", constant.ChannelTypeZhipu_v4, "glm-coding-plan"}, - {"zhipu_coding_international", "chat_native", "", constant.ChannelTypeZhipu_v4, "glm-coding-plan-international"}, + {"claude", "claude", "", constant.ChannelTypeAnthropic, "", false}, + {"gemini", "gemini", "", constant.ChannelTypeGemini, "", false}, + {"openai", "responses", "", constant.ChannelTypeOpenAI, "", false}, + {"custom_none", "responses", dto.AdvancedCustomConverterNone, constant.ChannelTypeAdvancedCustom, "", false}, + {"custom_claude", "claude", dto.AdvancedCustomConverterOpenAIChatCompletionsToAnthropicMessages, constant.ChannelTypeAdvancedCustom, "", false}, + {"custom_gemini", "gemini", dto.AdvancedCustomConverterOpenAIChatCompletionsToGeminiGenerateContent, constant.ChannelTypeAdvancedCustom, "", false}, + {"custom_responses", "responses", dto.AdvancedCustomConverterOpenAIChatCompletionsToOpenAIResponses, constant.ChannelTypeAdvancedCustom, "", false}, + {"zhipu_standard", "responses", "", constant.ChannelTypeZhipu_v4, "", false}, + {"zhipu_coding", "chat_native", "", constant.ChannelTypeZhipu_v4, "glm-coding-plan", false}, + {"zhipu_coding_international", "chat_native", "", constant.ChannelTypeZhipu_v4, "glm-coding-plan-international", false}, + {"volc_standard", "responses", "", constant.ChannelTypeVolcEngine, "", false}, + {"volc_kimi_coding", "chat_native", "", constant.ChannelTypeVolcEngine, "kimi-coding-plan", false}, + {"volc_doubao_coding", "chat_native", "", constant.ChannelTypeVolcEngine, "doubao-coding-plan", false}, + {"claude_zhipu_standard", "responses", "", constant.ChannelTypeZhipu_v4, "", true}, + {"claude_volc_standard", "responses", "", constant.ChannelTypeVolcEngine, "", true}, + {"claude_zhipu_coding", "claude", "", constant.ChannelTypeZhipu_v4, "glm-coding-plan", true}, + {"claude_zhipu_coding_international", "claude", "", constant.ChannelTypeZhipu_v4, "glm-coding-plan-international", true}, + {"claude_volc_kimi_coding", "claude", "", constant.ChannelTypeVolcEngine, "kimi-coding-plan", true}, + {"claude_volc_doubao_coding", "claude", "", constant.ChannelTypeVolcEngine, "doubao-coding-plan", true}, } for _, provider := range providers { modes := []bool{true} @@ -83,13 +94,21 @@ func TestGlobalChatResponsesProtocolRetainsOriginalFunding(t *testing.T) { original := constant.ChannelSpecialBases[provider.baseAlias] synthetic := original synthetic.OpenAIBaseURL = upstream.URL + "/coding" + synthetic.ClaudeBaseURL = upstream.URL + "/coding" constant.ChannelSpecialBases[provider.baseAlias] = synthetic t.Cleanup(func() { constant.ChannelSpecialBases[provider.baseAlias] = original }) c.Set(string(constant.ContextKeyChannelBaseUrl), provider.baseAlias) // The adaptor still resolves the real alias, but its HTTP target is local. } info.IsStream = stream - info.Request.(*dto.GeneralOpenAIRequest).Stream = common.GetPointer(stream) + if provider.claudeClient { + info.RelayFormat = types.RelayFormatClaude + info.Request = &dto.ClaudeRequest{Model: "gpt-test", Messages: []dto.ClaudeMessage{{Role: "user", Content: "test"}}, MaxTokens: common.GetPointer(uint(64)), Stream: common.GetPointer(stream)} + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + info.RequestURLPath = "/v1/messages" + } else { + info.Request.(*dto.GeneralOpenAIRequest).Stream = common.GetPointer(stream) + } model_setting.GetGlobalSettings().ChatCompletionsToResponsesPolicy = model_setting.ChatCompletionsToResponsesPolicy{Enabled: enabled, AllChannels: true, ModelPatterns: []string{"^gpt-test$"}} if provider.channelType == constant.ChannelTypeAdvancedCustom { path := "/v1/responses" @@ -101,10 +120,20 @@ func TestGlobalChatResponsesProtocolRetainsOriginalFunding(t *testing.T) { c.Set(string(constant.ContextKeyChannelOtherSetting), dto.ChannelOtherSettings{AdvancedCustom: &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{IncomingPath: "/v1/chat/completions", UpstreamPath: path, Converter: provider.converter}}}}) } require.Nil(t, service.PreConsumeBilling(c, 100, info)) - apiErr := TextHelper(c, info) + var apiErr *types.MaxAPIError + if provider.claudeClient { + apiErr = ClaudeHelper(c, info) + } else { + apiErr = TextHelper(c, info) + } require.Nil(t, apiErr, "global policy must not send native replies to a Responses parser") require.Contains(t, recorder.Body.String(), "partial output") - require.Contains(t, recorder.Body.String(), "chat.completion") + if provider.claudeClient { + require.Equal(t, types.RelayFormat(types.RelayFormatClaude), info.RelayFormat) + require.Contains(t, recorder.Body.String(), `"type":"message`) + } else { + require.Contains(t, recorder.Body.String(), "chat.completion") + } require.EqualValues(t, 1, calls.Load()) select { case request := <-requests: @@ -126,6 +155,8 @@ func TestGlobalChatResponsesProtocolRetainsOriginalFunding(t *testing.T) { path := "/v1/responses" if provider.channelType == constant.ChannelTypeZhipu_v4 { path = "/api/v1/responses" + } else if provider.channelType == constant.ChannelTypeVolcEngine { + path = "/api/v3/responses" } require.Equal(t, path, request.path) require.Contains(t, request.body, "input") @@ -242,65 +273,71 @@ func TestGlobalResponsesBridgeRejectsNativeBeforeUpstreamAndRefunds(t *testing.T } } -func TestGlobalResponsesBridgeZhipuEndpointEligibility(t *testing.T) { +func TestGlobalResponsesBridgeEndpointEligibility(t *testing.T) { initRequestReplayTests(t) bases := []string{"", "https://open.bigmodel.cn", "https://synthetic.invalid"} for alias := range constant.ChannelSpecialBases { bases = append(bases, alias) } - for _, base := range bases { - t.Run(base, func(t *testing.T) { - c, info, _ := streamFundingContext(t, "chat_native", "zhipu-capability", base, constant.ChannelTypeZhipu_v4, streamFundingCases()[3]) - info.InitChannelMeta(c) - adaptor := GetAdaptor(info.ApiType) - adaptor.Init(info) - info.RelayMode = relayconstant.RelayModeResponses - _, err := adaptor.GetRequestURL(info) - info.RelayMode = relayconstant.RelayModeChatCompletions - model_setting.GetGlobalSettings().ChatCompletionsToResponsesPolicy = model_setting.ChatCompletionsToResponsesPolicy{Enabled: true, AllChannels: true, ModelPatterns: []string{"^gpt-test$"}} - require.Equal(t, err == nil, shouldUseGlobalResponsesBridge(c, info), "eligibility must match the actual adaptor URL capability") - }) + for _, channelType := range []int{constant.ChannelTypeZhipu_v4, constant.ChannelTypeVolcEngine} { + for _, base := range bases { + t.Run(fmt.Sprintf("channel_%d/%s", channelType, base), func(t *testing.T) { + c, info, _ := streamFundingContext(t, "chat_native", "endpoint-capability", base, channelType, streamFundingCases()[3]) + info.InitChannelMeta(c) + adaptor := GetAdaptor(info.ApiType) + adaptor.Init(info) + info.RelayMode = relayconstant.RelayModeResponses + requestURL, err := adaptor.GetRequestURL(info) + parsed, parseErr := url.ParseRequestURI(requestURL) + expected := err == nil && parseErr == nil && parsed.Scheme != "" && parsed.Host != "" + info.RelayMode = relayconstant.RelayModeChatCompletions + model_setting.GetGlobalSettings().ChatCompletionsToResponsesPolicy = model_setting.ChatCompletionsToResponsesPolicy{Enabled: true, AllChannels: true, ModelPatterns: []string{"^gpt-test$"}} + require.Equal(t, expected, shouldUseGlobalResponsesBridge(c, info), "eligibility must match the actual adaptor URL capability") + }) + } } } -func TestDirectResponsesBridgeRejectsZhipuCodingAliasesAndRefunds(t *testing.T) { +func TestDirectResponsesBridgeRejectsCodingAliasesAndRefunds(t *testing.T) { initRequestReplayTests(t) - for alias := range constant.ChannelSpecialBases { - for _, funding := range streamFundingCases() { - t.Run(alias+"/"+funding.name, func(t *testing.T) { - db := setupStreamFundingLedger(t, funding) - var calls atomic.Int32 - upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { calls.Add(1) })) - t.Cleanup(upstream.Close) - original := constant.ChannelSpecialBases[alias] - synthetic := original - synthetic.OpenAIBaseURL = upstream.URL - constant.ChannelSpecialBases[alias] = synthetic - t.Cleanup(func() { constant.ChannelSpecialBases[alias] = original }) - requestID := "rejected-zhipu-coding" - c, info, recorder := streamFundingContext(t, "chat_native", requestID, alias, constant.ChannelTypeZhipu_v4, funding) - info.InitChannelMeta(c) - info.RequestURLPath = "/v1/chat/completions" - adaptor := &responsesConversionProbe{Adaptor: GetAdaptor(info.ApiType)} - adaptor.Init(info) - require.Nil(t, service.PreConsumeBilling(c, 100, info)) - usage, apiErr := chatCompletionsViaResponses(c, info, adaptor, info.Request.(*dto.GeneralOpenAIRequest)) - require.Zero(t, adaptor.conversions, "unsupported endpoint must be rejected before converting the request") - require.Nil(t, usage) - require.NotNil(t, apiErr) - require.Equal(t, types.ErrorCodeConvertRequestFailed, apiErr.GetErrorCode()) - require.True(t, types.IsSkipRetryError(apiErr)) - require.Zero(t, calls.Load()) - require.Empty(t, recorder.Body.String()) - require.Equal(t, relayconstant.RelayModeChatCompletions, info.RelayMode) - require.Equal(t, "/v1/chat/completions", info.RequestURLPath) - service.HandleFailedBilling(c, info, apiErr) - service.HandleFailedBilling(c, info, apiErr) - assertStreamFundingBalances(t, db, funding, requestID, 0, false, 0, 0) - before := requestReplaySnapshot(t, db) - model.ProcessPendingBillingSettlementsOnce() - require.Equal(t, before, requestReplaySnapshot(t, db)) - }) + for _, channelType := range []int{constant.ChannelTypeZhipu_v4, constant.ChannelTypeVolcEngine} { + for alias := range constant.ChannelSpecialBases { + for _, funding := range streamFundingCases() { + t.Run(fmt.Sprintf("channel_%d/%s/%s", channelType, alias, funding.name), func(t *testing.T) { + db := setupStreamFundingLedger(t, funding) + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { calls.Add(1) })) + t.Cleanup(upstream.Close) + original := constant.ChannelSpecialBases[alias] + synthetic := original + synthetic.OpenAIBaseURL = upstream.URL + constant.ChannelSpecialBases[alias] = synthetic + t.Cleanup(func() { constant.ChannelSpecialBases[alias] = original }) + requestID := "rejected-zhipu-coding" + c, info, recorder := streamFundingContext(t, "chat_native", requestID, alias, channelType, funding) + info.InitChannelMeta(c) + info.RequestURLPath = "/v1/chat/completions" + adaptor := &responsesConversionProbe{Adaptor: GetAdaptor(info.ApiType)} + adaptor.Init(info) + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + usage, apiErr := chatCompletionsViaResponses(c, info, adaptor, info.Request.(*dto.GeneralOpenAIRequest)) + require.Zero(t, adaptor.conversions, "unsupported endpoint must be rejected before converting the request") + require.Nil(t, usage) + require.NotNil(t, apiErr) + require.Equal(t, types.ErrorCodeConvertRequestFailed, apiErr.GetErrorCode()) + require.True(t, types.IsSkipRetryError(apiErr)) + require.Zero(t, calls.Load()) + require.Empty(t, recorder.Body.String()) + require.Equal(t, relayconstant.RelayModeChatCompletions, info.RelayMode) + require.Equal(t, "/v1/chat/completions", info.RequestURLPath) + service.HandleFailedBilling(c, info, apiErr) + service.HandleFailedBilling(c, info, apiErr) + assertStreamFundingBalances(t, db, funding, requestID, 0, false, 0, 0) + before := requestReplaySnapshot(t, db) + model.ProcessPendingBillingSettlementsOnce() + require.Equal(t, before, requestReplaySnapshot(t, db)) + }) + } } } } diff --git a/relay/claude_handler.go b/relay/claude_handler.go index dfa1ea31..dfc33a29 100644 --- a/relay/claude_handler.go +++ b/relay/claude_handler.go @@ -92,6 +92,7 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (maxAPIError *typ if !model_setting.GetGlobalSettings().PassThroughRequestEnabled && !info.ChannelSetting.PassThroughBodyEnabled && + service.SupportsOpenAIResponsesUpstreamForBaseURL(info.ChannelType, info.ChannelBaseUrl) && service.ShouldChatCompletionsUseResponsesGlobal(info.ChannelId, info.ChannelType, info.OriginModelName) { openAIRequest, convErr := service.ClaudeToOpenAIRequest(*request, info) if convErr != nil { diff --git a/relay/file_search_stream_settlement_test.go b/relay/file_search_stream_settlement_test.go index d58b49a8..d1274aea 100644 --- a/relay/file_search_stream_settlement_test.go +++ b/relay/file_search_stream_settlement_test.go @@ -31,20 +31,29 @@ type fileSearchCancelWriter struct { recorder *httptest.ResponseRecorder cancel context.CancelFunc triggered atomic.Bool + callType string } func (w *fileSearchCancelWriter) Flush() { w.ResponseWriter.Flush() - if strings.Contains(w.recorder.Body.String(), `"type":"file_search_call"`) && w.triggered.CompareAndSwap(false, true) { + callType := w.callType + if callType == "" { + callType = dto.BuildInCallFileSearchCall + } + if strings.Contains(w.recorder.Body.String(), `"type":"`+callType+`"`) && w.triggered.CompareAndSwap(false, true) { w.cancel() } } func setFileSearchLedgerPriceForTest(t *testing.T, price float64) { + setSearchLedgerPriceForTest(t, dto.BuildInToolFileSearch, price) +} + +func setSearchLedgerPriceForTest(t *testing.T, toolName string, price float64) { t.Helper() prices := config.GlobalConfig.Get("tool_price_setting").(*operation_setting.ToolPriceSetting) oldPrices := prices.Prices - prices.Prices = map[string]float64{dto.BuildInToolFileSearch: price} + prices.Prices = map[string]float64{toolName: price} operation_setting.RebuildToolPriceIndex() t.Cleanup(func() { prices.Prices = oldPrices @@ -53,12 +62,21 @@ func setFileSearchLedgerPriceForTest(t *testing.T, price float64) { } func TestFileSearchInterruptedStreamRetainsOriginalFunding(t *testing.T) { + testSearchInterruptedStreamRetainsOriginalFunding(t, dto.BuildInToolFileSearch, dto.BuildInCallFileSearchCall) +} + +func TestWebSearchInterruptedStreamRetainsOriginalFunding(t *testing.T) { + testSearchInterruptedStreamRetainsOriginalFunding(t, dto.BuildInToolWebSearch, dto.BuildInCallWebSearchCall) +} + +func testSearchInterruptedStreamRetainsOriginalFunding(t *testing.T, toolName, callType string) { + t.Helper() gin.SetMode(gin.TestMode) service.InitHttpClient() service.InitTokenEncoders() oldEmpty, oldRetries := common.EmptyCompletionRetryEnabled, common.RetryTimes common.RetryTimes = 2 - setFileSearchLedgerPriceForTest(t, 0.02) + setSearchLedgerPriceForTest(t, toolName, 0.02) t.Cleanup(func() { common.EmptyCompletionRetryEnabled, common.RetryTimes = oldEmpty, oldRetries }) @@ -71,7 +89,7 @@ func TestFileSearchInterruptedStreamRetainsOriginalFunding(t *testing.T) { common.EmptyCompletionRetryEnabled = retry db := setupStreamFundingLedger(t, funding) _, frames := metadataFundingFrames(t, "responses_native", usageKind) - item := map[string]any{"type": "file_search_call", "id": "fs_completed", "queries": []string{"synthetic query"}, + item := map[string]any{"type": callType, "id": "search_completed", "queries": []string{"synthetic query"}, "results": []any{map[string]any{"file_id": "file_synthetic", "text": "synthetic-search-result"}}} if status == "completed" { item["status"] = "completed" @@ -106,11 +124,30 @@ func TestFileSearchInterruptedStreamRetainsOriginalFunding(t *testing.T) { fingerprint := sha256.Sum256([]byte(t.Name())) requestID := fmt.Sprintf("file-search-%x", fingerprint[:16]) c, info, recorder := streamFundingContext(t, "responses_native", requestID, upstream.URL, constant.ChannelTypeOpenAI, funding) - info.Request.(*dto.OpenAIResponsesRequest).Tools = []byte(`[{"type":"file_search","vector_store_ids":["vs_synthetic"]}]`) + if callType == dto.BuildInCallFileSearchCall { + info.Request.(*dto.OpenAIResponsesRequest).Tools = []byte(`[{"type":"file_search","vector_store_ids":["vs_synthetic"]}]`) + } else { + info.Request.(*dto.OpenAIResponsesRequest).Tools = []byte(`[{"type":"web_search"}]`) + } info.ResponsesUsageInfo = &relaycommon.ResponsesUsageInfo{BuiltInTools: map[string]*relaycommon.BuildInToolInfo{ - dto.BuildInToolFileSearch: {ToolName: dto.BuildInToolFileSearch}, + toolName: {ToolName: toolName}, }} + if callType == dto.BuildInCallWebSearchCall { + info.ToolUsage = relaycommon.NewToolUsageLedger(info.OriginModelName) + } require.Nil(t, service.PreConsumeBilling(c, 100, info)) + if callType == dto.BuildInCallWebSearchCall { + prices := config.GlobalConfig.Get("tool_price_setting").(*operation_setting.ToolPriceSetting) + original := prices.Prices + changed := make(map[string]float64, len(original)) + for name, price := range original { + changed[name] = price + } + changed[toolName] = 99 + prices.Prices = changed + operation_setting.RebuildToolPriceIndex() + t.Cleanup(func() { prices.Prices = original; operation_setting.RebuildToolPriceIndex() }) + } require.Equal(t, funding.source, info.BillingSource) require.NoError(t, db.Create(&model.UserSubscription{Id: 956, UserId: 951, PlanId: 954, Status: "active", AmountTotal: 1000, StartTime: time.Now().Unix() - 1800, EndTime: time.Now().Unix() + 1800}).Error) @@ -123,7 +160,7 @@ func TestFileSearchInterruptedStreamRetainsOriginalFunding(t *testing.T) { ctx, cancel := context.WithCancel(c.Request.Context()) t.Cleanup(cancel) c.Request = c.Request.WithContext(ctx) - cancelWriter = &fileSearchCancelWriter{ResponseWriter: c.Writer, recorder: recorder, cancel: cancel} + cancelWriter = &fileSearchCancelWriter{ResponseWriter: c.Writer, recorder: recorder, cancel: cancel, callType: callType} c.Writer = cancelWriter } apiErr := ResponsesHelper(c, info) @@ -145,7 +182,7 @@ func TestFileSearchInterruptedStreamRetainsOriginalFunding(t *testing.T) { require.True(t, types.IsSkipRetryError(apiErr)) require.True(t, info.HasRecordedChannelFirstResult()) require.False(t, info.FirstResultTime.IsZero()) - require.Equal(t, 1, info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolFileSearch].CallCount) + require.Equal(t, 1, info.ResponsesUsageInfo.BuiltInTools[toolName].CallCount) require.Contains(t, recorder.Body.String(), "synthetic-search-result") require.NotContains(t, recorder.Body.String(), "response.completed") require.EqualValues(t, 1, calls.Load()) @@ -164,8 +201,16 @@ func TestFileSearchInterruptedStreamRetainsOriginalFunding(t *testing.T) { require.NoError(t, db.Where("type = ?", model.LogTypeConsume).First(&log).Error) var other map[string]any require.NoError(t, common.UnmarshalJsonStr(log.Other, &other)) - require.EqualValues(t, 1, other["file_search_call_count"]) - require.EqualValues(t, 0.02, other["file_search_price"]) + require.EqualValues(t, 1, other[toolName+"_call_count"]) + require.EqualValues(t, 0.02, other[toolName+"_price"]) + if callType == dto.BuildInCallWebSearchCall { + toolUsage, ok := other["web_search_tool_usage"].(map[string]any) + require.True(t, ok, "per-tool search usage must be a JSON object") + details, ok := toolUsage[toolName].(map[string]any) + require.True(t, ok, "the completed search kind must be present in the audit record") + require.EqualValues(t, 1, details["call_count"]) + require.EqualValues(t, 0.02, details["price"]) + } var outbox int64 require.NoError(t, db.Model(&model.CacheInvalidationTask{}).Count(&outbox).Error) for range 2 { @@ -197,52 +242,71 @@ func TestFileSearchInterruptedStreamRetainsOriginalFunding(t *testing.T) { } func TestFileSearchLifecycleItemsRefundOriginalFunding(t *testing.T) { + testSearchLifecycleItemsRefundOriginalFunding(t, dto.BuildInToolFileSearch, dto.BuildInCallFileSearchCall) +} + +func TestWebSearchLifecycleItemsRefundOriginalFunding(t *testing.T) { + testSearchLifecycleItemsRefundOriginalFunding(t, dto.BuildInToolWebSearch, dto.BuildInCallWebSearchCall) +} + +func testSearchLifecycleItemsRefundOriginalFunding(t *testing.T, toolName, callType string) { + t.Helper() gin.SetMode(gin.TestMode) service.InitHttpClient() service.InitTokenEncoders() oldEmpty := common.EmptyCompletionRetryEnabled common.EmptyCompletionRetryEnabled = false t.Cleanup(func() { common.EmptyCompletionRetryEnabled = oldEmpty }) - setFileSearchLedgerPriceForTest(t, 0.02) - for _, funding := range streamFundingCases() { - for _, tc := range []struct{ name, event, status string }{ - {"added_running", dto.ResponsesOutputTypeItemAdded, "in_progress"}, - {"added_completed", dto.ResponsesOutputTypeItemAdded, "completed"}, - {"done_running", dto.ResponsesOutputTypeItemDone, "in_progress"}, - {"done_failed", dto.ResponsesOutputTypeItemDone, "failed"}, - {"done_cancelled", dto.ResponsesOutputTypeItemDone, "cancelled"}, - {"done_unknown", dto.ResponsesOutputTypeItemDone, "future_status"}, - {"nil_item", dto.ResponsesOutputTypeItemDone, ""}, - } { - t.Run(funding.name+"/"+tc.name, func(t *testing.T) { - db := setupStreamFundingLedger(t, funding) - _, frames := metadataFundingFrames(t, "responses_native", "known") - event := dto.ResponsesStreamResponse{Type: tc.event, Item: &dto.ResponsesOutput{Type: dto.BuildInCallFileSearchCall, ID: "fs_metadata", Status: tc.status}} - if tc.name == "nil_item" { - event.Item = nil + setSearchLedgerPriceForTest(t, toolName, 0.02) + usageKinds := []string{"known"} + if callType == dto.BuildInCallWebSearchCall { + usageKinds = append(usageKinds, "zero", "missing") + } + for _, usageKind := range usageKinds { + for _, funding := range streamFundingCases() { + for _, tc := range []struct{ name, event, status string }{ + {"added_running", dto.ResponsesOutputTypeItemAdded, "in_progress"}, + {"added_completed", dto.ResponsesOutputTypeItemAdded, "completed"}, + {"done_running", dto.ResponsesOutputTypeItemDone, "in_progress"}, + {"done_failed", dto.ResponsesOutputTypeItemDone, "failed"}, + {"done_cancelled", dto.ResponsesOutputTypeItemDone, "cancelled"}, + {"done_unknown", dto.ResponsesOutputTypeItemDone, "future_status"}, + {"nil_item", dto.ResponsesOutputTypeItemDone, ""}, + } { + name := funding.name + "/" + tc.name + if callType == dto.BuildInCallWebSearchCall { + name = usageKind + "/" + name } - data, err := common.Marshal(event) - require.NoError(t, err) - frames = append(frames, string(data)) - upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { - w.Header().Set("Content-Type", "text/event-stream") - _, _ = fmt.Fprint(w, "data: "+strings.Join(frames, "\n\ndata: ")+"\n\n") - })) - t.Cleanup(upstream.Close) - requestID := "file-search-lifecycle-" + funding.name + "-" + tc.name - c, info, _ := streamFundingContext(t, "responses_native", requestID, upstream.URL, constant.ChannelTypeOpenAI, funding) - info.ResponsesUsageInfo = &relaycommon.ResponsesUsageInfo{BuiltInTools: map[string]*relaycommon.BuildInToolInfo{ - dto.BuildInToolFileSearch: {ToolName: dto.BuildInToolFileSearch}, - }} - require.Nil(t, service.PreConsumeBilling(c, 100, info)) - apiErr := ResponsesHelper(c, info) - require.NotNil(t, apiErr) - service.HandleFailedBilling(c, info, apiErr) - service.HandleFailedBilling(c, info, apiErr) - require.False(t, info.HasRecordedChannelFirstResult()) - require.Zero(t, info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolFileSearch].CallCount) - assertStreamFundingBalances(t, db, funding, requestID, 0, false, 0, 0) - }) + t.Run(name, func(t *testing.T) { + db := setupStreamFundingLedger(t, funding) + _, frames := metadataFundingFrames(t, "responses_native", usageKind) + event := dto.ResponsesStreamResponse{Type: tc.event, Item: &dto.ResponsesOutput{Type: callType, ID: "search_metadata", Status: tc.status}} + if tc.name == "nil_item" { + event.Item = nil + } + data, err := common.Marshal(event) + require.NoError(t, err) + frames = append(frames, string(data)) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, "data: "+strings.Join(frames, "\n\ndata: ")+"\n\n") + })) + t.Cleanup(upstream.Close) + requestID := "file-search-lifecycle-" + funding.name + "-" + tc.name + c, info, _ := streamFundingContext(t, "responses_native", requestID, upstream.URL, constant.ChannelTypeOpenAI, funding) + info.ResponsesUsageInfo = &relaycommon.ResponsesUsageInfo{BuiltInTools: map[string]*relaycommon.BuildInToolInfo{ + toolName: {ToolName: toolName}, + }} + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + apiErr := ResponsesHelper(c, info) + require.NotNil(t, apiErr) + service.HandleFailedBilling(c, info, apiErr) + service.HandleFailedBilling(c, info, apiErr) + require.False(t, info.HasRecordedChannelFirstResult()) + require.Zero(t, info.ResponsesUsageInfo.BuiltInTools[toolName].CallCount) + assertStreamFundingBalances(t, db, funding, requestID, 0, false, 0, 0) + }) + } } } } diff --git a/relay/helper/common.go b/relay/helper/common.go index f968deaf..5efb1f3b 100644 --- a/relay/helper/common.go +++ b/relay/helper/common.go @@ -163,7 +163,7 @@ func ResponsesStreamEventHasOutput(event dto.ResponsesStreamResponse) bool { return event.Delta != "" case dto.ResponsesOutputTypeItemAdded, dto.ResponsesOutputTypeItemDone: if event.Item != nil { - if event.Item.Type == dto.BuildInCallFileSearchCall { + if event.Item.Type == dto.BuildInCallFileSearchCall || event.Item.Type == dto.BuildInCallWebSearchCall { // A search lifecycle/failed item is metadata. Only a finished // invocation carries a delivered result and its existing fee. status := strings.ToLower(strings.TrimSpace(event.Item.Status)) diff --git a/relay/metadata_funding_settlement_test.go b/relay/metadata_funding_settlement_test.go index 3ec8b8c8..5f5f903d 100644 --- a/relay/metadata_funding_settlement_test.go +++ b/relay/metadata_funding_settlement_test.go @@ -65,7 +65,12 @@ func metadataFundingFrames(t *testing.T, route, usage string) (int, []string) { if route == "responses_to_chat_custom" { channel = constant.ChannelTypeAdvancedCustom } - return channel, []string{encode(map[string]any{"type": "response.created", "response": response})} + frames := []string{encode(map[string]any{"type": "response.created", "response": response})} + if route == "responses_web_search" { + frames = append(frames, encode(map[string]any{"type": "response.output_item.added", "output_index": 0, + "item": map[string]any{"type": "web_search_call", "id": "ws_metadata", "status": "in_progress"}})) + } + return channel, frames } } @@ -91,7 +96,7 @@ func TestMetadataOnlyFailureRetainsOriginalFunding(t *testing.T) { oldEmptyRetry := common.EmptyCompletionRetryEnabled common.EmptyCompletionRetryEnabled = false t.Cleanup(func() { common.EmptyCompletionRetryEnabled = oldEmptyRetry }) - for _, route := range []string{"responses_native", "chat_native", "chat_to_responses", "responses_to_chat_custom", "responses_to_chat_global", "claude_to_responses", "gemini_to_responses"} { + for _, route := range []string{"responses_native", "responses_web_search", "chat_native", "chat_to_responses", "responses_to_chat_custom", "responses_to_chat_global", "claude_to_responses", "gemini_to_responses"} { for _, funding := range streamFundingCases() { for _, usage := range []string{"known", "zero", "missing"} { for _, ending := range []string{"eof", "transport", "malformed", "cancel", "timeout", "eof_terminal_write_failure"} { diff --git a/service/openaicompat/policy.go b/service/openaicompat/policy.go index 77951977..2c69bded 100644 --- a/service/openaicompat/policy.go +++ b/service/openaicompat/policy.go @@ -29,7 +29,7 @@ func SupportsOpenAIResponsesUpstreamForBaseURL(channelType int, baseURL string) return false } apiType, _ := common.ChannelType2APIType(channelType) - if apiType == constant.APITypeZhipuV4 { + if apiType == constant.APITypeZhipuV4 || apiType == constant.APITypeVolcEngine { _, codingPlan := constant.ChannelSpecialBases[baseURL] return !codingPlan } diff --git a/service/text_quota.go b/service/text_quota.go index 6fd4214b..db698e44 100644 --- a/service/text_quota.go +++ b/service/text_quota.go @@ -50,6 +50,7 @@ type textQuotaSummary struct { UsageSemantic string WebSearchPrice float64 WebSearchCallCount int + WebSearchToolUsage map[string]webSearchToolUsage ClaudeWebSearchPrice float64 ClaudeWebSearchCallCount int FileSearchPrice float64 @@ -60,6 +61,11 @@ type textQuotaSummary struct { ToolCallSurchargeQuota decimal.Decimal } +type webSearchToolUsage struct { + CallCount int `json:"call_count"` + Price float64 `json:"price"` +} + func (s *textQuotaSummary) hasBillableUsage(relayInfo *relaycommon.RelayInfo) bool { if s.TotalTokens > 0 || !s.ToolCallSurchargeQuota.IsZero() { return true @@ -128,14 +134,34 @@ func calculateTextToolCallSurcharge(ctx *gin.Context, relayInfo *relaycommon.Rel summary.CustomToolUsage, surcharge = calculateCustomToolCallSurcharge(relayInfo, summary.GroupRatio) if relayInfo.ResponsesUsageInfo != nil { - if webSearchTool, exists := relayInfo.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolWebSearchPreview]; exists && webSearchTool.CallCount > 0 { - summary.WebSearchCallCount = webSearchTool.CallCount - summary.WebSearchPrice = relayInfo.FrozenToolPrice("web_search_preview") - surcharge = surcharge.Add(decimal.NewFromFloat(summary.WebSearchPrice). - Mul(decimal.NewFromInt(int64(webSearchTool.CallCount))). + var webSearchCost decimal.Decimal + webSearchCalls := 0 + var webSearchUsage map[string]webSearchToolUsage + for _, toolName := range []string{dto.BuildInToolWebSearchPreview, dto.BuildInToolWebSearch} { + webSearchTool := relayInfo.ResponsesUsageInfo.BuiltInTools[toolName] + if webSearchTool == nil || webSearchTool.CallCount <= 0 { + continue + } + price := relayInfo.FrozenToolPrice(toolName) + if webSearchUsage == nil { + webSearchUsage = make(map[string]webSearchToolUsage, 2) + } + webSearchUsage[toolName] = webSearchToolUsage{CallCount: webSearchTool.CallCount, Price: price} + webSearchCalls += webSearchTool.CallCount + webSearchCost = webSearchCost.Add(decimal.NewFromFloat(price).Mul(decimal.NewFromInt(int64(webSearchTool.CallCount)))) + } + summary.WebSearchCallCount = webSearchCalls + summary.WebSearchToolUsage = webSearchUsage + if webSearchCalls > 0 { + // Keep the aggregate log fields compatible; charge the exact sum + // from frozen per-tool prices, rather than this display average. + summary.WebSearchPrice = webSearchCost.Div(decimal.NewFromInt(int64(summary.WebSearchCallCount))).InexactFloat64() + surcharge = surcharge.Add(webSearchCost. Div(decimal.NewFromInt(1000)). Mul(dGroupRatio). Mul(dQuotaPerUnit)) + } else { + summary.WebSearchPrice = 0 } } else if strings.HasSuffix(summary.ModelName, "search-preview") { summary.WebSearchCallCount = 1 @@ -547,6 +573,7 @@ func postTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us other["web_search"] = true other["web_search_call_count"] = summary.WebSearchCallCount other["web_search_price"] = summary.WebSearchPrice + other["web_search_tool_usage"] = summary.WebSearchToolUsage } else if summary.ClaudeWebSearchCallCount > 0 { other["web_search"] = true other["web_search_call_count"] = summary.ClaudeWebSearchCallCount diff --git a/service/tool_usage_billing_test.go b/service/tool_usage_billing_test.go index 831d48b1..32550399 100644 --- a/service/tool_usage_billing_test.go +++ b/service/tool_usage_billing_test.go @@ -1,6 +1,7 @@ package service import ( + "fmt" "net/http/httptest" "testing" "time" @@ -170,6 +171,58 @@ func TestCalculateTextToolCallSurchargeSkipsPerRequestModels(t *testing.T) { require.Empty(t, summary.CustomToolUsage.Items) } +func TestWebSearchSurchargeRetainsFrozenPricesAndDistinctCalls(t *testing.T) { + for _, tc := range []struct { + name string + previewCalls, webCalls int + webPrice, groupRatio float64 + }{ + {"preview", 2, 0, 0.02, 1}, + {"web", 0, 3, 0.02, 1}, + {"both", 2, 3, 0.02, 1}, + {"free_group", 2, 3, 0.02, 0}, + {"free_web", 0, 3, 0, 1}, + } { + t.Run(tc.name, func(t *testing.T) { + setServiceToolPricesForTest(t, map[string]float64{dto.BuildInToolWebSearchPreview: 0.03, dto.BuildInToolWebSearch: tc.webPrice}) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + info := &relaycommon.RelayInfo{OriginModelName: "gpt-test", ToolUsage: relaycommon.NewToolUsageLedger("gpt-test"), + ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{BuiltInTools: map[string]*relaycommon.BuildInToolInfo{ + dto.BuildInToolWebSearchPreview: {ToolName: dto.BuildInToolWebSearchPreview}, + dto.BuildInToolWebSearch: {ToolName: dto.BuildInToolWebSearch}, + }}} + info.ToolUsage.BeginAttempt(0) + for toolName, count := range map[string]int{dto.BuildInToolWebSearchPreview: tc.previewCalls, dto.BuildInToolWebSearch: tc.webCalls} { + for i := 0; i < count; i++ { + identity := relaycommon.ToolCallIdentity{Scope: "openai-responses", CallID: fmt.Sprintf("%s-%d", toolName, i)} + billable := toolName != dto.BuildInToolWebSearch || tc.webPrice > 0 + require.Equal(t, billable, info.ObserveBuiltInToolCall(toolName, identity)) + require.False(t, info.ObserveBuiltInToolCall(toolName, identity)) + } + } + setServiceToolPricesForTest(t, map[string]float64{dto.BuildInToolWebSearchPreview: 99, dto.BuildInToolWebSearch: 99}) + summary := textQuotaSummary{ModelName: "gpt-test", GroupRatio: tc.groupRatio} + cost := decimal.NewFromFloat(0.03).Mul(decimal.NewFromInt(int64(tc.previewCalls))).Add(decimal.NewFromFloat(tc.webPrice).Mul(decimal.NewFromInt(int64(tc.webCalls)))) + expected := cost.Div(decimal.NewFromInt(1000)).Mul(decimal.NewFromFloat(common.QuotaPerUnit)).Mul(decimal.NewFromFloat(tc.groupRatio)) + actual := calculateTextToolCallSurcharge(ctx, info, &summary) + require.True(t, expected.Equal(actual), "expected %s, actual %s; frozen preview=%g, web=%g", expected, actual, info.FrozenToolPrice(dto.BuildInToolWebSearchPreview), info.FrozenToolPrice(dto.BuildInToolWebSearch)) + calls := tc.previewCalls + tc.webCalls + if tc.webPrice == 0 { + calls -= tc.webCalls // The existing ledger excludes free tool calls. + } + require.Equal(t, calls, summary.WebSearchCallCount) + require.True(t, expected.Equal(calculateTextToolCallSurcharge(ctx, info, &summary))) + require.Equal(t, calls, summary.WebSearchCallCount, "recomputing a summary must not duplicate observed calls") + if calls == 0 { + require.Zero(t, summary.WebSearchPrice) + return + } + average, _ := cost.Div(decimal.NewFromInt(int64(calls))).Float64() + require.Equal(t, average, summary.WebSearchPrice) + }) + } +} + func TestCalculateTextQuotaSummaryBillsProjectedResponsesBuiltInToolOnce(t *testing.T) { setServiceToolPricesForTest(t, map[string]float64{ dto.BuildInToolWebSearchPreview: 5, From 9bbec9e7d5eaba83deeb3f8be935ee46fb992e5f Mon Sep 17 00:00:00 2001 From: CSCITech Date: Tue, 13 Oct 2026 23:02:03 +0800 Subject: [PATCH 22/24] feat: support separate SQLite log files --- .env.example | 5 + docs/installation/sqlite-log-database.md | 26 +++ model/main.go | 88 ++++++---- model/sqlite_log_config_test.go | 207 +++++++++++++++++++++++ service/billing_process_recovery_test.go | 59 ++++++- 5 files changed, 348 insertions(+), 37 deletions(-) create mode 100644 docs/installation/sqlite-log-database.md create mode 100644 model/sqlite_log_config_test.go diff --git a/.env.example b/.env.example index fd94aea1..9c45d6e3 100644 --- a/.env.example +++ b/.env.example @@ -27,6 +27,11 @@ # LOG_SQL_DSN=user:password@tcp(127.0.0.1:3306)/logdb?parseTime=true # SQLite数据库路径 # SQLITE_PATH=/path/to/sqlite.db +# 独立SQLite日志文件（可选；LOG_SQL_DSN留空或为local，默认日志与主库共用） +# Optional separate SQLite log file; use with an empty or local LOG_SQL_DSN. +# LOG_SQLITE_PATH=/path/to/logs.db +# 切换既有日志库前先恢复待处理结算并备份日志及billing_log_receipts；不自动迁移旧数据。 +# Reconcile pending effects and back up logs with billing_log_receipts before switching; no automatic data copy. # 数据库最大空闲连接数 # SQL_MAX_IDLE_CONNS=100 # 数据库最大打开连接数 diff --git a/docs/installation/sqlite-log-database.md b/docs/installation/sqlite-log-database.md new file mode 100644 index 00000000..f757ba7a --- /dev/null +++ b/docs/installation/sqlite-log-database.md @@ -0,0 +1,26 @@ +# 独立 SQLite 日志文件 / Separate SQLite log database + +`LOG_SQLITE_PATH` 为可选的 SQLite 日志文件路径。只设置 `SQLITE_PATH` 时，日志与业务数据继续使用同一个数据库。配置 `LOG_SQLITE_PATH` 后，日志及结算日志 receipt 使用独立 SQLite 文件，业务资金和结算操作仍保存在主库。 + +`LOG_SQLITE_PATH` optionally selects a separate SQLite file for logs and billing log receipts. Without it, the existing shared-database behavior is preserved. Funding and durable settlement operations remain in the primary database. + +```dotenv +SQLITE_PATH=/data/ledger.sqlite +LOG_SQLITE_PATH=/data/logs.sqlite +``` + +- `LOG_SQL_DSN` 留空或设为 `local` 时，可使用 `LOG_SQLITE_PATH`。旧配置 `LOG_SQL_DSN=local` 且未配置新路径时，继续使用 `SQLITE_PATH`。 +- `LOG_SQL_DSN` 为 MySQL/PostgreSQL 连接串时，使用该连接串并保持现有行为；同时设置 `LOG_SQLITE_PATH` 会返回配置错误。二选一配置日志目标。 +- 提前创建文件的父目录，并为运行进程提供读写权限。显式日志库打开或迁移失败会阻止启动，修复路径、权限或数据库后再启动。 +- 路径支持 SQLite 驱动的文件连接参数，建议使用持久文件；两个 SQLite 文件应放在本地可靠存储上。本配置不代替多主机或外部数据库验收。 + +Leave `LOG_SQL_DSN` empty or set it to `local` when using `LOG_SQLITE_PATH`. Existing `LOG_SQL_DSN=local` deployments without the new setting keep using `SQLITE_PATH`. A remote MySQL/PostgreSQL log DSN and `LOG_SQLITE_PATH` together are rejected. Create the parent directory and grant the service access before startup. Opening or migrating the configured log database must succeed before startup can complete. + +## 切换既有部署 / Changing an existing deployment + +1. 暂停新请求及相关后台写入，在原配置下完成待恢复的资金与日志副作用，核对结算状态。 +2. 备份业务库、日志库和配置。日志与 `billing_log_receipts` 属于同一恢复历史，应一起保留；单独复制日志或创建空日志库不能证明既有记录已经迁移。 +3. 使用经过验证的离线迁移过程准备目标文件，再更新配置。此设置只选择连接目标，不自动搬迁、删除或合并旧数据；新文件初始为空时，界面读取的是新文件中的日志。 +4. 启动后核对历史日志、receipt、待恢复结算、余额与计数，并保留原文件用于审计和恢复。切回旧配置前，应先核对期间新增的数据与待恢复操作。 + +Quiesce requests and background writers, finish pending settlement effects with the original database, and back up all databases and configuration. Preserve logs and `billing_log_receipts` together. Prepare the target file with a verified offline migration before changing configuration. This setting selects a database; it does not copy, delete, or merge existing data. Verify history, receipts, settlement state, balances, and counters after startup, and reconcile new writes before switching back. diff --git a/model/main.go b/model/main.go index 7e1f07e3..d06027ad 100644 --- a/model/main.go +++ b/model/main.go @@ -52,14 +52,30 @@ var userOAuthIdentityMigrations = []userOAuthIdentityMigration{ {column: "linux_do_id", indexName: "ux_users_linux_do_id", mysqlGeneratedColumn: "linux_do_id_unique"}, } +func hasSeparateLogDatabase() bool { + return os.Getenv("LOG_SQL_DSN") != "" || strings.TrimSpace(os.Getenv("LOG_SQLITE_PATH")) != "" +} + +func sqliteDatabasePath(isLog bool) string { + if isLog { + if path := strings.TrimSpace(os.Getenv("LOG_SQLITE_PATH")); path != "" { + return path + } + } + return common.SQLitePath +} + func logRetryMarkerBackfillCompletionKey() string { logDSN := strings.TrimSpace(os.Getenv("LOG_SQL_DSN")) - if logDSN == "" { + if !hasSeparateLogDatabase() { return logRetryMarkerBackfillOptionKey } + if logDSN == "" { + logDSN = "local" + } if strings.HasPrefix(logDSN, "local") { - logDSN = logDSN + "\n" + common.SQLitePath + logDSN = logDSN + "\n" + sqliteDatabasePath(true) } fingerprint := common.Sha1([]byte(common.LogSqlType + "\n" + logDSN)) @@ -79,7 +95,7 @@ func initCol() { commonTrueVal = "1" commonFalseVal = "0" } - if os.Getenv("LOG_SQL_DSN") != "" { + if hasSeparateLogDatabase() { switch common.LogSqlType { case common.DatabaseTypePostgreSQL: logGroupCol = `"group"` @@ -92,7 +108,7 @@ func initCol() { logKeyCol = "`key`" } } else { - // LOG_SQL_DSN 为空时，日志数据库与主数据库相同 + // No separate log database is configured. if common.UsingPostgreSQL { logGroupCol = `"group"` logKeyCol = `"key"` @@ -171,6 +187,9 @@ func chooseDB(envName string, isLog bool) (*gorm.DB, error) { initCol() }() dsn := os.Getenv(envName) + if isLog && strings.TrimSpace(os.Getenv("LOG_SQLITE_PATH")) != "" && dsn != "" && !strings.HasPrefix(dsn, "local") { + return nil, errors.New("LOG_SQLITE_PATH requires an empty or local LOG_SQL_DSN") + } if dsn != "" { if strings.HasPrefix(dsn, "postgres://") || strings.HasPrefix(dsn, "postgresql://") { // Use PostgreSQL @@ -194,7 +213,7 @@ func chooseDB(envName string, isLog bool) (*gorm.DB, error) { } else { common.LogSqlType = common.DatabaseTypeSQLite } - return gorm.Open(sqlite.Open(common.SQLitePath), &gorm.Config{ + return gorm.Open(sqlite.Open(sqliteDatabasePath(isLog)), &gorm.Config{ PrepareStmt: true, // precompile SQL }) } @@ -224,7 +243,7 @@ func chooseDB(envName string, isLog bool) (*gorm.DB, error) { } else { common.LogSqlType = common.DatabaseTypeSQLite } - return gorm.Open(sqlite.Open(common.SQLitePath), &gorm.Config{ + return gorm.Open(sqlite.Open(sqliteDatabasePath(isLog)), &gorm.Config{ PrepareStmt: true, // precompile SQL }) } @@ -266,40 +285,47 @@ func InitDB() (err error) { } func InitLogDB() (err error) { - if os.Getenv("LOG_SQL_DSN") == "" { + if !hasSeparateLogDatabase() { LOG_DB = DB return } db, err := chooseDB("LOG_SQL_DSN", true) - if err == nil { - if common.DebugEnabled { - db = db.Debug() - } - LOG_DB = db - // If log DB is MySQL, also ensure Chinese-capable charset - if common.LogSqlType == common.DatabaseTypeMySQL { - if err := checkMySQLChineseSupport(LOG_DB); err != nil { - panic(err) - } + if err != nil { + if db != nil { + _ = closeDB(db) } - sqlDB, err := LOG_DB.DB() + return fmt.Errorf("failed to initialize log database: %w", err) + } + previousLogDB := LOG_DB + defer func() { if err != nil { - return err + _ = closeDB(db) + LOG_DB = previousLogDB } - sqlDB.SetMaxIdleConns(common.GetEnvOrDefault("SQL_MAX_IDLE_CONNS", 100)) - sqlDB.SetMaxOpenConns(common.GetEnvOrDefault("SQL_MAX_OPEN_CONNS", 1000)) - sqlDB.SetConnMaxLifetime(time.Second * time.Duration(common.GetEnvOrDefault("SQL_MAX_LIFETIME", 60))) - - if !common.IsMasterNode { - return nil + }() + if common.DebugEnabled { + db = db.Debug() + } + LOG_DB = db + // If log DB is MySQL, also ensure Chinese-capable charset + if common.LogSqlType == common.DatabaseTypeMySQL { + if err := checkMySQLChineseSupport(LOG_DB); err != nil { + panic(err) } - common.SysLog("database migration started") - err = migrateLOGDB() + } + sqlDB, err := LOG_DB.DB() + if err != nil { return err - } else { - common.FatalLog(err) } - return err + sqlDB.SetMaxIdleConns(common.GetEnvOrDefault("SQL_MAX_IDLE_CONNS", 100)) + sqlDB.SetMaxOpenConns(common.GetEnvOrDefault("SQL_MAX_OPEN_CONNS", 1000)) + sqlDB.SetConnMaxLifetime(time.Second * time.Duration(common.GetEnvOrDefault("SQL_MAX_LIFETIME", 60))) + + if !common.IsMasterNode { + return nil + } + common.SysLog("database migration started") + return migrateLOGDB() } func migrateDB() error { @@ -403,7 +429,7 @@ func migrateDB() error { if err := migrateSubscriptionPlanPolicyDefaults(planTableExisted, planWalletColumnExisted, planDowngradeColumnExisted); err != nil { return err } - if os.Getenv("LOG_SQL_DSN") == "" { + if !hasSeparateLogDatabase() { LOG_DB = DB scheduleLogRetryMarkerBackfill() } diff --git a/model/sqlite_log_config_test.go b/model/sqlite_log_config_test.go new file mode 100644 index 00000000..d76e1330 --- /dev/null +++ b/model/sqlite_log_config_test.go @@ -0,0 +1,207 @@ +package model + +import ( + "os" + "path/filepath" + "testing" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupSQLiteLogConfigTest(t *testing.T) string { + t.Helper() + previousDB, previousLogDB := DB, LOG_DB + previousPath, previousLogType := common.SQLitePath, common.LogSqlType + previousSQLite, previousMySQL, previousPostgres := common.UsingSQLite, common.UsingMySQL, common.UsingPostgreSQL + previousMaster, previousRunner := common.IsMasterNode, logRetryMarkerBackfillAsyncRunner + t.Cleanup(func() { + DB, LOG_DB = previousDB, previousLogDB + common.SQLitePath, common.LogSqlType = previousPath, previousLogType + common.UsingSQLite, common.UsingMySQL, common.UsingPostgreSQL = previousSQLite, previousMySQL, previousPostgres + common.IsMasterNode, logRetryMarkerBackfillAsyncRunner = previousMaster, previousRunner + initCol() + }) + t.Setenv("SQL_DSN", "local") + t.Setenv("LOG_SQL_DSN", "") + t.Setenv("LOG_SQLITE_PATH", "") + t.Setenv("QUOTA_DATA_AGGREGATE_MIGRATION_ENABLED", "false") + dir := t.TempDir() + common.SQLitePath = filepath.Join(dir, "ledger.sqlite") + common.UsingSQLite, common.UsingMySQL, common.UsingPostgreSQL = true, false, false + common.LogSqlType = common.DatabaseTypeSQLite + common.IsMasterNode = true + logRetryMarkerBackfillAsyncRunner = func(fn func()) { fn() } + return dir +} + +func TestSQLiteLogConfigInitialization(t *testing.T) { + for _, mode := range []string{"default", "legacy_local", "path_only", "local_and_path"} { + t.Run(mode, func(t *testing.T) { + dir := setupSQLiteLogConfigTest(t) + if mode == "legacy_local" || mode == "local_and_path" { + t.Setenv("LOG_SQL_DSN", "local") + } + logPath := filepath.Join(dir, "logs 中文 space.sqlite") + separate := mode == "path_only" || mode == "local_and_path" + if separate { + t.Setenv("LOG_SQLITE_PATH", logPath) + } + require.NoError(t, InitDB()) + t.Cleanup(func() { require.NoError(t, closeDB(DB)) }) + require.NoError(t, InitLogDB()) + if LOG_DB != DB { + t.Cleanup(func() { require.NoError(t, closeDB(LOG_DB)) }) + } + require.True(t, LOG_DB.Migrator().HasTable(&Log{})) + require.True(t, LOG_DB.Migrator().HasTable(&BillingLogReceipt{})) + require.NoError(t, DB.Create(&User{Id: 9541, Username: "sqlite-config-user", Quota: 1000}).Error) + log := Log{UserId: 9541, Type: LogTypeConsume, Quota: 7} + require.NoError(t, LOG_DB.Create(&log).Error) + var logs int64 + require.NoError(t, DB.Model(&Log{}).Count(&logs).Error) + if separate { + require.EqualValues(t, 0, logs, "log writes must not reach the business file") + require.False(t, LOG_DB.Migrator().HasTable(&User{})) + require.False(t, LOG_DB.Migrator().HasTable(&BillingSettlement{})) + require.FileExists(t, logPath) + } else { + require.EqualValues(t, 1, logs) + if mode == "default" { + require.Same(t, DB, LOG_DB) + } + } + require.NoError(t, CloseDB()) + mainSQL, err := DB.DB() + require.NoError(t, err) + require.Error(t, mainSQL.Ping()) + logSQL, err := LOG_DB.DB() + require.NoError(t, err) + require.Error(t, logSQL.Ping()) + require.NoError(t, InitDB()) + require.NoError(t, InitLogDB()) + var user User + require.NoError(t, DB.First(&user, 9541).Error) + require.Equal(t, int64(1000), user.Quota) + var persisted Log + require.NoError(t, LOG_DB.First(&persisted, log.Id).Error) + require.Equal(t, 7, persisted.Quota) + }) + } +} + +func TestSQLiteLogConfigBackfillIdentity(t *testing.T) { + dir := setupSQLiteLogConfigTest(t) + require.Equal(t, logRetryMarkerBackfillOptionKey, logRetryMarkerBackfillCompletionKey()) + t.Setenv("LOG_SQL_DSN", "local") + legacyKey := logRetryMarkerBackfillOptionKey + ":" + common.Sha1([]byte(common.DatabaseTypeSQLite+"\nlocal\n"+common.SQLitePath)) + require.Equal(t, legacyKey, logRetryMarkerBackfillCompletionKey()) + t.Setenv("LOG_SQLITE_PATH", filepath.Join(dir, "logs-one.sqlite")) + first := logRetryMarkerBackfillCompletionKey() + require.NotEqual(t, legacyKey, first) + require.NotContains(t, first, dir) + t.Setenv("LOG_SQLITE_PATH", filepath.Join(dir, "logs-two.sqlite")) + require.NotEqual(t, first, logRetryMarkerBackfillCompletionKey(), "a new file must get its own backfill") + t.Setenv("LOG_SQL_DSN", "") + require.NotEqual(t, logRetryMarkerBackfillOptionKey, logRetryMarkerBackfillCompletionKey()) +} + +func TestSQLiteLogConfigBackfillOnFileSwitch(t *testing.T) { + dir := setupSQLiteLogConfigTest(t) + require.NoError(t, InitDB()) + t.Cleanup(func() { require.NoError(t, closeDB(DB)) }) + for _, name := range []string{"logs-one.sqlite", "logs-two.sqlite"} { + path := filepath.Join(dir, name) + db, err := gorm.Open(sqlite.Open(path), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&Log{})) + log := Log{UserId: 1, Type: LogTypeConsume, Other: common.MapToJsonStr(map[string]interface{}{"retry_log": true})} + require.NoError(t, db.Create(&log).Error) + require.NoError(t, db.Model(&Log{}).Where("id = ?", log.Id).UpdateColumn("is_retry", false).Error) + require.NoError(t, closeDB(db)) + t.Setenv("LOG_SQLITE_PATH", path) + require.NoError(t, InitLogDB()) + var actual Log + require.NoError(t, LOG_DB.First(&actual, log.Id).Error) + require.True(t, actual.IsRetry, "each configured file requires its own legacy marker backfill") + require.NoError(t, closeDB(LOG_DB)) + } +} + +func TestSQLiteLogConfigOpenFailureDoesNotFallback(t *testing.T) { + dir := setupSQLiteLogConfigTest(t) + t.Setenv("LOG_SQLITE_PATH", filepath.Join(dir, "missing-directory", "logs.sqlite")) + db, err := chooseDB("LOG_SQL_DSN", true) + if db != nil { + t.Cleanup(func() { _ = closeDB(db) }) + } + require.Error(t, err, "an invalid explicit file must not open the business file") + require.NoFileExists(t, common.SQLitePath) +} + +func TestSQLiteLogConfigRemoteConflictIsRedacted(t *testing.T) { + dir := setupSQLiteLogConfigTest(t) + t.Setenv("LOG_SQLITE_PATH", filepath.Join(dir, "logs.sqlite")) + for _, dsn := range []string{"postgres://synthetic:config-test-secret@127.0.0.1:1/logs", "synthetic:config-test-secret@tcp(127.0.0.1:1)/logs"} { + t.Setenv("LOG_SQL_DSN", dsn) + db, err := chooseDB("LOG_SQL_DSN", true) + require.Error(t, err) + require.Nil(t, db) + require.Contains(t, err.Error(), "LOG_SQLITE_PATH") + require.NotContains(t, err.Error(), "config-test-secret") + require.NotContains(t, err.Error(), "127.0.0.1") + } + require.NoFileExists(t, os.Getenv("LOG_SQLITE_PATH")) +} + +func TestSQLiteLogConfigInitFailureDoesNotFallback(t *testing.T) { + for _, mode := range []string{"open", "migration", "remote_conflict"} { + t.Run(mode, func(t *testing.T) { + dir := setupSQLiteLogConfigTest(t) + DB, LOG_DB = nil, nil + path := filepath.Join(dir, "logs.sqlite") + if mode == "open" { + path = filepath.Join(dir, "missing-directory", "logs.sqlite") + } + t.Setenv("LOG_SQLITE_PATH", path) + if mode == "remote_conflict" { + t.Setenv("LOG_SQL_DSN", "postgres://synthetic:config-test-secret@127.0.0.1:1/logs") + } + require.NoError(t, InitDB()) + t.Cleanup(func() { require.NoError(t, closeDB(DB)) }) + if mode == "migration" { + logDB, err := gorm.Open(sqlite.Open(path), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, logDB.Exec("CREATE VIEW logs AS SELECT 1 AS id").Error) + require.NoError(t, closeDB(logDB)) + } + err := InitLogDB() + require.Error(t, err) + require.Nil(t, LOG_DB, "failed startup must not silently publish the business connection") + require.NotContains(t, err.Error(), "config-test-secret") + mainSQL, sqlErr := DB.DB() + require.NoError(t, sqlErr) + require.NoError(t, mainSQL.Ping(), "a failed log connection must not close the business connection") + if mode == "migration" { + require.NoError(t, os.Rename(path, path+".failed"), "failed migration must release the SQLite file handle") + } + }) + } +} + +func TestSQLiteLogConfigUsesLogDialectColumns(t *testing.T) { + dir := setupSQLiteLogConfigTest(t) + common.UsingSQLite, common.UsingPostgreSQL = false, true + common.IsMasterNode = false + t.Setenv("LOG_SQLITE_PATH", filepath.Join(dir, "logs.sqlite")) + require.NoError(t, InitLogDB()) + t.Cleanup(func() { require.NoError(t, closeDB(LOG_DB)) }) + require.Equal(t, `"group"`, commonGroupCol) + require.Equal(t, "`group`", logGroupCol) + require.Equal(t, "`key`", logKeyCol) + require.Equal(t, common.DatabaseTypeSQLite, common.LogSqlType) + require.False(t, common.UsingSQLite, "log initialization must not change the main database flags") + require.True(t, common.UsingPostgreSQL) +} diff --git a/service/billing_process_recovery_test.go b/service/billing_process_recovery_test.go index df6b908a..a95dccc3 100644 --- a/service/billing_process_recovery_test.go +++ b/service/billing_process_recovery_test.go @@ -83,10 +83,23 @@ func TestBillingProcessHelper(t *testing.T) { var cfg billingProcessConfig require.NoError(t, common.Unmarshal(data, &cfg)) require.Equal(t, filepath.Join(filepath.Dir(configPath), "ledger.sqlite"), cfg.Database) - db := openBillingProcessDB(t, cfg.Database) - logDB := db + var db, logDB *gorm.DB if cfg.SeparateLogDB { - logDB = openBillingProcessDB(t, filepath.Join(filepath.Dir(cfg.Database), "logs.sqlite")) + // Use the real startup configuration instead of injecting two handles. + t.Setenv("SQL_DSN", "local") + t.Setenv("LOG_SQL_DSN", "") + t.Setenv("LOG_SQLITE_PATH", filepath.Join(filepath.Dir(cfg.Database), "logs.sqlite")) + common.SQLitePath = cfg.Database + common.IsMasterNode = false // The isolated fixture already owns its schema. + require.NoError(t, model.InitDB()) + require.NoError(t, model.InitLogDB()) + db, logDB = model.DB, model.LOG_DB + require.NotSame(t, db, logDB) + configureBillingProcessDB(t, db) + configureBillingProcessDB(t, logDB) + } else { + db = openBillingProcessDB(t, cfg.Database) + logDB = db } model.DB, model.LOG_DB = db, logDB if cfg.LogUnavailable { @@ -252,6 +265,13 @@ func openBillingProcessDB(t *testing.T, path string) *gorm.DB { t.Helper() db, err := gorm.Open(sqlite.Open(path), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) require.NoError(t, err) + configureBillingProcessDB(t, db) + return db +} + +func configureBillingProcessDB(t *testing.T, db *gorm.DB) { + t.Helper() + db.Logger = logger.Default.LogMode(logger.Silent) sqlDB, err := db.DB() require.NoError(t, err) sqlDB.SetMaxOpenConns(1) @@ -259,7 +279,6 @@ func openBillingProcessDB(t *testing.T, path string) *gorm.DB { // synchronous=OFF or memory/WAL shortcuts weaken the interruption test. require.NoError(t, db.Exec("PRAGMA busy_timeout = 30000").Error) t.Cleanup(func() { require.NoError(t, sqlDB.Close()) }) - return db } func seedBillingProcessDB(t *testing.T, funding billingProcessFunding) (string, *gorm.DB) { @@ -319,7 +338,7 @@ func installBillingProcessBarrier(t *testing.T, db, logDB *gorm.DB, cfg billingP return } once.Do(func() { - _, inTx := tx.Statement.ConnPool.(*sql.Tx) + inTx := billingProcessHasSQLTransaction(tx.Statement.ConnPool) pauseBillingProcess(t, cfg, billingProcessBarrier{Phase: cfg.Phase, Table: tx.Statement.Table, Database: database, InsideTransaction: inTx}) }) } @@ -465,6 +484,33 @@ func awaitBillingProcessBarrier(t *testing.T, p *billingTestProcess) billingProc } } +func billingProcessHasSQLTransaction(pool gorm.ConnPool) bool { + switch tx := pool.(type) { + case *sql.Tx: + return tx != nil + case *gorm.PreparedStmtTX: + return tx != nil && billingProcessHasSQLTransaction(tx.Tx) + default: + return false + } +} + +func TestBillingProcessTransactionDetection(t *testing.T) { + db := openBillingProcessDB(t, filepath.Join(t.TempDir(), "transaction.sqlite")) + for _, prepared := range []bool{false, true} { + t.Run(fmt.Sprintf("prepared_%t", prepared), func(t *testing.T) { + connection := db.Session(&gorm.Session{PrepareStmt: prepared}) + require.False(t, billingProcessHasSQLTransaction(connection.Statement.ConnPool), "an open connection is not a transaction") + tx := connection.Begin() + require.NoError(t, tx.Error) + t.Cleanup(func() { require.NoError(t, tx.Rollback().Error) }) + require.True(t, billingProcessHasSQLTransaction(tx.Statement.ConnPool), "the marker must unwrap to a real SQL transaction") + }) + } + require.False(t, billingProcessHasSQLTransaction((*sql.Tx)(nil))) + require.False(t, billingProcessHasSQLTransaction((*gorm.PreparedStmtTX)(nil))) +} + func billingProcessMarkerNotReady(err error) bool { // Windows can temporarily deny access while the atomic rename is observed. // Retry only ERROR_SHARING_VIOLATION (32), never arbitrary I/O failures. @@ -489,7 +535,8 @@ func killBillingTestProcess(t *testing.T, p *billingTestProcess, phase string) b func finishBillingTestProcess(t *testing.T, p *billingTestProcess) { t.Helper() - require.NoError(t, <-p.done, "child process failed\n%s", p.output.text()) + err := <-p.done + require.NoError(t, err, "child process failed\n%s", p.output.text()) } func billingProcessRecord(t *testing.T, db *gorm.DB) model.BillingSettlement { From 33eb872b51d9c56bffcd7c9beb45cb4fba6120cb Mon Sep 17 00:00:00 2001 From: CSCITech Date: Wed, 14 Oct 2026 00:29:23 +0800 Subject: [PATCH 23/24] fix: preserve terminal snapshots and completed tool output --- relay/channel/openai/empty_completion.go | 2 +- relay/channel/openai/relay_responses.go | 18 +- .../openai/responses_terminal_repair_test.go | 29 +++ relay/metadata_funding_settlement_test.go | 179 ++++++++++++++++++ .../chat_to_responses_response.go | 3 + .../openaicompat/custom_tool_finalize_test.go | 48 +++++ 6 files changed, 277 insertions(+), 2 deletions(-) diff --git a/relay/channel/openai/empty_completion.go b/relay/channel/openai/empty_completion.go index 525517e8..87847cc4 100644 --- a/relay/channel/openai/empty_completion.go +++ b/relay/channel/openai/empty_completion.go @@ -240,7 +240,7 @@ func responsesStreamHasVisibleOutput(streamResp *dto.ResponsesStreamResponse) bo return helper.ResponsesStreamEventHasOutput(*streamResp) } return responsesOutputHasVisiblePayload(streamResp.Item) - case "response.completed": + case "response.completed", "response.incomplete", "response.done": return responsesHasVisibleOutput(streamResp.Response) default: return false diff --git a/relay/channel/openai/relay_responses.go b/relay/channel/openai/relay_responses.go index 0b0086a7..a52445ff 100644 --- a/relay/channel/openai/relay_responses.go +++ b/relay/channel/openai/relay_responses.go @@ -25,6 +25,7 @@ const maxPendingResponsesStreamEvents = 32 // arguments twice. This estimate never replaces upstream-reported token counts. type responsesOutputEstimate struct { text strings.Builder + textSnapshot string tools []*responsesToolEstimate toolsByAlias map[string]*responsesToolEstimate } @@ -112,16 +113,31 @@ func (e *responsesOutputEstimate) observe(event dto.ResponsesStreamResponse) { } case "response.completed", "response.incomplete", "response.done": if event.Response != nil { + var snapshot strings.Builder for i, item := range event.Response.Output { e.observe(dto.ResponsesStreamResponse{Type: "response.output_item.done", Item: &item, OutputIndex: &i}) + if item.Type == "message" && (item.Role == "" || item.Role == "assistant") { + for _, content := range item.Content { + snapshot.WriteString(content.Text) + snapshot.WriteString(content.Refusal) + } + } } + e.textSnapshot = snapshot.String() } } } func (e *responsesOutputEstimate) textForTokens() string { var text strings.Builder - text.WriteString(e.text.String()) + deliveredText := e.text.String() + // Terminal-only text is delivered in the full snapshot. Matching deltas + // are a prefix of that snapshot, so count its extension once. Preserve + // delivered deltas when a snapshot is empty or contradicts them. + if strings.HasPrefix(e.textSnapshot, deliveredText) { + deliveredText = e.textSnapshot + } + text.WriteString(deliveredText) for _, tool := range e.tools { text.WriteString(tool.name) arguments := tool.deltas.String() diff --git a/relay/channel/openai/responses_terminal_repair_test.go b/relay/channel/openai/responses_terminal_repair_test.go index 9b4acadc..a5dbe71b 100644 --- a/relay/channel/openai/responses_terminal_repair_test.go +++ b/relay/channel/openai/responses_terminal_repair_test.go @@ -16,6 +16,35 @@ import ( "github.com/stretchr/testify/require" ) +func TestNativeResponsesDeliveredTextSnapshotEstimate(t *testing.T) { + for _, kind := range []string{"output_text", "refusal"} { + for _, delta := range []string{"", "terminal", "terminal snapshot output", "contradicting delta"} { + t.Run(kind+"/delta_"+delta, func(t *testing.T) { + var estimate responsesOutputEstimate + if delta != "" { + estimate.observe(dto.ResponsesStreamResponse{Type: "response." + map[string]string{"output_text": "output_text", "refusal": "refusal"}[kind] + ".delta", Delta: delta}) + } + content := dto.ResponsesOutputContent{Type: kind} + if kind == "refusal" { + content.Refusal = "terminal snapshot output" + } else { + content.Text = "terminal snapshot output" + } + estimate.observe(dto.ResponsesStreamResponse{Type: "response.incomplete", Response: &dto.OpenAIResponsesResponse{Output: []dto.ResponsesOutput{ + {Type: "message", Role: "assistant", Content: []dto.ResponsesOutputContent{content}}, + }}}) + expected := "terminal snapshot output" + if delta == "contradicting delta" { + expected = delta + } + require.Equal(t, expected, estimate.textForTokens(), "snapshots extend matching delivered prefixes once") + estimate.observe(dto.ResponsesStreamResponse{Type: "response.done", Response: &dto.OpenAIResponsesResponse{}}) + require.Equal(t, delta, estimate.textForTokens(), "empty snapshots preserve already delivered deltas") + }) + } + } +} + func TestResponsesToChatLegalTerminalEvents(t *testing.T) { for _, terminal := range []struct { name, event, status, reason, finish string diff --git a/relay/metadata_funding_settlement_test.go b/relay/metadata_funding_settlement_test.go index 5f5f903d..b1f4f7c6 100644 --- a/relay/metadata_funding_settlement_test.go +++ b/relay/metadata_funding_settlement_test.go @@ -13,6 +13,7 @@ import ( "github.com/MAX-API-Next/MAX-API/common" "github.com/MAX-API-Next/MAX-API/constant" + "github.com/MAX-API-Next/MAX-API/dto" "github.com/MAX-API-Next/MAX-API/model" relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" "github.com/MAX-API-Next/MAX-API/service" @@ -248,6 +249,96 @@ func TestMetadataOnlyFailureRetainsOriginalFunding(t *testing.T) { } } +func TestChatToResponsesDoneToolFunding(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitHttpClient() + service.InitTokenEncoders() + oldRetry := common.EmptyCompletionRetryEnabled + common.EmptyCompletionRetryEnabled = false + t.Cleanup(func() { common.EmptyCompletionRetryEnabled = oldRetry }) + for _, kind := range []string{"unnamed", "function", "custom"} { + for _, usageKind := range []string{"known", "zero", "missing"} { + for _, funding := range streamFundingCases() { + t.Run(kind+"/"+usageKind+"/"+funding.name, func(t *testing.T) { + var baseline [3]int + for _, late := range []bool{false, true} { + t.Run(fmt.Sprintf("late_%t", late), func(t *testing.T) { + db := setupStreamFundingLedger(t, funding) + name := "lookup" + if kind == "unnamed" { + name = "" + } + encode := func(value map[string]any) string { + data, err := common.Marshal(value) + require.NoError(t, err) + return string(data) + } + frames := []string{encode(map[string]any{"id": "chat_done_tool", "choices": []any{map[string]any{"index": 0, "delta": map[string]any{ + "content": "kept text", "tool_calls": []any{map[string]any{"index": 0, "id": "call_done", "type": "function", "function": map[string]any{"name": name, "arguments": `{"input":"original"}`}}}, + }}}}), encode(map[string]any{"id": "chat_done_tool", "choices": []any{map[string]any{"index": 0, "delta": map[string]any{}, "finish_reason": "tool_calls"}}})} + if late { + frames = append(frames, encode(map[string]any{"id": "chat_done_tool", "choices": []any{map[string]any{"index": 0, "delta": map[string]any{ + "tool_calls": []any{map[string]any{"index": 0, "id": "call_done", "type": "function", "function": map[string]any{"name": "lookup", "arguments": "late-arguments"}}}, + }}}})) + } + if usageKind != "missing" { + prompt, completion := 4, 2 + if usageKind == "zero" { + prompt, completion = 0, 0 + } + frames = append(frames, encode(map[string]any{"id": "chat_done_tool", "choices": []any{}, "usage": map[string]any{"prompt_tokens": prompt, "completion_tokens": completion, "total_tokens": prompt + completion}})) + } + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, "data: "+strings.Join(frames, "\n\ndata: ")+"\n\ndata: [DONE]\n\n") + })) + t.Cleanup(upstream.Close) + fingerprint := sha256.Sum256([]byte(t.Name())) + requestID := fmt.Sprintf("done-tool-%x", fingerprint[:16]) + c, info, recorder := streamFundingContext(t, "chat_to_responses", requestID, upstream.URL, constant.ChannelTypeAdvancedCustom, funding) + if kind == "custom" { + info.Request.(*dto.OpenAIResponsesRequest).Tools = []byte(`[{"type":"custom","name":"lookup"}]`) + } + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + require.Nil(t, ResponsesHelper(c, info)) + require.NotContains(t, recorder.Body.String(), "late-arguments") + require.Contains(t, recorder.Body.String(), "kept text") + var settlement model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).First(&settlement).Error) + var effect model.BillingSettlementEffect + require.NoError(t, common.UnmarshalJsonStr(settlement.EffectPayload, &effect)) + actual := [3]int{int(effect.Quota), effect.PromptTokens, effect.CompletionTokens} + if !late { + baseline = actual + } else { + require.Equal(t, baseline, actual, "stale tool fragments must not change accepted output billing") + } + if usageKind == "known" { + require.Equal(t, [3]int{6, 4, 2}, actual) + } + if usageKind == "zero" { + require.Equal(t, [3]int{}, actual) + } + assertStreamFundingBalances(t, db, funding, requestID, actual[0], true, actual[1], actual[2]) + replayInfo := &relaycommon.RelayInfo{RequestId: requestID, UserId: info.UserId, TokenId: info.TokenId, TokenKey: info.TokenKey, OriginModelName: info.OriginModelName, ForcePreConsume: true, UserSetting: info.UserSetting} + replay, replayErr := service.NewBillingSession(c, replayInfo, 100) + require.Nil(t, replayErr) + require.NoError(t, replay.SettleWithEffect(actual[0], &effect)) + require.NoError(t, replay.SettleWithEffect(actual[0], &effect)) + replay.Refund(c) + model.ProcessPendingBillingSettlementsOnce() + assertStreamFundingBalances(t, db, funding, requestID, actual[0], true, actual[1], actual[2]) + require.EqualValues(t, 1, calls.Load()) + }) + } + }) + } + } + } +} + func TestMetadataOnlyCompletedResponsesPreservesReportedUsage(t *testing.T) { gin.SetMode(gin.TestMode) service.InitHttpClient() @@ -272,6 +363,94 @@ func TestMetadataOnlyCompletedResponsesPreservesReportedUsage(t *testing.T) { } } +func TestNativeResponsesTerminalSnapshotFunding(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitHttpClient() + service.InitTokenEncoders() + oldRetry, oldTimes := common.EmptyCompletionRetryEnabled, common.RetryTimes + t.Cleanup(func() { common.EmptyCompletionRetryEnabled, common.RetryTimes = oldRetry, oldTimes }) + common.RetryTimes = 1 + const payload = "terminal snapshot output" + for _, event := range []string{"response.completed", "response.incomplete", "response.done"} { + for _, kind := range []string{"text", "function_call", "custom_tool_call"} { + for _, usageKind := range []string{"known", "zero", "missing"} { + for _, retry := range []bool{false, true} { + for _, funding := range streamFundingCases() { + t.Run(fmt.Sprintf("%s/%s/%s/retry_%t/%s", event, kind, usageKind, retry, funding.name), func(t *testing.T) { + common.EmptyCompletionRetryEnabled = retry + db := setupStreamFundingLedger(t, funding) + output := map[string]any{"type": kind, "id": "item_snapshot", "status": "completed"} + estimatedText := payload + if kind == "text" { + output["type"], output["role"] = "message", "assistant" + output["content"] = []any{map[string]any{"type": "output_text", "text": payload}} + } else { + output["name"], output["call_id"] = "lookup", "call_snapshot" + field := "arguments" + if kind == "custom_tool_call" { + field = "input" + } + output[field] = payload + estimatedText = "lookup" + payload + } + response := map[string]any{"id": "resp_snapshot", "status": "completed", "output": []any{output}} + if event == "response.incomplete" { + response["status"] = "incomplete" + response["incomplete_details"] = map[string]any{"reason": "max_output_tokens"} + } + prompt, completion := 4, 2 + if usageKind == "zero" { + prompt, completion = 0, 0 + } + if usageKind != "missing" { + response["usage"] = map[string]int{"input_tokens": prompt, "output_tokens": completion, "total_tokens": prompt + completion} + } else { + prompt, completion = 10, service.CountTextToken(estimatedText, "gpt-test") + } + frame, err := common.Marshal(map[string]any{"type": event, "response": response}) + require.NoError(t, err) + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + calls.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + _, _ = fmt.Fprint(w, "data: "+string(frame)+"\n\n") + })) + t.Cleanup(upstream.Close) + fingerprint := sha256.Sum256([]byte(t.Name())) + requestID := fmt.Sprintf("terminal-snapshot-%x", fingerprint[:16]) + c, info, recorder := streamFundingContext(t, "responses_native", requestID, upstream.URL, constant.ChannelTypeOpenAI, funding) + require.Nil(t, service.PreConsumeBilling(c, 100, info)) + require.Nil(t, ResponsesHelper(c, info), "terminal snapshot output must not enter empty retry") + require.True(t, info.HasRecordedChannelFirstResult()) + require.Contains(t, recorder.Body.String(), payload) + require.EqualValues(t, 1, calls.Load()) + quota := prompt + completion + assertStreamFundingBalances(t, db, funding, requestID, quota, true, prompt, completion) + var settlement model.BillingSettlement + require.NoError(t, db.Where("operation_key = ?", model.BillingRequestFinalizeOperationKey(requestID)).First(&settlement).Error) + require.Equal(t, funding.source, settlement.Source) + require.EqualValues(t, quota-100, settlement.AppliedFundingDelta) + var effect model.BillingSettlementEffect + require.NoError(t, common.UnmarshalJsonStr(settlement.EffectPayload, &effect)) + require.True(t, effect.QuotaIsActual) + require.EqualValues(t, quota, effect.Quota) + replayInfo := &relaycommon.RelayInfo{RequestId: requestID, UserId: info.UserId, TokenId: info.TokenId, TokenKey: info.TokenKey, + OriginModelName: info.OriginModelName, ForcePreConsume: true, UserSetting: info.UserSetting} + replay, replayErr := service.NewBillingSession(c, replayInfo, 100) + require.Nil(t, replayErr) + require.NoError(t, replay.SettleWithEffect(quota, &effect)) + require.NoError(t, replay.SettleWithEffect(quota, &effect)) + replay.Refund(c) + model.ProcessPendingBillingSettlementsOnce() + assertStreamFundingBalances(t, db, funding, requestID, quota, true, prompt, completion) + }) + } + } + } + } + } +} + // A terminal event name cannot override a failed response status. Exercise the // native adapter and original funding ledger, including durable failure replay. func TestNativeResponsesTerminalFailureRetainsOriginalFunding(t *testing.T) { diff --git a/service/openaicompat/chat_to_responses_response.go b/service/openaicompat/chat_to_responses_response.go index 896133eb..0d8fe117 100644 --- a/service/openaicompat/chat_to_responses_response.go +++ b/service/openaicompat/chat_to_responses_response.go @@ -382,6 +382,9 @@ func (s *ChatToResponsesStreamState) appendToolCallDelta(toolCall dto.ToolCallRe chatIndex = *toolCall.Index } tool := s.toolsByIndex[chatIndex] + if tool != nil && tool.Done { + return nil, nil + } events := make([]ChatToResponsesStreamEvent, 0, 2) if tool == nil { tool = &chatToResponsesStreamTool{ diff --git a/service/openaicompat/custom_tool_finalize_test.go b/service/openaicompat/custom_tool_finalize_test.go index 36b7fec3..e6f2e064 100644 --- a/service/openaicompat/custom_tool_finalize_test.go +++ b/service/openaicompat/custom_tool_finalize_test.go @@ -46,6 +46,54 @@ func TestFlushChatCompletionsStreamToResponsesOutputDefersTerminalUsage(t *testi } } +func TestChatToResponsesDoneToolIgnoresLateDeltas(t *testing.T) { + for _, kind := range []string{"unnamed", "function", "custom"} { + for _, changedIdentity := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/changed_identity_%t", kind, changedIdentity), func(t *testing.T) { + state := NewChatToResponsesStreamState("resp_late", "test") + state.CustomToolNames = map[string]struct{}{"lookup": {}} + if kind != "custom" { + state.CustomToolNames = nil + } + name := "lookup" + if kind == "unnamed" { + name = "" + } + _, err := ChatCompletionsStreamChunkToResponsesEvents(&dto.ChatCompletionsStreamResponse{Choices: []dto.ChatCompletionsStreamResponseChoice{ + {Delta: dto.ChatCompletionsStreamResponseChoiceDelta{ToolCalls: []dto.ToolCallResponse{{Index: common.GetPointer(0), ID: "call_late", Function: dto.FunctionResponse{Name: name, Arguments: `{"input":"original"}`}}}}}, + }}, state) + require.NoError(t, err) + _, err = ChatCompletionsStreamChunkToResponsesEvents(&dto.ChatCompletionsStreamResponse{Choices: []dto.ChatCompletionsStreamResponseChoice{ + {FinishReason: common.GetPointer("tool_calls")}, + }}, state) + require.NoError(t, err) + before, err := common.Marshal(state.finalResponse()) + require.NoError(t, err) + id, lateName := "call_late", "lookup" + if changedIdentity { + id, lateName = "call_changed", "changed_name" + } + events, err := ChatCompletionsStreamChunkToResponsesEvents(&dto.ChatCompletionsStreamResponse{Choices: []dto.ChatCompletionsStreamResponseChoice{ + {Delta: dto.ChatCompletionsStreamResponseChoiceDelta{ToolCalls: []dto.ToolCallResponse{{Index: common.GetPointer(0), ID: id, Function: dto.FunctionResponse{Name: lateName, Arguments: "late-arguments"}}}}}, + }}, state) + require.NoError(t, err, "completed tool identity is frozen; stale fragments have no authority") + require.Empty(t, events, "no new item or arguments may follow tool done") + after, err := common.Marshal(state.finalResponse()) + require.NoError(t, err) + require.JSONEq(t, string(before), string(after)) + terminal := FinalizeChatCompletionsStreamToResponses(state) + require.Len(t, terminal, 1) + if kind == "unnamed" { + require.Equal(t, "response.incomplete", terminal[0].Type) + require.Equal(t, "missing_tool_name", terminal[0].Payload.Response.IncompleteDetails.Reason) + require.Empty(t, terminal[0].Payload.Response.Output) + } + require.Empty(t, FinalizeChatCompletionsStreamToResponses(state)) + }) + } + } +} + func TestChatToResponsesUnnamedToolCallCannotComplete(t *testing.T) { for _, finish := range []string{"", "tool_calls", "length"} { for _, flushFirst := range []bool{false, true} { From 98709043536005d24dd4fff21c97738246fc9fbc Mon Sep 17 00:00:00 2001 From: CSCITech Date: Wed, 14 Oct 2026 01:36:51 +0800 Subject: [PATCH 24/24] fix: preserve explicit Claude custom converter routes --- ...chat_responses_protocol_settlement_test.go | 30 +++++++++++++++---- relay/claude_handler.go | 3 +- 2 files changed, 25 insertions(+), 8 deletions(-) diff --git a/relay/chat_responses_protocol_settlement_test.go b/relay/chat_responses_protocol_settlement_test.go index 89de9e81..6ce1a2ff 100644 --- a/relay/chat_responses_protocol_settlement_test.go +++ b/relay/chat_responses_protocol_settlement_test.go @@ -41,6 +41,8 @@ func TestGlobalChatResponsesProtocolRetainsOriginalFunding(t *testing.T) { {"custom_claude", "claude", dto.AdvancedCustomConverterOpenAIChatCompletionsToAnthropicMessages, constant.ChannelTypeAdvancedCustom, "", false}, {"custom_gemini", "gemini", dto.AdvancedCustomConverterOpenAIChatCompletionsToGeminiGenerateContent, constant.ChannelTypeAdvancedCustom, "", false}, {"custom_responses", "responses", dto.AdvancedCustomConverterOpenAIChatCompletionsToOpenAIResponses, constant.ChannelTypeAdvancedCustom, "", false}, + {"claude_custom_chat", "chat_complete", dto.AdvancedCustomConverterAnthropicMessagesToOpenAIChatCompletions, constant.ChannelTypeAdvancedCustom, "", true}, + {"claude_custom_none", "responses", dto.AdvancedCustomConverterNone, constant.ChannelTypeAdvancedCustom, "", true}, {"zhipu_standard", "responses", "", constant.ChannelTypeZhipu_v4, "", false}, {"zhipu_coding", "chat_native", "", constant.ChannelTypeZhipu_v4, "glm-coding-plan", false}, {"zhipu_coding_international", "chat_native", "", constant.ChannelTypeZhipu_v4, "glm-coding-plan-international", false}, @@ -56,7 +58,7 @@ func TestGlobalChatResponsesProtocolRetainsOriginalFunding(t *testing.T) { } for _, provider := range providers { modes := []bool{true} - if provider.name == "claude" || provider.name == "gemini" || provider.baseAlias != "" { + if provider.name == "claude" || provider.name == "gemini" || provider.name == "claude_custom_chat" || provider.baseAlias != "" { modes = []bool{false, true} } for _, enabled := range modes { @@ -103,6 +105,8 @@ func TestGlobalChatResponsesProtocolRetainsOriginalFunding(t *testing.T) { info.IsStream = stream if provider.claudeClient { info.RelayFormat = types.RelayFormatClaude + info.IsFirstRequest = true + info.ClaudeConvertInfo = &relaycommon.ClaudeConvertInfo{LastMessagesType: relaycommon.LastMessageTypeNone} info.Request = &dto.ClaudeRequest{Model: "gpt-test", Messages: []dto.ClaudeMessage{{Role: "user", Content: "test"}}, MaxTokens: common.GetPointer(uint(64)), Stream: common.GetPointer(stream)} c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) info.RequestURLPath = "/v1/messages" @@ -114,10 +118,12 @@ func TestGlobalChatResponsesProtocolRetainsOriginalFunding(t *testing.T) { path := "/v1/responses" if provider.wire == "claude" { path = "/v1/messages" + } else if provider.wire == "chat_complete" { + path = "/v1/chat/completions" } else if provider.wire == "gemini" { path = "/v1beta/models/{model}:generateContent" } - c.Set(string(constant.ContextKeyChannelOtherSetting), dto.ChannelOtherSettings{AdvancedCustom: &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{IncomingPath: "/v1/chat/completions", UpstreamPath: path, Converter: provider.converter}}}}) + c.Set(string(constant.ContextKeyChannelOtherSetting), dto.ChannelOtherSettings{AdvancedCustom: &dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{IncomingPath: c.Request.URL.Path, UpstreamPath: path, Converter: provider.converter}}}}) } require.Nil(t, service.PreConsumeBilling(c, 100, info)) var apiErr *types.MaxAPIError @@ -147,8 +153,12 @@ func TestGlobalChatResponsesProtocolRetainsOriginalFunding(t *testing.T) { require.Contains(t, strings.ToLower(request.path), "generatecontent") require.Contains(t, request.body, "contents") require.NotContains(t, request.body, "input") - case "chat_native": - require.Equal(t, "/coding/chat/completions", request.path) + case "chat_native", "chat_complete": + path := "/coding/chat/completions" + if provider.channelType == constant.ChannelTypeAdvancedCustom { + path = "/v1/chat/completions" + } + require.Equal(t, path, request.path) require.Contains(t, request.body, "messages") require.NotContains(t, request.body, "input") default: @@ -356,7 +366,7 @@ func globalBridgeWireBody(t *testing.T, wire string, stream bool) string { t.Helper() if !stream { switch wire { - case "chat_native": + case "chat_native", "chat_complete": return `{"id":"chat_funding","object":"chat.completion","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"partial output"},"finish_reason":"stop"}],"usage":{"prompt_tokens":4,"completion_tokens":2,"total_tokens":6}}` case "claude": return `{"id":"msg_funding","type":"message","role":"assistant","model":"gpt-test","content":[{"type":"text","text":"partial output"}],"stop_reason":"end_turn","usage":{"input_tokens":4,"output_tokens":2}}` @@ -368,8 +378,16 @@ func globalBridgeWireBody(t *testing.T, wire string, stream bool) string { } var frames []string switch wire { - case "chat_native": + case "chat_native", "chat_complete": _, frames = streamFundingFrames(t, "chat_native", "known") + if wire == "chat_complete" { + // The Claude reply converter requires a complete Chat protocol, not + // the deliberately interrupted funding fixture's content/usage pair. + role := `{"id":"chat_funding","choices":[{"index":0,"delta":{"role":"assistant"}}]}` + finish := `{"id":"chat_funding","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}` + usage := `{"id":"chat_funding","choices":[],"usage":{"prompt_tokens":4,"completion_tokens":2,"total_tokens":6}}` + frames = []string{role, frames[0], finish, usage} + } for i, frame := range frames { var chunk map[string]any require.NoError(t, common.UnmarshalJsonStr(frame, &chunk)) diff --git a/relay/claude_handler.go b/relay/claude_handler.go index dfc33a29..77a6410b 100644 --- a/relay/claude_handler.go +++ b/relay/claude_handler.go @@ -92,8 +92,7 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (maxAPIError *typ if !model_setting.GetGlobalSettings().PassThroughRequestEnabled && !info.ChannelSetting.PassThroughBodyEnabled && - service.SupportsOpenAIResponsesUpstreamForBaseURL(info.ChannelType, info.ChannelBaseUrl) && - service.ShouldChatCompletionsUseResponsesGlobal(info.ChannelId, info.ChannelType, info.OriginModelName) { + shouldUseGlobalResponsesBridge(c, info) { openAIRequest, convErr := service.ClaudeToOpenAIRequest(*request, info) if convErr != nil { return types.NewError(convErr, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())