diff --git a/.gitignore b/.gitignore index 529ad2bd..a491d466 100644 --- a/.gitignore +++ b/.gitignore @@ -41,4 +41,5 @@ skills-lock.json Operation.md /home .tools - +AGENTS.md +CLAUDE.md diff --git a/AGENTS.md b/AGENTS.md index e46dfa0b..86e509ba 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -144,3 +144,39 @@ This includes investigation notes, review outputs, scratch plans, working drafts Do NOT place these files in the repository root, `docs/`, `.github/`, or source directories unless the content is explicitly intended to be public project documentation. `.tmp/` is ignored by Git; use it to prevent accidental publication of sensitive or internal process information to the open-source GitHub repository. If process documents are found outside `.tmp/`, move them into `.tmp/` before committing. Local agent workflow files, including `.agents/skills/` and `.agents/upstream-sync/`, are private project-maintenance material for this workspace and MUST NOT be committed to the public GitHub repository. Keep reusable private agent instructions there locally, or move public-safe documentation into normal project docs when it is intentionally meant to be open source. + +### Rule 9: Monitor Independent New API Upstream Modules and Plugins + +New API-derived capabilities that evolve outside MAX-API's normal dependency graph MUST be checked explicitly. This includes independent Go modules such as `relaykit`, bundled Task Plugins under `plugins/`, and their host-side plugin or protocol-conversion frameworks. + +Run the private upstream component check before any of the following work: + +- changing relay DTOs, protocol conversion, streaming conversion, usage normalization, or finish-reason mapping; +- changing Task Plugin loading, routing, sandboxing, artifact handling, usage facts, or plugin billing integration; +- preparing a release, performing an upstream feature/security comparison, or claiming parity with the latest New API; +- adopting or upgrading an independent New API module or plugin. + +Required procedure: + +1. Read `.agents/upstream-sync/components.json`, `.agents/upstream-sync/new-api-sync.md`, and the living implementation plan at `new-api/.tmp/MAX_API_UPSTREAM_IMPLEMENTATION_PLAN.md`. +2. Run the checker from the repository root with the cross-platform `pwsh` executable when available: + `pwsh -NoProfile -ExecutionPolicy Bypass -File .agents/upstream-sync/check-new-api.ps1`. + On Windows hosts without `pwsh`, use the equivalent Windows PowerShell fallback: + `powershell -NoProfile -ExecutionPolicy Bypass -File .agents/upstream-sync/check-new-api.ps1`. +3. Treat the ignored `new-api/new-api` source copy as reference material only. Its presence does not prove that it matches upstream HEAD. +4. If the remote state cannot be verified, report `需要进一步验证`; do not describe the local copy as current. +5. Detection is not acceptance: never automatically merge or wholesale copy an upstream change. Classify it as `pending`, `accepted`, `ported`, `skipped`, `superseded`, or `blocked`, with rationale. +6. Advance the reviewed baseline only after the affected code and compatibility gates have been inspected and the relevant tests have passed. Merely detecting a newer commit MUST NOT advance the baseline. +7. Keep generated comparison reports under `.tmp/`. Keep the reusable checker, component registry, and local sync state under `.agents/upstream-sync/`; these remain private and uncommitted under Rule 8. +8. Before implementing a listed capability, update the living plan with the current MAX/upstream baselines, revalidated evidence, scope, dependency, license decision, risks, and test plan; mark the item `实施中`. +9. After each related iteration, update the same plan with actual files and key logic, test results, residual verification gaps, upstream component status, delivery status, change log, and the next iteration entry. A code change is not complete until this plan update is performed. + +RelayKit-specific requirements: + +- Describe RelayKit accurately as an independent, full-matrix, auditable protocol-conversion module. MAX-API already has multiple conversion paths and `BillingUsage`; do not claim that MAX-API has no protocol conversion. +- Review request, non-stream response, and streaming response conversion separately, including finalization behavior. +- Verify explicit zero-value preservation, tool calls, multimodal input, reasoning content, finish reasons, and usage mapping. +- Review changes to converter identity, quality, steps, usage, and diagnostics as compatibility-sensitive interface changes. +- RelayKit does not own billing. Preserve MAX-API's billing, reservation, settlement, retry, and logging invariants, and prove that normalized usage is applied exactly once. +- If the private checker or registry is unavailable in a fresh workspace, perform the equivalent remote commit/path comparison manually and record the missing maintenance prerequisite instead of silently skipping the check. +- If the living plan is missing in a fresh workspace, recreate it under `new-api/.tmp/` from the current code and upstream records before continuing; never substitute an old chat summary for the plan. diff --git a/controller/relay.go b/controller/relay.go index d05d87d8..039b1ec1 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -683,7 +683,7 @@ func RelayTask(c *gin.Context) { } c.Request.Body = io.NopCloser(bodyStorage) - result, taskErr = relay.RelayTaskSubmit(c, relayInfo) + result, taskErr = relayTaskSubmitAttempt(c, relayInfo, retryParam.GetRetry()) if taskErr == nil { break } @@ -816,6 +816,11 @@ func respondTaskError(c *gin.Context, taskErr *dto.TaskError) { c.JSON(taskErr.StatusCode, taskErr) } +func relayTaskSubmitAttempt(c *gin.Context, relayInfo *relaycommon.RelayInfo, retryIndex int) (*relay.TaskSubmitResult, *dto.TaskError) { + relayInfo.RetryIndex = retryIndex + return relay.RelayTaskSubmit(c, relayInfo) +} + func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *dto.TaskError, retryTimes int) bool { if taskErr == nil { return false diff --git a/controller/relay_retry_test.go b/controller/relay_retry_test.go index 94e8c39d..5f7a7523 100644 --- a/controller/relay_retry_test.go +++ b/controller/relay_retry_test.go @@ -8,6 +8,7 @@ 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/setting/operation_setting" "github.com/MAX-API-Next/MAX-API/types" @@ -15,6 +16,23 @@ import ( "github.com/stretchr/testify/require" ) +func TestRelayTaskSubmitAttemptAdvancesToolUsageRetryIndex(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Set("platform", "invalid-test-platform") + info := &relaycommon.RelayInfo{ + OriginModelName: "task-test", + RetryIndex: 0, + ToolUsage: relaycommon.NewToolUsageLedger("task-test"), + } + info.ToolUsage.BeginAttempt(0) + + _, taskErr := relayTaskSubmitAttempt(c, info, 1) + require.NotNil(t, taskErr) + require.Equal(t, 1, info.RetryIndex) + require.Equal(t, 1, info.ToolUsageSnapshot().Attempt) +} + func TestShouldRetryEmptyResponseRequiresFeatureFlag(t *testing.T) { gin.SetMode(gin.TestMode) c, _ := gin.CreateTestContext(httptest.NewRecorder()) diff --git a/controller/subscription.go b/controller/subscription.go index 51adb312..d77f75a1 100644 --- a/controller/subscription.go +++ b/controller/subscription.go @@ -138,6 +138,37 @@ type AdminUpsertSubscriptionPlanRequest struct { Plan model.SubscriptionPlan `json:"plan"` } +func subscriptionPlanUpdateMap(plan model.SubscriptionPlan) map[string]interface{} { + updates := map[string]interface{}{ + "title": plan.Title, + "subtitle": plan.Subtitle, + "price_amount": plan.PriceAmount, + "currency": plan.Currency, + "duration_unit": plan.DurationUnit, + "duration_value": plan.DurationValue, + "custom_seconds": plan.CustomSeconds, + "enabled": plan.Enabled, + "sort_order": plan.SortOrder, + "stripe_price_id": plan.StripePriceId, + "creem_product_id": plan.CreemProductId, + "waffo_pancake_product_id": plan.WaffoPancakeProductId, + "max_purchase_per_user": plan.MaxPurchasePerUser, + "total_amount": plan.TotalAmount, + "upgrade_group": plan.UpgradeGroup, + "downgrade_group": plan.DowngradeGroup, + "quota_reset_period": plan.QuotaResetPeriod, + "quota_reset_custom_seconds": plan.QuotaResetCustomSeconds, + "updated_at": common.GetTimestamp(), + } + if plan.AllowBalancePay != nil { + updates["allow_balance_pay"] = *plan.AllowBalancePay + } + if plan.AllowWalletOverflow != nil { + updates["allow_wallet_overflow"] = *plan.AllowWalletOverflow + } + return updates +} + func AdminCreateSubscriptionPlan(c *gin.Context) { if !requirePaymentCompliance(c) { return @@ -168,6 +199,9 @@ func AdminCreateSubscriptionPlan(c *gin.Context) { if req.Plan.AllowBalancePay == nil { req.Plan.AllowBalancePay = common.GetPointer(true) } + if req.Plan.AllowWalletOverflow == nil { + req.Plan.AllowWalletOverflow = common.GetPointer(true) + } if req.Plan.DurationUnit == "" { req.Plan.DurationUnit = model.SubscriptionDurationMonth } @@ -189,6 +223,13 @@ func AdminCreateSubscriptionPlan(c *gin.Context) { return } } + req.Plan.DowngradeGroup = strings.TrimSpace(req.Plan.DowngradeGroup) + if req.Plan.DowngradeGroup != "" { + if _, ok := ratio_setting.GetGroupRatioCopy()[req.Plan.DowngradeGroup]; !ok { + common.ApiErrorMsg(c, "降级分组不存在") + return + } + } req.Plan.QuotaResetPeriod = model.NormalizeResetPeriod(req.Plan.QuotaResetPeriod) if req.Plan.QuotaResetPeriod == model.SubscriptionResetCustom && req.Plan.QuotaResetCustomSeconds <= 0 { common.ApiErrorMsg(c, "自定义重置周期需大于0秒") @@ -256,6 +297,13 @@ func AdminUpdateSubscriptionPlan(c *gin.Context) { return } } + req.Plan.DowngradeGroup = strings.TrimSpace(req.Plan.DowngradeGroup) + if req.Plan.DowngradeGroup != "" { + if _, ok := ratio_setting.GetGroupRatioCopy()[req.Plan.DowngradeGroup]; !ok { + common.ApiErrorMsg(c, "降级分组不存在") + return + } + } req.Plan.QuotaResetPeriod = model.NormalizeResetPeriod(req.Plan.QuotaResetPeriod) if req.Plan.QuotaResetPeriod == model.SubscriptionResetCustom && req.Plan.QuotaResetCustomSeconds <= 0 { common.ApiErrorMsg(c, "自定义重置周期需大于0秒") @@ -264,29 +312,7 @@ func AdminUpdateSubscriptionPlan(c *gin.Context) { err := model.DB.Transaction(func(tx *gorm.DB) error { // update plan (allow zero values updates with map) - updateMap := map[string]interface{}{ - "title": req.Plan.Title, - "subtitle": req.Plan.Subtitle, - "price_amount": req.Plan.PriceAmount, - "currency": req.Plan.Currency, - "duration_unit": req.Plan.DurationUnit, - "duration_value": req.Plan.DurationValue, - "custom_seconds": req.Plan.CustomSeconds, - "enabled": req.Plan.Enabled, - "sort_order": req.Plan.SortOrder, - "stripe_price_id": req.Plan.StripePriceId, - "creem_product_id": req.Plan.CreemProductId, - "waffo_pancake_product_id": req.Plan.WaffoPancakeProductId, - "max_purchase_per_user": req.Plan.MaxPurchasePerUser, - "total_amount": req.Plan.TotalAmount, - "upgrade_group": req.Plan.UpgradeGroup, - "quota_reset_period": req.Plan.QuotaResetPeriod, - "quota_reset_custom_seconds": req.Plan.QuotaResetCustomSeconds, - "updated_at": common.GetTimestamp(), - } - if req.Plan.AllowBalancePay != nil { - updateMap["allow_balance_pay"] = *req.Plan.AllowBalancePay - } + updateMap := subscriptionPlanUpdateMap(req.Plan) if err := tx.Model(&model.SubscriptionPlan{}).Where("id = ?", id).Updates(updateMap).Error; err != nil { return err } diff --git a/controller/subscription_test.go b/controller/subscription_test.go index 1c359768..8fac391d 100644 --- a/controller/subscription_test.go +++ b/controller/subscription_test.go @@ -4,13 +4,16 @@ import ( "bytes" "net/http" "net/http/httptest" + "sync" "testing" "github.com/MAX-API-Next/MAX-API/common" appi18n "github.com/MAX-API-Next/MAX-API/i18n" + "github.com/MAX-API-Next/MAX-API/model" "github.com/MAX-API-Next/MAX-API/setting/operation_setting" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" + "gorm.io/gorm/schema" ) type subscriptionAPIResponse struct { @@ -82,3 +85,25 @@ func TestAdminResetPlanSubscriptionsRequiresPaymentCompliance(t *testing.T) { require.False(t, response.Success) require.NotEmpty(t, response.Message) } + +func TestSubscriptionPlanUpdateMapUsesOnlyModelColumns(t *testing.T) { + parsed, err := schema.Parse(&model.SubscriptionPlan{}, &sync.Map{}, schema.NamingStrategy{}) + require.NoError(t, err) + columns := make(map[string]struct{}, len(parsed.DBNames)) + for _, name := range parsed.DBNames { + columns[name] = struct{}{} + } + allowBalance := false + allowWallet := false + updates := subscriptionPlanUpdateMap(model.SubscriptionPlan{ + Title: "plan", AllowBalancePay: &allowBalance, AllowWalletOverflow: &allowWallet, + DowngradeGroup: "default", + }) + for key := range updates { + if _, ok := columns[key]; !ok { + t.Fatalf("subscription plan update key %q is not a model column", key) + } + } + require.Contains(t, updates, "allow_wallet_overflow") + require.Contains(t, updates, "downgrade_group") +} diff --git a/dto/gemini.go b/dto/gemini.go index 6a445d6d..79c75ec5 100644 --- a/dto/gemini.go +++ b/dto/gemini.go @@ -238,8 +238,10 @@ func (g *GeminiInlineData) UnmarshalJSON(data []byte) error { } type FunctionCall struct { + ID string `json:"id,omitempty"` FunctionName string `json:"name"` Arguments any `json:"args"` + WillContinue *bool `json:"willContinue,omitempty"` } type GeminiFunctionResponse struct { diff --git a/dto/openai_response.go b/dto/openai_response.go index 720ac566..1cebfacf 100644 --- a/dto/openai_response.go +++ b/dto/openai_response.go @@ -404,13 +404,17 @@ type ResponsesOutput struct { CallId string `json:"call_id,omitempty"` Name string `json:"name,omitempty"` Arguments json.RawMessage `json:"arguments,omitempty"` + Input string `json:"input,omitempty"` } -// ArgumentsString returns function call arguments in the string form expected by Chat Completions. +// ArgumentsString returns function arguments or custom-tool input in the string form expected by Chat Completions. func (r *ResponsesOutput) ArgumentsString() string { if r == nil { return "" } + if r.Type == BuildInCallCustomToolCall { + return r.Input + } return ResponsesArgumentsString(r.Arguments) } @@ -433,11 +437,18 @@ type ResponsesReasoningSummaryPart struct { const ( BuildInToolWebSearchPreview = "web_search_preview" + BuildInToolWebSearch = "web_search" BuildInToolFileSearch = "file_search" + BuildInToolGoogleSearch = "google_search" + BuildInToolImageGeneration = "image_generation" ) const ( - BuildInCallWebSearchCall = "web_search_call" + BuildInCallWebSearchCall = "web_search_call" + BuildInCallFileSearchCall = "file_search_call" + BuildInCallFunctionCall = "function_call" + BuildInCallCustomToolCall = "custom_tool_call" + BuildInCallToolUse = "tool_use" ) const ( @@ -449,6 +460,9 @@ const ( 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"` // - response.function_call_arguments.delta diff --git a/model/main.go b/model/main.go index 84592406..e63838af 100644 --- a/model/main.go +++ b/model/main.go @@ -303,6 +303,19 @@ func InitLogDB() (err error) { } func migrateDB() error { + // Capture whether the policy columns existed before AutoMigrate so legacy + // rows can be initialized to the historical allow-wallet default instead of + // inheriting a bool zero value when the column is first introduced. + userSubscriptionsTableExisted := DB != nil && DB.Migrator().HasTable(&UserSubscription{}) + userSubscriptionWalletColumnExisted := userSubscriptionsTableExisted && DB.Migrator().HasColumn(&UserSubscription{}, "allow_wallet_overflow") + userSubscriptionDowngradeColumnExisted := userSubscriptionsTableExisted && DB.Migrator().HasColumn(&UserSubscription{}, "downgrade_group") + planTableExisted := DB != nil && DB.Migrator().HasTable(&SubscriptionPlan{}) + planWalletColumnExisted := planTableExisted && DB.Migrator().HasColumn(&SubscriptionPlan{}, "allow_wallet_overflow") + planDowngradeColumnExisted := planTableExisted && DB.Migrator().HasColumn(&SubscriptionPlan{}, "downgrade_group") + if err := ensureUserSubscriptionPolicyColumns(userSubscriptionsTableExisted, userSubscriptionWalletColumnExisted, userSubscriptionDowngradeColumnExisted); err != nil { + return err + } + // Migrate model_limits column from varchar to text for existing tables if err := migrateTokenModelLimitsToText(); err != nil { return err @@ -355,6 +368,9 @@ func migrateDB() error { if err != nil { return err } + if err := migrateUserSubscriptionPolicyDefaults(userSubscriptionsTableExisted, userSubscriptionWalletColumnExisted, userSubscriptionDowngradeColumnExisted); err != nil { + return err + } if err := migrateMetadataNameKeys(); err != nil { return err } @@ -376,10 +392,16 @@ func migrateDB() error { return err } } else { + if err := ensureSubscriptionPlanPolicyColumns(planTableExisted, planWalletColumnExisted, planDowngradeColumnExisted); err != nil { + return err + } if err := DB.AutoMigrate(&SubscriptionPlan{}); err != nil { return err } } + if err := migrateSubscriptionPlanPolicyDefaults(planTableExisted, planWalletColumnExisted, planDowngradeColumnExisted); err != nil { + return err + } if os.Getenv("LOG_SQL_DSN") == "" { LOG_DB = DB scheduleLogRetryMarkerBackfill() @@ -726,6 +748,128 @@ type sqliteColumnDef struct { DDL string } +func ensureUserSubscriptionPolicyColumns(tableExisted, walletColumnExisted, downgradeColumnExisted bool) error { + if DB == nil || !tableExisted { + return nil + } + addColumn := func(name string, exists bool) error { + return addSubscriptionPolicyColumn(&UserSubscription{}, "user_subscriptions", name, exists) + } + if err := addColumn("allow_wallet_overflow", walletColumnExisted); err != nil { + return err + } + return addColumn("downgrade_group", downgradeColumnExisted) +} + +func ensureSubscriptionPlanPolicyColumns(tableExisted, walletColumnExisted, downgradeColumnExisted bool) error { + if DB == nil || !tableExisted { + return nil + } + addColumn := func(name string, exists bool) error { + return addSubscriptionPolicyColumn(&SubscriptionPlan{}, "subscription_plans", name, exists) + } + if err := addColumn("allow_wallet_overflow", walletColumnExisted); err != nil { + return err + } + return addColumn("downgrade_group", downgradeColumnExisted) +} + +func addSubscriptionPolicyColumn(modelValue any, tableName, columnName string, existed bool) error { + if existed { + return nil + } + if err := execSubscriptionPolicyColumnDDL(tableName, columnName); err != nil { + // Another application node can add the same column after this node's + // startup HasColumn snapshot. Suppress only that successful concurrent + // outcome; every other DDL failure keeps its original error context. + if DB.Migrator().HasColumn(modelValue, columnName) { + return nil + } + return fmt.Errorf("failed to add %s.%s: %w", tableName, columnName, err) + } + return nil +} + +func execSubscriptionPolicyColumnDDL(tableName, columnName string) error { + if !common.UsingPostgreSQL && !common.UsingMySQL && !common.UsingSQLite { + return fmt.Errorf("database dialect is not configured for subscription policy column %s.%s", tableName, columnName) + } + switch tableName { + case "user_subscriptions": + switch columnName { + case "allow_wallet_overflow": + if common.UsingPostgreSQL { + return DB.Exec(`ALTER TABLE "user_subscriptions" ADD COLUMN "allow_wallet_overflow" boolean NOT NULL DEFAULT true`).Error + } else if common.UsingMySQL { + return DB.Exec("ALTER TABLE `user_subscriptions` ADD COLUMN `allow_wallet_overflow` boolean NOT NULL DEFAULT TRUE").Error + } else { + return DB.Exec("ALTER TABLE `user_subscriptions` ADD COLUMN `allow_wallet_overflow` numeric NOT NULL DEFAULT 1").Error + } + case "downgrade_group": + if common.UsingPostgreSQL { + return DB.Exec(`ALTER TABLE "user_subscriptions" ADD COLUMN "downgrade_group" varchar(64) NOT NULL DEFAULT ''`).Error + } else { + return DB.Exec("ALTER TABLE `user_subscriptions` ADD COLUMN `downgrade_group` varchar(64) NOT NULL DEFAULT ''").Error + } + } + case "subscription_plans": + switch columnName { + case "allow_wallet_overflow": + if common.UsingPostgreSQL { + return DB.Exec(`ALTER TABLE "subscription_plans" ADD COLUMN "allow_wallet_overflow" boolean DEFAULT true`).Error + } else if common.UsingMySQL { + return DB.Exec("ALTER TABLE `subscription_plans` ADD COLUMN `allow_wallet_overflow` boolean DEFAULT TRUE").Error + } else { + return DB.Exec("ALTER TABLE `subscription_plans` ADD COLUMN `allow_wallet_overflow` numeric DEFAULT 1").Error + } + case "downgrade_group": + if common.UsingPostgreSQL { + return DB.Exec(`ALTER TABLE "subscription_plans" ADD COLUMN "downgrade_group" varchar(64) NOT NULL DEFAULT ''`).Error + } else { + return DB.Exec("ALTER TABLE `subscription_plans` ADD COLUMN `downgrade_group` varchar(64) NOT NULL DEFAULT ''").Error + } + } + } + return fmt.Errorf("unsupported subscription policy column %s.%s", tableName, columnName) +} + +func migrateUserSubscriptionPolicyDefaults(tableExisted, walletColumnExisted, downgradeColumnExisted bool) error { + if DB == nil || !DB.Migrator().HasTable(&UserSubscription{}) { + return nil + } + // A newly added bool column cannot distinguish legacy rows from an explicit + // false value after the fact. Initialize only rows that predate the column, + // and repair nullable legacy rows, preserving every persisted false value. + if !walletColumnExisted { + if tableExisted { + if err := DB.Model(&UserSubscription{}).Where("allow_wallet_overflow IS NULL").Update("allow_wallet_overflow", true).Error; err != nil { + return fmt.Errorf("failed to backfill user subscription wallet policy: %w", err) + } + } + } else if err := DB.Model(&UserSubscription{}).Where("allow_wallet_overflow IS NULL").Update("allow_wallet_overflow", true).Error; err != nil { + return fmt.Errorf("failed to repair user subscription wallet policy: %w", err) + } + if err := DB.Model(&UserSubscription{}).Where("downgrade_group IS NULL").Update("downgrade_group", "").Error; err != nil { + return fmt.Errorf("failed to backfill user subscription downgrade policy: %w", err) + } + return nil +} + +func migrateSubscriptionPlanPolicyDefaults(tableExisted, walletColumnExisted, downgradeColumnExisted bool) error { + if DB == nil || !DB.Migrator().HasTable(&SubscriptionPlan{}) { + return nil + } + if walletColumnExisted || tableExisted { + if err := DB.Model(&SubscriptionPlan{}).Where("allow_wallet_overflow IS NULL").Update("allow_wallet_overflow", true).Error; err != nil { + return fmt.Errorf("failed to backfill subscription plan wallet policy: %w", err) + } + } + if err := DB.Model(&SubscriptionPlan{}).Where("downgrade_group IS NULL").Update("downgrade_group", "").Error; err != nil { + return fmt.Errorf("failed to backfill subscription plan downgrade policy: %w", err) + } + return nil +} + func ensureSubscriptionPlanTableSQLite() error { if !common.UsingSQLite { return nil @@ -744,11 +888,13 @@ func ensureSubscriptionPlanTableSQLite() error { ` + "`enabled`" + ` numeric DEFAULT 1, ` + "`sort_order`" + ` integer DEFAULT 0, ` + "`allow_balance_pay`" + ` numeric DEFAULT 1, +` + "`allow_wallet_overflow`" + ` numeric DEFAULT 1, ` + "`stripe_price_id`" + ` varchar(128) DEFAULT '', ` + "`creem_product_id`" + ` varchar(128) DEFAULT '', ` + "`waffo_pancake_product_id`" + ` varchar(128) DEFAULT '', ` + "`max_purchase_per_user`" + ` integer DEFAULT 0, ` + "`upgrade_group`" + ` varchar(64) DEFAULT '', +` + "`downgrade_group`" + ` varchar(64) DEFAULT '', ` + "`total_amount`" + ` bigint NOT NULL DEFAULT 0, ` + "`quota_reset_period`" + ` varchar(16) DEFAULT 'never', ` + "`quota_reset_custom_seconds`" + ` bigint DEFAULT 0, @@ -779,11 +925,13 @@ PRIMARY KEY (` + "`id`" + `) {Name: "enabled", DDL: "`enabled` numeric DEFAULT 1"}, {Name: "sort_order", DDL: "`sort_order` integer DEFAULT 0"}, {Name: "allow_balance_pay", DDL: "`allow_balance_pay` numeric DEFAULT 1"}, + {Name: "allow_wallet_overflow", DDL: "`allow_wallet_overflow` numeric DEFAULT 1"}, {Name: "stripe_price_id", DDL: "`stripe_price_id` varchar(128) DEFAULT ''"}, {Name: "creem_product_id", DDL: "`creem_product_id` varchar(128) DEFAULT ''"}, {Name: "waffo_pancake_product_id", DDL: "`waffo_pancake_product_id` varchar(128) DEFAULT ''"}, {Name: "max_purchase_per_user", DDL: "`max_purchase_per_user` integer DEFAULT 0"}, {Name: "upgrade_group", DDL: "`upgrade_group` varchar(64) DEFAULT ''"}, + {Name: "downgrade_group", DDL: "`downgrade_group` varchar(64) DEFAULT ''"}, {Name: "total_amount", DDL: "`total_amount` bigint NOT NULL DEFAULT 0"}, {Name: "quota_reset_period", DDL: "`quota_reset_period` varchar(16) DEFAULT 'never'"}, {Name: "quota_reset_custom_seconds", DDL: "`quota_reset_custom_seconds` bigint DEFAULT 0"}, diff --git a/model/main_migration_test.go b/model/main_migration_test.go index dec3ccaa..8039f3bf 100644 --- a/model/main_migration_test.go +++ b/model/main_migration_test.go @@ -9,6 +9,8 @@ import ( "github.com/MAX-API-Next/MAX-API/common" "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "gorm.io/gorm" ) @@ -101,3 +103,208 @@ func TestMigrateDBRunsCoreAutoMigrateBeforeSubscriptionPlanPriceMigrationFailure t.Fatal("expected core schema auto-migration to run before subscription price migration failure") } } + +func TestUserSubscriptionPolicyMigrationAddsSafeDefaultsForLegacyRows(t *testing.T) { + previousDB := DB + previousLogDB := LOG_DB + t.Cleanup(func() { + DB = previousDB + LOG_DB = previousLogDB + }) + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + DB = db + LOG_DB = db + require.NoError(t, db.Exec(`CREATE TABLE user_subscriptions (id integer primary key, user_id integer, status varchar(32), end_time bigint, updated_at bigint)`).Error) + require.NoError(t, db.Exec(`INSERT INTO user_subscriptions (id, user_id, status, end_time) VALUES (1, 7, 'active', 9999999999)`).Error) + + require.NoError(t, ensureUserSubscriptionPolicyColumns(true, false, false)) + require.NoError(t, migrateUserSubscriptionPolicyDefaults(true, false, false)) + + var got struct { + AllowWalletOverflow bool `gorm:"column:allow_wallet_overflow"` + DowngradeGroup string `gorm:"column:downgrade_group"` + } + require.NoError(t, db.Table("user_subscriptions").First(&got, 1).Error) + assert.True(t, got.AllowWalletOverflow) + assert.Equal(t, "", got.DowngradeGroup) +} + +func TestSubscriptionPolicyMigrationUsesConfiguredDatabaseFlags(t *testing.T) { + previousDB := DB + previousLogDB := LOG_DB + previousSQLite := common.UsingSQLite + previousMySQL := common.UsingMySQL + previousPostgreSQL := common.UsingPostgreSQL + t.Cleanup(func() { + DB = previousDB + LOG_DB = previousLogDB + common.UsingSQLite = previousSQLite + common.UsingMySQL = previousMySQL + common.UsingPostgreSQL = previousPostgreSQL + }) + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + DB = db + LOG_DB = db + // Deliberately keep the SQLite driver while selecting the configured MySQL + // branch. The migration must follow the project flags rather than infer a + // different dialect from DB.Dialector.Name(). + common.UsingSQLite = false + common.UsingMySQL = true + common.UsingPostgreSQL = false + require.NoError(t, db.Exec(`CREATE TABLE user_subscriptions (id integer primary key)`).Error) + require.NoError(t, ensureUserSubscriptionPolicyColumns(true, false, false)) + assert.True(t, db.Migrator().HasColumn("user_subscriptions", "allow_wallet_overflow")) + assert.True(t, db.Migrator().HasColumn("user_subscriptions", "downgrade_group")) + var tableSQL string + require.NoError(t, db.Raw(`SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'user_subscriptions'`).Scan(&tableSQL).Error) + assert.Contains(t, tableSQL, "allow_wallet_overflow` boolean NOT NULL DEFAULT TRUE") + assert.NotContains(t, tableSQL, "allow_wallet_overflow` numeric NOT NULL DEFAULT 1") +} + +func TestSubscriptionPolicyColumnAddToleratesConcurrentCreate(t *testing.T) { + previousDB := DB + previousLogDB := LOG_DB + previousSQLite := common.UsingSQLite + previousMySQL := common.UsingMySQL + previousPostgreSQL := common.UsingPostgreSQL + t.Cleanup(func() { + DB = previousDB + LOG_DB = previousLogDB + common.UsingSQLite = previousSQLite + common.UsingMySQL = previousMySQL + common.UsingPostgreSQL = previousPostgreSQL + }) + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + DB = db + LOG_DB = db + common.UsingSQLite = true + common.UsingMySQL = false + common.UsingPostgreSQL = false + require.NoError(t, db.Exec(`CREATE TABLE user_subscriptions ( + id integer primary key, + allow_wallet_overflow numeric NOT NULL DEFAULT 1, + downgrade_group varchar(64) NOT NULL DEFAULT '' + )`).Error) + + // Simulate another node adding both columns after this node captured stale + // HasColumn=false values but before it executes the ALTER statements. + require.NoError(t, ensureUserSubscriptionPolicyColumns(true, false, false)) + + // The subscription_plans helper is used only on non-SQLite production paths. + // Select the MySQL branch while retaining the SQLite test driver so the same + // deterministic stale-check scenario exercises that path. + common.UsingSQLite = false + common.UsingMySQL = true + require.NoError(t, db.Exec(`CREATE TABLE subscription_plans ( + id integer primary key, + allow_wallet_overflow boolean DEFAULT TRUE, + downgrade_group varchar(64) NOT NULL DEFAULT '' + )`).Error) + require.NoError(t, ensureSubscriptionPlanPolicyColumns(true, false, false)) +} + +func TestSubscriptionPolicyMigrationRejectsMissingDatabaseFlags(t *testing.T) { + previousDB := DB + previousSQLite := common.UsingSQLite + previousMySQL := common.UsingMySQL + previousPostgreSQL := common.UsingPostgreSQL + t.Cleanup(func() { + DB = previousDB + common.UsingSQLite = previousSQLite + common.UsingMySQL = previousMySQL + common.UsingPostgreSQL = previousPostgreSQL + }) + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + DB = db + common.UsingSQLite = false + common.UsingMySQL = false + common.UsingPostgreSQL = false + + err = execSubscriptionPolicyColumnDDL("user_subscriptions", "allow_wallet_overflow") + require.ErrorContains(t, err, "database dialect is not configured") + err = ensureUserSubscriptionPolicyColumns(true, false, true) + require.ErrorContains(t, err, "failed to add user_subscriptions.allow_wallet_overflow") + require.ErrorContains(t, err, "database dialect is not configured") +} + +func TestSubscriptionPlanPolicyMigrationIsIdempotentOnSQLite(t *testing.T) { + previousDB := DB + previousLogDB := LOG_DB + previousSQLite := common.UsingSQLite + t.Cleanup(func() { + DB = previousDB + LOG_DB = previousLogDB + common.UsingSQLite = previousSQLite + }) + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + DB = db + LOG_DB = db + common.UsingSQLite = true + require.NoError(t, db.Exec(`CREATE TABLE subscription_plans (id integer primary key, title varchar(128) NOT NULL, price_amount decimal(10,6) NOT NULL)`).Error) + require.NoError(t, db.Exec(`INSERT INTO subscription_plans (id, title, price_amount) VALUES (1, 'legacy', 1)`).Error) + + require.NoError(t, ensureSubscriptionPlanTableSQLite()) + require.NoError(t, migrateSubscriptionPlanPolicyDefaults(true, false, false)) + require.NoError(t, ensureSubscriptionPlanTableSQLite()) + require.NoError(t, migrateSubscriptionPlanPolicyDefaults(true, true, true)) + + var got struct { + AllowWalletOverflow *bool `gorm:"column:allow_wallet_overflow"` + DowngradeGroup string `gorm:"column:downgrade_group"` + } + require.NoError(t, db.Table("subscription_plans").First(&got, 1).Error) + require.NotNil(t, got.AllowWalletOverflow) + assert.True(t, *got.AllowWalletOverflow) + assert.Equal(t, "", got.DowngradeGroup) +} + +func TestUserSubscriptionPolicyMigrationPreservesExplicitFalse(t *testing.T) { + previousDB := DB + previousLogDB := LOG_DB + t.Cleanup(func() { + DB = previousDB + LOG_DB = previousLogDB + }) + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + DB = db + LOG_DB = db + require.NoError(t, db.Exec(`CREATE TABLE user_subscriptions ( + id integer primary key, + user_id integer, + status varchar(32), + end_time bigint, + allow_wallet_overflow numeric, + downgrade_group varchar(64), + updated_at bigint + )`).Error) + require.NoError(t, db.Exec(`INSERT INTO user_subscriptions (id, user_id, status, end_time, allow_wallet_overflow, downgrade_group) VALUES + (1, 7, 'active', 9999999999, 0, NULL), + (2, 8, 'active', 9999999999, NULL, 'vip')`).Error) + + require.NoError(t, migrateUserSubscriptionPolicyDefaults(true, true, true)) + + var rows []struct { + AllowWalletOverflow *bool `gorm:"column:allow_wallet_overflow"` + DowngradeGroup string `gorm:"column:downgrade_group"` + } + require.NoError(t, db.Table("user_subscriptions").Order("id").Find(&rows).Error) + require.Len(t, rows, 2) + require.NotNil(t, rows[0].AllowWalletOverflow) + assert.False(t, *rows[0].AllowWalletOverflow) + assert.Equal(t, "", rows[0].DowngradeGroup) + require.NotNil(t, rows[1].AllowWalletOverflow) + assert.True(t, *rows[1].AllowWalletOverflow) + assert.Equal(t, "vip", rows[1].DowngradeGroup) +} diff --git a/model/subscription.go b/model/subscription.go index 04db9e92..15680a31 100644 --- a/model/subscription.go +++ b/model/subscription.go @@ -234,6 +234,10 @@ type SubscriptionPlan struct { AllowBalancePay *bool `json:"allow_balance_pay" gorm:"default:true"` + // Allow falling back to wallet balance after subscription quota is exhausted. + // nil preserves the historical default of allowing wallet overflow. + AllowWalletOverflow *bool `json:"allow_wallet_overflow" gorm:"default:true"` + StripePriceId string `json:"stripe_price_id" gorm:"type:varchar(128);default:''"` CreemProductId string `json:"creem_product_id" gorm:"type:varchar(128);default:''"` WaffoPancakeProductId string `json:"waffo_pancake_product_id" gorm:"type:varchar(128);default:''"` @@ -244,6 +248,9 @@ type SubscriptionPlan struct { // Upgrade user group after purchase (empty = no change) UpgradeGroup string `json:"upgrade_group" gorm:"type:varchar(64);default:''"` + // Downgrade user group on expiry (empty = revert to the group held before purchase) + DowngradeGroup string `json:"downgrade_group" gorm:"type:varchar(64);default:''"` + // Total quota (amount in quota units, 0 = unlimited) TotalAmount int64 `json:"total_amount" gorm:"type:bigint;not null;default:0"` @@ -271,6 +278,9 @@ func (p *SubscriptionPlan) NormalizeDefaults() { if p.AllowBalancePay == nil { p.AllowBalancePay = common.GetPointer(true) } + if p.AllowWalletOverflow == nil { + p.AllowWalletOverflow = common.GetPointer(true) + } } // Subscription order (payment -> webhook -> create UserSubscription) @@ -333,6 +343,11 @@ type UserSubscription struct { UpgradeGroup string `json:"upgrade_group" gorm:"type:varchar(64);default:''"` PrevUserGroup string `json:"prev_user_group" gorm:"type:varchar(64);default:''"` + // Snapshots are frozen when the subscription is created so later plan edits + // cannot change an already purchased subscription's billing or lifecycle policy. + DowngradeGroup string `json:"downgrade_group" gorm:"type:varchar(64);default:''"` + AllowWalletOverflow bool `json:"allow_wallet_overflow" gorm:"not null"` + CreatedAt int64 `json:"created_at" gorm:"bigint"` UpdatedAt int64 `json:"updated_at" gorm:"bigint"` } @@ -488,7 +503,7 @@ func getUserGroupByIdTx(tx *gorm.DB, userId int) (string, error) { tx = DB } var group string - if err := tx.Model(&User{}).Where("id = ?", userId).Select(commonGroupCol).Find(&group).Error; err != nil { + if err := withRowLock(tx).Model(&User{}).Where("id = ?", userId).Select(commonGroupCol).Find(&group).Error; err != nil { return "", err } return group, nil @@ -498,17 +513,15 @@ func downgradeUserGroupForSubscriptionTx(tx *gorm.DB, sub *UserSubscription, now if tx == nil || sub == nil { return "", errors.New("invalid downgrade args") } + downgradeGroup := strings.TrimSpace(sub.DowngradeGroup) upgradeGroup := strings.TrimSpace(sub.UpgradeGroup) - if upgradeGroup == "" { + if downgradeGroup == "" && upgradeGroup == "" { return "", nil } currentGroup, err := getUserGroupByIdTx(tx, sub.UserId) if err != nil { return "", err } - if currentGroup != upgradeGroup { - return "", nil - } activeQuery := tx.Where("user_id = ? AND status = ? AND end_time > ? AND id <> ? AND upgrade_group <> ''", sub.UserId, "active", now, sub.Id). Where("TRIM(upgrade_group) = ?", currentGroup). @@ -520,15 +533,27 @@ func downgradeUserGroupForSubscriptionTx(tx *gorm.DB, sub *UserSubscription, now if activeQuery.RowsAffected > 0 { return "", nil } - prevGroup := strings.TrimSpace(sub.PrevUserGroup) - if prevGroup == "" || prevGroup == currentGroup { + target := downgradeGroup + if target == "" { + // Legacy behavior only reverts when this subscription actually elevated + // the user. Keep that guard for manually changed user groups. + if upgradeGroup == "" || currentGroup != upgradeGroup { + return "", nil + } + target = strings.TrimSpace(sub.PrevUserGroup) + } else if upgradeGroup != "" && currentGroup != upgradeGroup { + // An explicit target still must not overwrite a group granted manually + // or by another non-subscription path. + return "", nil + } + if target == "" || target == currentGroup { return "", nil } if err := tx.Model(&User{}).Where("id = ?", sub.UserId). - Update("group", prevGroup).Error; err != nil { + Update("group", target).Error; err != nil { return "", err } - return prevGroup, nil + return target, nil } func CreateUserSubscriptionFromPlanTx(tx *gorm.DB, userId int, plan *SubscriptionPlan, source string) (*UserSubscription, error) { @@ -570,6 +595,10 @@ func CreateUserSubscriptionFromPlanTx(tx *gorm.DB, userId int, plan *Subscriptio lastReset = now.Unix() } upgradeGroup := strings.TrimSpace(plan.UpgradeGroup) + allowWalletOverflow := true + if plan.AllowWalletOverflow != nil { + allowWalletOverflow = *plan.AllowWalletOverflow + } prevGroup := "" if upgradeGroup != "" { currentGroup, err := getUserGroupByIdTx(tx, userId) @@ -585,20 +614,22 @@ func CreateUserSubscriptionFromPlanTx(tx *gorm.DB, userId int, plan *Subscriptio } } sub := &UserSubscription{ - UserId: userId, - PlanId: plan.Id, - AmountTotal: plan.TotalAmount, - AmountUsed: 0, - StartTime: now.Unix(), - EndTime: endUnix, - Status: "active", - Source: source, - LastResetTime: lastReset, - NextResetTime: nextReset, - UpgradeGroup: upgradeGroup, - PrevUserGroup: prevGroup, - CreatedAt: common.GetTimestamp(), - UpdatedAt: common.GetTimestamp(), + UserId: userId, + PlanId: plan.Id, + AmountTotal: plan.TotalAmount, + AmountUsed: 0, + StartTime: now.Unix(), + EndTime: endUnix, + Status: "active", + Source: source, + LastResetTime: lastReset, + NextResetTime: nextReset, + UpgradeGroup: upgradeGroup, + PrevUserGroup: prevGroup, + DowngradeGroup: strings.TrimSpace(plan.DowngradeGroup), + AllowWalletOverflow: allowWalletOverflow, + CreatedAt: common.GetTimestamp(), + UpdatedAt: common.GetTimestamp(), } if err := tx.Create(sub).Error; err != nil { return nil, err @@ -903,6 +934,26 @@ func HasActiveUserSubscription(userId int) (bool, error) { return count > 0, nil } +// UserActiveSubscriptionsAllowWalletOverflow reports whether a subscription-first +// request may fall back to the wallet after subscription quota is exhausted. +// The policy is fail-closed: one active subscription with an explicit false +// snapshot blocks the fallback. Historical rows are migrated with a true +// default, so the bool column remains non-null and deterministic across DBs. +func UserActiveSubscriptionsAllowWalletOverflow(userId int) (bool, error) { + if userId <= 0 { + return false, errors.New("invalid userId") + } + now := common.GetTimestamp() + var strictCount int64 + if err := DB.Model(&UserSubscription{}). + Where("user_id = ? AND status = ? AND end_time > ? AND allow_wallet_overflow = ?", + userId, "active", now, false). + Count(&strictCount).Error; err != nil { + return false, err + } + return strictCount == 0, nil +} + // GetAllUserSubscriptions returns all subscriptions (active and expired) for a user. func GetAllUserSubscriptions(userId int) ([]SubscriptionSummary, error) { if userId <= 0 { @@ -1171,6 +1222,20 @@ func ExpireDueSubscriptions(limit int) (int, error) { for userId := range userIds { cacheGroup := "" err := DB.Transaction(func(tx *gorm.DB) error { + // Lock and snapshot only the rows that this transaction is about to + // expire. The downgrade decision must never replay a policy from an + // older expiration processed by a previous worker. + var dueSubs []UserSubscription + dueQuery := withRowLock(tx). + Where("user_id = ? AND status = ? AND end_time > 0 AND end_time <= ?", userId, "active", now). + Order("end_time asc, id asc"). + Find(&dueSubs) + if dueQuery.Error != nil { + return dueQuery.Error + } + if len(dueSubs) == 0 { + return nil + } res := tx.Model(&UserSubscription{}). Where("user_id = ? AND status = ? AND end_time > 0 AND end_time <= ?", userId, "active", now). Updates(map[string]interface{}{ @@ -1182,27 +1247,36 @@ func ExpireDueSubscriptions(limit int) (int, error) { } expiredCount += int(res.RowsAffected) - // Select the most recent expired upgrade and restore its previous group - // unless an active subscription still grants the current group. - var lastExpired UserSubscription - expiredQuery := tx.Where("user_id = ? AND status = ? AND upgrade_group <> ''", - userId, "expired"). - Order("end_time desc, id desc"). - Limit(1). - Find(&lastExpired) - if expiredQuery.Error != nil || expiredQuery.RowsAffected == 0 { + // Select the most recent row from this transaction's expiration batch + // carrying either an explicit downgrade target or a legacy upgrade snapshot. + var lastExpired *UserSubscription + for i := len(dueSubs) - 1; i >= 0; i-- { + candidate := &dueSubs[i] + if strings.TrimSpace(candidate.UpgradeGroup) != "" || strings.TrimSpace(candidate.DowngradeGroup) != "" { + lastExpired = candidate + break + } + } + if lastExpired == nil { return nil } upgradeGroup := strings.TrimSpace(lastExpired.UpgradeGroup) + downgradeGroup := strings.TrimSpace(lastExpired.DowngradeGroup) prevGroup := strings.TrimSpace(lastExpired.PrevUserGroup) - if upgradeGroup == "" || prevGroup == "" { + if upgradeGroup == "" && downgradeGroup == "" { return nil } currentGroup, err := getUserGroupByIdTx(tx, userId) if err != nil { return err } - if currentGroup != upgradeGroup || currentGroup == prevGroup { + target := downgradeGroup + if target == "" { + if prevGroup == "" || upgradeGroup == "" || currentGroup != upgradeGroup || currentGroup == prevGroup { + return nil + } + target = prevGroup + } else if upgradeGroup != "" && currentGroup != upgradeGroup { return nil } // Keep the current group only when another active subscription grants @@ -1219,11 +1293,14 @@ func ExpireDueSubscriptions(limit int) (int, error) { if activeQuery.RowsAffected > 0 { return nil } + if target == "" || target == currentGroup { + return nil + } if err := tx.Model(&User{}).Where("id = ?", userId). - Update("group", prevGroup).Error; err != nil { + Update("group", target).Error; err != nil { return err } - cacheGroup = prevGroup + cacheGroup = target return nil }) if err != nil { diff --git a/model/subscription_policy_test.go b/model/subscription_policy_test.go new file mode 100644 index 00000000..74fc34c4 --- /dev/null +++ b/model/subscription_policy_test.go @@ -0,0 +1,222 @@ +package model + +import ( + "testing" + "time" + + "github.com/MAX-API-Next/MAX-API/common" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func TestSubscriptionPlanNormalizeDefaultsAllowsWalletOverflowByDefault(t *testing.T) { + var plan SubscriptionPlan + plan.NormalizeDefaults() + require.NotNil(t, plan.AllowWalletOverflow) + assert.True(t, *plan.AllowWalletOverflow) + + allow := false + plan.AllowWalletOverflow = &allow + plan.NormalizeDefaults() + require.NotNil(t, plan.AllowWalletOverflow) + assert.False(t, *plan.AllowWalletOverflow) +} + +func TestCreateUserSubscriptionSnapshotsWalletOverflowAndDowngradePolicy(t *testing.T) { + truncateTables(t) + now := GetDBTimestamp() + user := &User{Id: 9801, Username: "subscription-policy-snapshot", Group: "default", Status: common.UserStatusEnabled} + require.NoError(t, DB.Create(user).Error) + allow := false + plan := &SubscriptionPlan{ + Id: 9802, Title: "Strict", Enabled: true, TotalAmount: 100, + DurationUnit: SubscriptionDurationMonth, DurationValue: 1, + AllowWalletOverflow: &allow, UpgradeGroup: "svip", DowngradeGroup: "vip", + } + require.NoError(t, DB.Create(plan).Error) + + var created *UserSubscription + require.NoError(t, DB.Transaction(func(tx *gorm.DB) error { + var err error + created, err = CreateUserSubscriptionFromPlanTx(tx, user.Id, plan, "test") + return err + })) + require.NotNil(t, created) + assert.False(t, created.AllowWalletOverflow) + assert.Equal(t, "vip", created.DowngradeGroup) + + // Later plan edits must not alter the already purchased snapshot. + allowLater := true + require.NoError(t, DB.Model(&SubscriptionPlan{}).Where("id = ?", plan.Id).Updates(map[string]interface{}{ + "allow_wallet_overflow": allowLater, + "downgrade_group": "default", + }).Error) + var stored UserSubscription + require.NoError(t, DB.First(&stored, created.Id).Error) + assert.False(t, stored.AllowWalletOverflow) + assert.Equal(t, "vip", stored.DowngradeGroup) + assert.Greater(t, stored.StartTime, int64(0)) + assert.Greater(t, stored.EndTime, now) +} + +func TestCreateUserSubscriptionDefaultsWalletOverflowWhenPlanUnset(t *testing.T) { + truncateTables(t) + user := &User{Id: 9811, Username: "subscription-policy-default", Group: "default", Status: common.UserStatusEnabled} + require.NoError(t, DB.Create(user).Error) + plan := &SubscriptionPlan{ + Id: 9812, Title: "Legacy", Enabled: true, TotalAmount: 100, + DurationUnit: SubscriptionDurationMonth, DurationValue: 1, + AllowWalletOverflow: nil, + } + require.NoError(t, DB.Create(plan).Error) + + var sub *UserSubscription + require.NoError(t, DB.Transaction(func(tx *gorm.DB) error { + var err error + sub, err = CreateUserSubscriptionFromPlanTx(tx, user.Id, plan, "test") + return err + })) + require.NotNil(t, sub) + assert.True(t, sub.AllowWalletOverflow) +} + +func TestUserActiveSubscriptionsAllowWalletOverflowBlocksOnAnyStrictSubscription(t *testing.T) { + truncateTables(t) + now := GetDBTimestamp() + user := &User{Id: 9821, Username: "subscription-policy-conflict", Group: "default", Status: common.UserStatusEnabled} + require.NoError(t, DB.Create(user).Error) + truePlan := &SubscriptionPlan{Id: 9822, Title: "Allowed", Enabled: true, DurationUnit: SubscriptionDurationMonth, DurationValue: 1} + falsePlan := &SubscriptionPlan{Id: 9823, Title: "Strict", Enabled: true, DurationUnit: SubscriptionDurationMonth, DurationValue: 1} + allow := true + deny := false + truePlan.AllowWalletOverflow = &allow + falsePlan.AllowWalletOverflow = &deny + require.NoError(t, DB.Create(truePlan).Error) + require.NoError(t, DB.Create(falsePlan).Error) + require.NoError(t, DB.Create(&UserSubscription{Id: 9824, UserId: user.Id, PlanId: truePlan.Id, Status: "active", StartTime: now - 10, EndTime: now + 3600, AllowWalletOverflow: true}).Error) + require.NoError(t, DB.Create(&UserSubscription{Id: 9825, UserId: user.Id, PlanId: falsePlan.Id, Status: "active", StartTime: now - 10, EndTime: now + 3600, AllowWalletOverflow: false}).Error) + + allowed, err := UserActiveSubscriptionsAllowWalletOverflow(user.Id) + require.NoError(t, err) + assert.False(t, allowed) + + require.NoError(t, DB.Model(&UserSubscription{}).Where("id = ?", 9825).Update("allow_wallet_overflow", true).Error) + allowed, err = UserActiveSubscriptionsAllowWalletOverflow(user.Id) + require.NoError(t, err) + assert.True(t, allowed) +} + +func TestDowngradeUserGroupUsesExplicitSnapshotBeforePreviousGroup(t *testing.T) { + truncateTables(t) + now := GetDBTimestamp() + user := &User{Id: 9831, Username: "subscription-explicit-downgrade", Group: "svip", Status: common.UserStatusEnabled} + sub := &UserSubscription{Id: 9832, UserId: user.Id, Status: "cancelled", EndTime: now + time.Hour.Milliseconds()/1000, UpgradeGroup: "svip", PrevUserGroup: "vip", DowngradeGroup: "default"} + require.NoError(t, DB.Create(user).Error) + require.NoError(t, DB.Create(sub).Error) + + var target string + require.NoError(t, DB.Transaction(func(tx *gorm.DB) error { + var err error + target, err = downgradeUserGroupForSubscriptionTx(tx, sub, now) + return err + })) + assert.Equal(t, "default", target) + var stored User + require.NoError(t, DB.First(&stored, user.Id).Error) + assert.Equal(t, "default", stored.Group) +} + +func TestAdminInvalidateUserSubscriptionUsesExplicitDowngradeWithoutUpgrade(t *testing.T) { + truncateTables(t) + now := GetDBTimestamp() + user := &User{Id: 9833, Username: "subscription-cancel-explicit-downgrade", Group: "svip", Status: common.UserStatusEnabled} + sub := &UserSubscription{Id: 9834, UserId: user.Id, Status: "active", EndTime: now + 3600, DowngradeGroup: "default"} + require.NoError(t, DB.Create(user).Error) + require.NoError(t, DB.Create(sub).Error) + + message, err := AdminInvalidateUserSubscription(sub.Id) + require.NoError(t, err) + assert.Contains(t, message, "default") + var stored User + require.NoError(t, DB.First(&stored, user.Id).Error) + assert.Equal(t, "default", stored.Group) +} + +func TestAdminDeleteUserSubscriptionUsesExplicitDowngradeWithoutUpgrade(t *testing.T) { + truncateTables(t) + now := GetDBTimestamp() + user := &User{Id: 9835, Username: "subscription-delete-explicit-downgrade", Group: "svip", Status: common.UserStatusEnabled} + sub := &UserSubscription{Id: 9836, UserId: user.Id, Status: "active", EndTime: now + 3600, DowngradeGroup: "default"} + require.NoError(t, DB.Create(user).Error) + require.NoError(t, DB.Create(sub).Error) + + message, err := AdminDeleteUserSubscription(sub.Id) + require.NoError(t, err) + assert.Contains(t, message, "default") + var stored User + require.NoError(t, DB.First(&stored, user.Id).Error) + assert.Equal(t, "default", stored.Group) + assert.ErrorIs(t, DB.First(&UserSubscription{}, sub.Id).Error, gorm.ErrRecordNotFound) +} + +func TestExpireDueSubscriptionsUsesExplicitDowngradeWithoutUpgrade(t *testing.T) { + truncateTables(t) + now := GetDBTimestamp() + user := &User{Id: 9837, Username: "subscription-expiry-explicit-downgrade", Group: "svip", Status: common.UserStatusEnabled} + sub := &UserSubscription{Id: 9838, UserId: user.Id, Status: "active", EndTime: now - 1, DowngradeGroup: "default"} + require.NoError(t, DB.Create(user).Error) + require.NoError(t, DB.Create(sub).Error) + + count, err := ExpireDueSubscriptions(10) + require.NoError(t, err) + assert.Equal(t, 1, count) + var stored User + require.NoError(t, DB.First(&stored, user.Id).Error) + assert.Equal(t, "default", stored.Group) + assert.Equal(t, "expired", getSubscriptionResetSub(t, sub.Id).Status) +} + +func TestExpireDueSubscriptionsUsesExplicitDowngradeSnapshot(t *testing.T) { + truncateTables(t) + now := GetDBTimestamp() + user := &User{Id: 9841, Username: "subscription-explicit-expiry", Group: "svip", Status: common.UserStatusEnabled} + expired := &UserSubscription{ + Id: 9842, UserId: user.Id, Status: "active", EndTime: now - 1, + UpgradeGroup: "svip", PrevUserGroup: "vip", DowngradeGroup: "default", + } + require.NoError(t, DB.Create(user).Error) + require.NoError(t, DB.Create(expired).Error) + + count, err := ExpireDueSubscriptions(10) + require.NoError(t, err) + assert.Equal(t, 1, count) + var stored User + require.NoError(t, DB.First(&stored, user.Id).Error) + assert.Equal(t, "default", stored.Group) + assert.Equal(t, "expired", getSubscriptionResetSub(t, expired.Id).Status) +} + +func TestExpireDueSubscriptionsDoesNotReplayHistoricalDowngradePolicy(t *testing.T) { + truncateTables(t) + now := GetDBTimestamp() + user := &User{Id: 9851, Username: "subscription-no-policy-expiry", Group: "svip", Status: common.UserStatusEnabled} + historical := &UserSubscription{ + Id: 9852, UserId: user.Id, Status: "expired", EndTime: now - 3600, + DowngradeGroup: "default", + } + dueWithoutPolicy := &UserSubscription{ + Id: 9853, UserId: user.Id, Status: "active", EndTime: now - 1, + } + require.NoError(t, DB.Create(user).Error) + require.NoError(t, DB.Create(historical).Error) + require.NoError(t, DB.Create(dueWithoutPolicy).Error) + + count, err := ExpireDueSubscriptions(10) + require.NoError(t, err) + assert.Equal(t, 1, count) + var stored User + require.NoError(t, DB.First(&stored, user.Id).Error) + assert.Equal(t, "svip", stored.Group) + assert.Equal(t, "expired", getSubscriptionResetSub(t, dueWithoutPolicy.Id).Status) +} diff --git a/relay/channel/claude/relay-claude.go b/relay/channel/claude/relay-claude.go index 392ee77c..ad80b1a1 100644 --- a/relay/channel/claude/relay-claude.go +++ b/relay/channel/claude/relay-claude.go @@ -645,6 +645,47 @@ type ClaudeResponseInfo struct { ResponseText strings.Builder Usage *dto.Usage Done bool + + pendingToolUses []claudePendingToolUse +} + +type claudePendingToolUse struct { + name string + identity relaycommon.ToolCallIdentity +} + +func (info *ClaudeResponseInfo) rememberToolUse(response *dto.ClaudeResponse) { + if info == nil || response == nil || response.ContentBlock == nil || response.ContentBlock.Type != dto.BuildInCallToolUse { + return + } + info.pendingToolUses = append(info.pendingToolUses, claudePendingToolUse{ + name: response.ContentBlock.Name, + identity: relaycommon.ToolCallIdentity{ + Scope: "claude-messages", + CallID: response.ContentBlock.Id, + Position: fmt.Sprintf("block:%d", response.GetIndex()), + }, + }) +} + +func (info *ClaudeResponseInfo) finishToolUses(relayInfo *relaycommon.RelayInfo, stopReason string) { + if info == nil { + return + } + pending := info.pendingToolUses + info.pendingToolUses = nil + if relayInfo == nil || !strings.EqualFold(strings.TrimSpace(stopReason), "tool_use") { + return + } + for _, toolUse := range pending { + relayInfo.ObserveCustomToolCall(toolUse.name, toolUse.identity) + } +} + +func (info *ClaudeResponseInfo) discardPendingToolUses() { + if info != nil { + info.pendingToolUses = nil + } } func cacheCreationTokensForOpenAIUsage(usage *dto.Usage) int { @@ -862,6 +903,7 @@ func HandleStreamResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud if claudeResponse.Delta != nil && claudeResponse.Delta.StopReason != nil { maybeMarkClaudeRefusal(c, *claudeResponse.Delta.StopReason) } + observeClaudeStreamToolUses(info, claudeInfo, &claudeResponse) if info.RelayFormat == types.RelayFormatClaude { FormatClaudeResponseInfo(&claudeResponse, nil, claudeInfo) @@ -894,6 +936,7 @@ func HandleStreamResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud } func HandleStreamFinalResponse(c *gin.Context, info *relaycommon.RelayInfo, claudeInfo *ClaudeResponseInfo) { + defer claudeInfo.discardPendingToolUses() if claudeInfo.Usage.PromptTokens == 0 { //上游出错 } @@ -942,6 +985,7 @@ func ClaudeStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon. ResponseText: strings.Builder{}, Usage: &dto.Usage{}, } + defer claudeInfo.discardPendingToolUses() var err *types.MaxAPIError helper.StreamScannerHandler(c, resp, info, func(data string, sr *helper.StreamResult) { err = HandleStreamResponseData(c, info, claudeInfo, data) @@ -984,6 +1028,7 @@ func HandleClaudeResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud claudeInfo.Usage.ClaudeCacheCreation5mTokens = claudeResponse.Usage.GetCacheCreation5mTokens() claudeInfo.Usage.ClaudeCacheCreation1hTokens = claudeResponse.Usage.GetCacheCreation1hTokens() } + observeClaudeToolUses(info, &claudeResponse) var responseData []byte var auditResponse *dto.OpenAITextResponse switch info.RelayFormat { @@ -1012,6 +1057,46 @@ func HandleClaudeResponseData(c *gin.Context, info *relaycommon.RelayInfo, claud return nil } +func observeClaudeToolUses(info *relaycommon.RelayInfo, response *dto.ClaudeResponse) { + if info == nil || response == nil { + return + } + if !strings.EqualFold(strings.TrimSpace(response.StopReason), "tool_use") { + return + } + for index := range response.Content { + block := &response.Content[index] + if block.Type != dto.BuildInCallToolUse { + continue + } + info.ObserveCustomToolCall(block.Name, relaycommon.ToolCallIdentity{ + Scope: "claude-messages", + CallID: block.Id, + Position: fmt.Sprintf("block:%d", index), + }) + } +} + +func observeClaudeStreamToolUses(info *relaycommon.RelayInfo, claudeInfo *ClaudeResponseInfo, response *dto.ClaudeResponse) { + if info == nil || claudeInfo == nil || response == nil { + return + } + if response.Type == "content_block_start" { + claudeInfo.rememberToolUse(response) + return + } + if response.Type != "message_delta" { + return + } + stopReason := response.StopReason + if response.Delta != nil && response.Delta.StopReason != nil { + stopReason = *response.Delta.StopReason + } + if strings.TrimSpace(stopReason) != "" { + claudeInfo.finishToolUses(info, stopReason) + } +} + func ClaudeHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*dto.Usage, *types.MaxAPIError) { defer service.CloseResponseBodyGracefully(resp) diff --git a/relay/channel/claude/relay_claude_test.go b/relay/channel/claude/relay_claude_test.go index 814839c4..05a87279 100644 --- a/relay/channel/claude/relay_claude_test.go +++ b/relay/channel/claude/relay_claude_test.go @@ -9,6 +9,8 @@ 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/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" @@ -18,6 +20,27 @@ func commonPointer[T any](value T) *T { return &value } +func setClaudeToolPricesForTest(t *testing.T, additions map[string]float64) { + t.Helper() + setting := config.GlobalConfig.Get("tool_price_setting").(*operation_setting.ToolPriceSetting) + original := make(map[string]float64, len(setting.Prices)) + for name, price := range setting.Prices { + original[name] = price + } + setting.Prices = make(map[string]float64, len(original)+len(additions)) + for name, price := range original { + setting.Prices[name] = price + } + for name, price := range additions { + setting.Prices[name] = price + } + operation_setting.RebuildToolPriceIndex() + t.Cleanup(func() { + setting.Prices = original + operation_setting.RebuildToolPriceIndex() + }) +} + func TestFormatClaudeResponseInfo_MessageStart(t *testing.T) { claudeInfo := &ClaudeResponseInfo{ Usage: &dto.Usage{}, @@ -57,6 +80,132 @@ func TestFormatClaudeResponseInfo_MessageStart(t *testing.T) { } } +func TestHandleClaudeResponseDataRecordsActualCustomToolUse(t *testing.T) { + setClaudeToolPricesForTest(t, map[string]float64{"lookup": 5}) + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + ledger := relaycommon.NewToolUsageLedger("claude-test") + ledger.BeginAttempt(0) + info := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatClaude, + OriginModelName: "claude-test", + ToolUsage: ledger, + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "claude-test"}, + } + response := dto.ClaudeResponse{ + Type: "message", + StopReason: "tool_use", + Content: []dto.ClaudeMediaMessage{{ + Type: "tool_use", + Id: "toolu-1", + Name: "lookup", + Input: map[string]any{}, + }}, + Usage: &dto.ClaudeUsage{InputTokens: 10, OutputTokens: 2}, + } + data, err := common.Marshal(response) + require.NoError(t, err) + + maxErr := HandleClaudeResponseData(c, info, &ClaudeResponseInfo{Usage: &dto.Usage{}}, nil, data) + require.Nil(t, maxErr) + require.True(t, info.CommitToolUsageAttempt()) + require.Equal(t, []relaycommon.ToolUsageItem{{Name: "lookup", CallCount: 1, PricePer1K: 5}}, info.ToolUsageSnapshot().Items) +} + +func TestHandleStreamResponseDataBillsCustomToolUseOnlyAfterToolUseTerminal(t *testing.T) { + setClaudeToolPricesForTest(t, map[string]float64{"lookup": 5}) + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + ledger := relaycommon.NewToolUsageLedger("claude-test") + ledger.BeginAttempt(0) + info := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatClaude, + OriginModelName: "claude-test", + ToolUsage: ledger, + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "claude-test"}, + } + claudeInfo := &ClaudeResponseInfo{Usage: &dto.Usage{}} + startData := `{"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu-1","name":"lookup","input":{}}}` + terminalData := `{"type":"message_delta","delta":{"stop_reason":"tool_use"},"usage":{"output_tokens":2}}` + + require.Nil(t, HandleStreamResponseData(c, info, claudeInfo, startData)) + require.Nil(t, HandleStreamResponseData(c, info, claudeInfo, startData)) + require.Nil(t, HandleStreamResponseData(c, info, claudeInfo, terminalData)) + require.True(t, info.CommitToolUsageAttempt()) + require.Equal(t, []relaycommon.ToolUsageItem{{Name: "lookup", CallCount: 1, PricePer1K: 5}}, info.ToolUsageSnapshot().Items) +} + +func TestHandleStreamResponseDataDoesNotBillIncompleteCustomToolUse(t *testing.T) { + setClaudeToolPricesForTest(t, map[string]float64{"lookup": 5}) + + tests := []struct { + name string + terminalData string + }{ + {name: "max tokens", terminalData: `{"type":"message_delta","delta":{"stop_reason":"max_tokens"},"usage":{"output_tokens":2}}`}, + {name: "refusal", terminalData: `{"type":"message_delta","delta":{"stop_reason":"refusal"},"usage":{"output_tokens":2}}`}, + {name: "missing terminal"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + ledger := relaycommon.NewToolUsageLedger("claude-test") + ledger.BeginAttempt(0) + info := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatClaude, + OriginModelName: "claude-test", + ToolUsage: ledger, + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "claude-test"}, + } + claudeInfo := &ClaudeResponseInfo{Usage: &dto.Usage{}} + + require.Nil(t, HandleStreamResponseData(c, info, claudeInfo, `{"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu-1","name":"lookup","input":{}}}`)) + if tt.terminalData != "" { + require.Nil(t, HandleStreamResponseData(c, info, claudeInfo, tt.terminalData)) + } + HandleStreamFinalResponse(c, info, claudeInfo) + require.True(t, info.CommitToolUsageAttempt()) + require.Empty(t, info.ToolUsageSnapshot().Items) + }) + } +} + +func TestHandleClaudeResponseDataDoesNotBillIncompleteCustomToolUse(t *testing.T) { + setClaudeToolPricesForTest(t, map[string]float64{"lookup": 5}) + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + ledger := relaycommon.NewToolUsageLedger("claude-test") + ledger.BeginAttempt(0) + info := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatClaude, + OriginModelName: "claude-test", + ToolUsage: ledger, + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "claude-test"}, + } + response := dto.ClaudeResponse{ + Type: "message", + StopReason: "max_tokens", + Content: []dto.ClaudeMediaMessage{{ + Type: "tool_use", + Id: "toolu-1", + Name: "lookup", + Input: map[string]any{}, + }}, + Usage: &dto.ClaudeUsage{InputTokens: 10, OutputTokens: 2}, + } + data, err := common.Marshal(response) + require.NoError(t, err) + + maxErr := HandleClaudeResponseData(c, info, &ClaudeResponseInfo{Usage: &dto.Usage{}}, nil, data) + require.Nil(t, maxErr) + require.True(t, info.CommitToolUsageAttempt()) + require.Empty(t, info.ToolUsageSnapshot().Items) +} + func TestFormatClaudeResponseInfo_MessageDelta_FullUsage(t *testing.T) { // message_start 先积累 usage claudeInfo := &ClaudeResponseInfo{ diff --git a/relay/channel/gemini/relay-gemini.go b/relay/channel/gemini/relay-gemini.go index 007776af..74fe499f 100644 --- a/relay/channel/gemini/relay-gemini.go +++ b/relay/channel/gemini/relay-gemini.go @@ -1171,6 +1171,7 @@ func buildUsageFromGeminiResponse(c *gin.Context, info *relaycommon.RelayInfo, r if response == nil { return dto.Usage{} } + observeGeminiFunctionCalls(info, response) if metadata := response.GetUsageMetadata(); dto.HasGeminiUsageMetadataTokens(metadata) { usage := buildUsageFromGeminiMetadata(*metadata, info.GetEstimatePromptTokens()) patchGeminiZeroCompletionUsage(c, info, &usage, geminiResponseUsageText(response), geminiResponseInlineImageCount(response)) @@ -1184,6 +1185,109 @@ func buildUsageFromGeminiResponse(c *gin.Context, info *relaycommon.RelayInfo, r return *usage } +func observeGeminiFunctionCalls(info *relaycommon.RelayInfo, response *dto.GeminiChatResponse) { + newGeminiFunctionCallObserver(info).Observe(response) +} + +type geminiFunctionCallIdentity struct { + name string + callID string +} + +type geminiFunctionCallObserver struct { + info *relaycommon.RelayInfo + pending map[string]geminiFunctionCallIdentity + ambiguous map[string]struct{} +} + +func newGeminiFunctionCallObserver(info *relaycommon.RelayInfo) *geminiFunctionCallObserver { + return &geminiFunctionCallObserver{info: info} +} + +func (o *geminiFunctionCallObserver) Observe(response *dto.GeminiChatResponse) { + if o == nil || o.info == nil || response == nil { + return + } + for _, candidate := range response.Candidates { + for partIndex, part := range candidate.Content.Parts { + if part.FunctionCall == nil { + continue + } + position := fmt.Sprintf("candidate:%d:part:%d", candidate.Index, partIndex) + identity := geminiFunctionCallIdentity{ + name: strings.TrimSpace(part.FunctionCall.FunctionName), + callID: strings.TrimSpace(part.FunctionCall.ID), + } + if part.FunctionCall.WillContinue != nil && *part.FunctionCall.WillContinue { + o.remember(position, identity) + continue + } + if _, blocked := o.ambiguous[position]; blocked { + delete(o.ambiguous, position) + delete(o.pending, position) + continue + } + if pending, ok := o.pending[position]; ok { + delete(o.pending, position) + var compatible bool + identity, compatible = mergeGeminiFunctionCallIdentity(pending, identity) + if !compatible { + continue + } + } + if identity.name == "" { + continue + } + o.info.ObserveCustomToolCall(identity.name, relaycommon.ToolCallIdentity{ + Scope: "gemini", + CallID: identity.callID, + Position: position, + }) + } + } +} + +func (o *geminiFunctionCallObserver) remember(position string, identity geminiFunctionCallIdentity) { + if _, blocked := o.ambiguous[position]; blocked { + return + } + if o.pending == nil { + o.pending = make(map[string]geminiFunctionCallIdentity) + } + if pending, ok := o.pending[position]; ok { + merged, compatible := mergeGeminiFunctionCallIdentity(pending, identity) + if !compatible { + delete(o.pending, position) + if o.ambiguous == nil { + o.ambiguous = make(map[string]struct{}) + } + o.ambiguous[position] = struct{}{} + return + } + o.pending[position] = merged + return + } + if identity.name != "" || identity.callID != "" { + o.pending[position] = identity + } +} + +func mergeGeminiFunctionCallIdentity(first, second geminiFunctionCallIdentity) (geminiFunctionCallIdentity, bool) { + if first.name != "" && second.name != "" && first.name != second.name { + return geminiFunctionCallIdentity{}, false + } + if first.callID != "" && second.callID != "" && first.callID != second.callID { + return geminiFunctionCallIdentity{}, false + } + if first.name == "" { + first.name = second.name + } + if first.callID == "" { + first.callID = second.callID + } + return first, true +} + func responseGeminiChat2OpenAI(c *gin.Context, response *dto.GeminiChatResponse) *dto.OpenAITextResponse { fullTextResponse := dto.OpenAITextResponse{ Id: helper.GetResponseID(c), @@ -1469,6 +1573,7 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http var imageCount int responseText := strings.Builder{} auditText := strings.Builder{} + functionCallObserver := newGeminiFunctionCallObserver(info) helper.StreamScannerHandler(c, resp, info, func(data string, sr *helper.StreamResult) { var geminiResponse dto.GeminiChatResponse @@ -1489,6 +1594,7 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http if len(geminiResponse.Candidates) == 0 && geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil { common.SetContextKey(c, constant.ContextKeyAdminRejectReason, fmt.Sprintf("gemini_block_reason=%s", *geminiResponse.PromptFeedback.BlockReason)) } + functionCallObserver.Observe(&geminiResponse) // 统计图片数量 for _, candidate := range geminiResponse.Candidates { diff --git a/relay/channel/gemini/relay_gemini_usage_test.go b/relay/channel/gemini/relay_gemini_usage_test.go index b9a72307..7d810bf1 100644 --- a/relay/channel/gemini/relay_gemini_usage_test.go +++ b/relay/channel/gemini/relay_gemini_usage_test.go @@ -2,6 +2,7 @@ package gemini import ( "bytes" + "context" "io" "net/http" "net/http/httptest" @@ -11,11 +12,163 @@ 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/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" ) +func setGeminiToolPricesForTest(t *testing.T, additions map[string]float64) { + t.Helper() + + setting := config.GlobalConfig.Get("tool_price_setting").(*operation_setting.ToolPriceSetting) + original := make(map[string]float64, len(setting.Prices)) + for name, price := range setting.Prices { + original[name] = price + } + setting.Prices = make(map[string]float64, len(original)+len(additions)) + for name, price := range original { + setting.Prices[name] = price + } + for name, price := range additions { + setting.Prices[name] = price + } + operation_setting.RebuildToolPriceIndex() + + t.Cleanup(func() { + setting.Prices = original + operation_setting.RebuildToolPriceIndex() + }) +} + +func newGeminiToolBillingContext(path string) (*gin.Context, *relaycommon.RelayInfo) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequestWithContext(context.Background(), http.MethodPost, path, nil) + ledger := relaycommon.NewToolUsageLedger("gemini-test") + ledger.BeginAttempt(0) + return c, &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatOpenAI, + OriginModelName: "gemini-test", + DisablePing: true, + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: "gemini-test", + }, + ToolUsage: ledger, + } +} + +func TestGeminiChatHandlerRecordsCompletedCustomFunctionCalls(t *testing.T) { + setGeminiToolPricesForTest(t, map[string]float64{"lookup": 5}) + c, info := newGeminiToolBillingContext("/v1/chat/completions") + + payload := `{ + "candidates":[{"index":0,"content":{"role":"model","parts":[ + {"functionCall":{"id":"call-1","name":"lookup","args":{"q":"max"}}}, + {"functionCall":{"id":"call-2","name":"unpriced","args":{}}} + ]}}], + "usageMetadata":{"promptTokenCount":1,"candidatesTokenCount":1,"totalTokenCount":2} + }` + resp := &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(payload))} + + _, maxErr := GeminiChatHandler(c, info, resp) + require.Nil(t, maxErr) + require.True(t, info.CommitToolUsageAttempt()) + require.Equal(t, []relaycommon.ToolUsageItem{{ + Name: "lookup", + CallCount: 1, + PricePer1K: 5, + }}, info.ToolUsageSnapshot().Items) +} + +func TestGeminiStreamHandlerWaitsForCompletedFunctionCallAndDeduplicatesID(t *testing.T) { + setGeminiToolPricesForTest(t, map[string]float64{"lookup": 5}) + c, info := newGeminiToolBillingContext("/v1/chat/completions") + + body := strings.Join([]string{ + `data: {"candidates":[{"index":0,"content":{"role":"model","parts":[{"functionCall":{"id":"call-1","name":"lookup","args":{"q":"max"},"willContinue":true}}]}}]}`, + `data: {"candidates":[{"index":0,"content":{"role":"model","parts":[{"functionCall":{"id":"call-1","name":"lookup","args":{"q":"max"},"willContinue":false}}]}}],"usageMetadata":{"promptTokenCount":1,"candidatesTokenCount":1,"totalTokenCount":2}}`, + `data: {"candidates":[{"index":0,"content":{"role":"model","parts":[{"functionCall":{"id":"call-1","name":"lookup","args":{"q":"max"},"willContinue":false}}]}}]}`, + `data: [DONE]`, + "", + }, "\n") + resp := &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(body))} + + _, maxErr := geminiStreamHandler(c, info, resp, func(_ string, _ *dto.GeminiChatResponse) bool { + return true + }) + require.Nil(t, maxErr) + require.True(t, info.CommitToolUsageAttempt()) + require.Equal(t, []relaycommon.ToolUsageItem{{ + Name: "lookup", + CallCount: 1, + PricePer1K: 5, + }}, info.ToolUsageSnapshot().Items) +} + +func TestGeminiStreamHandlerCarriesFunctionIdentityToEmptyCompletionChunk(t *testing.T) { + setGeminiToolPricesForTest(t, map[string]float64{"lookup": 5}) + c, info := newGeminiToolBillingContext("/v1/chat/completions") + + body := strings.Join([]string{ + `data: {"candidates":[{"index":0,"content":{"role":"model","parts":[{"functionCall":{"id":"call-1","name":"lookup","args":{"q":"max"},"willContinue":true}}]}}]}`, + `data: {"candidates":[{"index":0,"content":{"role":"model","parts":[{"functionCall":{"willContinue":false}}]}}],"usageMetadata":{"promptTokenCount":1,"candidatesTokenCount":1,"totalTokenCount":2}}`, + `data: [DONE]`, + "", + }, "\n") + resp := &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(body))} + + _, maxErr := geminiStreamHandler(c, info, resp, func(_ string, _ *dto.GeminiChatResponse) bool { + return true + }) + require.Nil(t, maxErr) + require.True(t, info.CommitToolUsageAttempt()) + require.Equal(t, []relaycommon.ToolUsageItem{{ + Name: "lookup", + CallCount: 1, + PricePer1K: 5, + }}, info.ToolUsageSnapshot().Items) +} + +func TestGeminiStreamHandlerDoesNotBillUnterminatedFunctionCall(t *testing.T) { + setGeminiToolPricesForTest(t, map[string]float64{"lookup": 5}) + c, info := newGeminiToolBillingContext("/v1/chat/completions") + + body := strings.Join([]string{ + `data: {"candidates":[{"index":0,"content":{"role":"model","parts":[{"functionCall":{"id":"call-1","name":"lookup","args":{"q":"max"},"willContinue":true}}]}}]}`, + `data: [DONE]`, + "", + }, "\n") + resp := &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(body))} + + _, maxErr := geminiStreamHandler(c, info, resp, func(_ string, _ *dto.GeminiChatResponse) bool { + return true + }) + require.Nil(t, maxErr) + require.True(t, info.CommitToolUsageAttempt()) + require.Empty(t, info.ToolUsageSnapshot().Items) +} + +func TestGeminiStreamHandlerIgnoresEmptyCompletionWithoutPendingIdentity(t *testing.T) { + setGeminiToolPricesForTest(t, map[string]float64{"lookup": 5}) + c, info := newGeminiToolBillingContext("/v1/chat/completions") + + body := strings.Join([]string{ + `data: {"candidates":[{"index":0,"content":{"role":"model","parts":[{"functionCall":{"willContinue":false}}]}}],"usageMetadata":{"promptTokenCount":1,"candidatesTokenCount":1,"totalTokenCount":2}}`, + `data: [DONE]`, + "", + }, "\n") + resp := &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(body))} + + _, maxErr := geminiStreamHandler(c, info, resp, func(_ string, _ *dto.GeminiChatResponse) bool { + return true + }) + require.Nil(t, maxErr) + require.True(t, info.CommitToolUsageAttempt()) + require.Empty(t, info.ToolUsageSnapshot().Items) +} + func TestGeminiChatHandlerCompletionTokensExcludeToolUsePromptTokens(t *testing.T) { t.Parallel() diff --git a/relay/channel/openai/chat_via_responses.go b/relay/channel/openai/chat_via_responses.go index 2967f3a0..abfc7adc 100644 --- a/relay/channel/openai/chat_via_responses.go +++ b/relay/channel/openai/chat_via_responses.go @@ -56,9 +56,13 @@ func OaiResponsesToChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) } + if terminalErr := responsesTerminalError(&responsesResp, resp.StatusCode, false); terminalErr != nil { + return nil, terminalErr + } if oaiError := responsesResp.GetOpenAIError(); oaiError != nil && oaiError.Type != "" { return nil, types.WithOpenAIError(*oaiError, resp.StatusCode) } + observeResponsesOutputs(info, responsesResp.Output) chatId := helper.GetResponseID(c) chatResp, usage, err := service.ResponsesResponseToChatCompletionsResponse(&responsesResp, chatId) @@ -148,6 +152,7 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo return true } 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 @@ -254,10 +259,6 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo if callID == "" { return true } - if outputText.Len() > 0 { - // Prefer streaming assistant text over tool calls to match non-stream behavior. - return true - } if !sendStartIfNeeded() { return false } @@ -420,9 +421,12 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo if streamResp.Item == nil { break } - if streamResp.Item.Type != "function_call" { + if streamResp.Item.Type != dto.BuildInCallFunctionCall && streamResp.Item.Type != dto.BuildInCallCustomToolCall { break } + if streamResp.Type == dto.ResponsesOutputTypeItemDone { + observeResponsesToolOutput(info, streamResp.Item, streamResp.OutputIndex) + } itemID := strings.TrimSpace(streamResp.Item.ID) callID := strings.TrimSpace(streamResp.Item.CallId) @@ -454,7 +458,7 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo return } - case "response.function_call_arguments.delta": + case "response.function_call_arguments.delta", "response.custom_tool_call_input.delta": itemID := strings.TrimSpace(streamResp.ItemID) callID := toolCallCanonicalIDByItemID[itemID] if callID == "" { @@ -469,11 +473,17 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo return } - case "response.function_call_arguments.done": + case "response.function_call_arguments.done", "response.custom_tool_call_input.done": case "response.completed": completedOutputTypes := []string(nil) if streamResp.Response != nil { + if terminalErr := responsesTerminalError(streamResp.Response, http.StatusInternalServerError, info.SendResponseCount > 0); terminalErr != nil { + streamErr = terminalErr + sr.Stop(streamErr) + return + } + observeResponsesOutputs(info, streamResp.Response.Output) completedOutputTypes = responsesOutputTypes(streamResp.Response) if streamResp.Response.Model != "" { model = streamResp.Response.Model @@ -515,20 +525,21 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo hasVisiblePayload = true } } - if !sawToolCall { - for _, out := range streamResp.Response.Output { - if out.Type != "function_call" { - continue - } - itemID := strings.TrimSpace(out.ID) - callID := strings.TrimSpace(out.CallId) - if callID == "" { - callID = itemID - } - if !sendToolCallDelta(callID, strings.TrimSpace(out.Name), out.ArgumentsString()) { - sr.Stop(streamErr) - return - } + for _, out := range streamResp.Response.Output { + if out.Type != dto.BuildInCallFunctionCall && out.Type != dto.BuildInCallCustomToolCall { + continue + } + itemID := strings.TrimSpace(out.ID) + callID := strings.TrimSpace(out.CallId) + if callID == "" { + callID = itemID + } + if _, alreadyEmitted := toolCallIndexByID[callID]; alreadyEmitted { + continue + } + if !sendToolCallDelta(callID, strings.TrimSpace(out.Name), out.ArgumentsString()) { + sr.Stop(streamErr) + return } } if streamResp.Response.HasImageGenerationCall() { @@ -554,7 +565,7 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo info.ClaudeConvertInfo.Usage = usage } finishReason := "stop" - if sawToolCall && outputText.Len() == 0 { + if sawToolCall { finishReason = "tool_calls" } stop := helper.GenerateStopResponse(responseId, createAt, model, finishReason) @@ -565,15 +576,9 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo sentStop = true } - case "response.error", "response.failed": - if streamResp.Response != nil { - if oaiErr := streamResp.Response.GetOpenAIError(); oaiErr != nil && oaiErr.Type != "" { - streamErr = types.WithOpenAIError(*oaiErr, http.StatusInternalServerError) - sr.Stop(streamErr) - return - } - } - streamErr = types.NewOpenAIError(fmt.Errorf("responses stream error: %s", streamResp.Type), types.ErrorCodeBadResponse, http.StatusInternalServerError) + case "error", "response.error", "response.failed", "response.cancelled", "response.canceled": + skipRetry := info.SendResponseCount > 0 + streamErr = responsesStreamTerminalError(&streamResp, skipRetry) sr.Stop(streamErr) return @@ -605,7 +610,7 @@ func OaiResponsesToChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo info.ClaudeConvertInfo.Usage = usage } finishReason := "stop" - if sawToolCall && outputText.Len() == 0 { + if sawToolCall { finishReason = "tool_calls" } stop := helper.GenerateStopResponse(responseId, createAt, model, finishReason) diff --git a/relay/channel/openai/relay-openai.go b/relay/channel/openai/relay-openai.go index 0f19a4c1..156923ea 100644 --- a/relay/channel/openai/relay-openai.go +++ b/relay/channel/openai/relay-openai.go @@ -122,6 +122,7 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re var usage = &dto.Usage{} var lastStreamData string var secondLastStreamData string // 存储倒数第二个stream data,用于音频模型 + toolCallObserver := newOpenAIStreamToolCallObserver(info) // 检查是否为音频模型 isAudioModel := strings.Contains(strings.ToLower(model), "audio") @@ -140,6 +141,7 @@ func OaiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Re } lastStreamData = data + observeOpenAIStreamToolCalls(toolCallObserver, data) if err := processTokenData(info.RelayMode, data, &responseTextBuilder, &toolCount); err != nil { logger.LogError(c, "error processing stream token data: "+err.Error()) sr.Error(err) @@ -228,6 +230,7 @@ func OpenaiHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respo if service.ResponseAuditEnabled() { service.SetRelayResponseAuditContent(info, service.BuildTextResponseAuditContent(&simpleResponse)) } + observeOpenAITextToolCalls(info, simpleResponse.Choices) for _, choice := range simpleResponse.Choices { if choice.FinishReason == constant.FinishReasonContentFilter { @@ -300,6 +303,164 @@ func OpenaiHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respo return &simpleResponse.Usage, nil } +func observeOpenAITextToolCalls(info *relaycommon.RelayInfo, choices []dto.OpenAITextResponseChoice) { + if info == nil { + return + } + for _, choice := range choices { + if choice.FinishReason != constant.FinishReasonToolCalls { + continue + } + for toolIndex, toolCall := range choice.Message.ParseToolCalls() { + info.ObserveCustomToolCall(toolCall.Function.Name, relaycommon.ToolCallIdentity{ + Scope: "openai-chat", + CallID: toolCall.ID, + Position: fmt.Sprintf("choice:%d:tool:%d", choice.Index, toolIndex), + }) + } + } +} + +type openAIStreamToolCallKey struct { + choiceIndex int + toolIndex int +} + +type openAIStreamToolCallIdentity struct { + name string + callID string +} + +type openAIStreamToolCallObserver struct { + info *relaycommon.RelayInfo + pending map[openAIStreamToolCallKey]openAIStreamToolCallIdentity + ambiguous map[openAIStreamToolCallKey]struct{} +} + +func newOpenAIStreamToolCallObserver(info *relaycommon.RelayInfo) *openAIStreamToolCallObserver { + return &openAIStreamToolCallObserver{info: info} +} + +func observeOpenAIStreamToolCalls(observer *openAIStreamToolCallObserver, data string) { + if observer == nil || data == "" || !openAIStreamChunkMayAffectToolBilling(data) { + return + } + var streamResponse dto.ChatCompletionsStreamResponse + if err := common.UnmarshalJsonStr(data, &streamResponse); err != nil { + return + } + observer.Observe(&streamResponse) +} + +func openAIStreamChunkMayAffectToolBilling(data string) bool { + if strings.Contains(data, `"tool_calls"`) { + return true + } + const finishReasonKey = `"finish_reason"` + for searchFrom := 0; searchFrom < len(data); { + keyOffset := strings.Index(data[searchFrom:], finishReasonKey) + if keyOffset < 0 { + return false + } + valueOffset := searchFrom + keyOffset + len(finishReasonKey) + for valueOffset < len(data) && (data[valueOffset] == ' ' || data[valueOffset] == '\t' || data[valueOffset] == '\r' || data[valueOffset] == '\n' || data[valueOffset] == ':') { + valueOffset++ + } + if valueOffset < len(data) && !strings.HasPrefix(data[valueOffset:], "null") { + return true + } + searchFrom = valueOffset + len("null") + } + return false +} + +func (o *openAIStreamToolCallObserver) Observe(streamResponse *dto.ChatCompletionsStreamResponse) { + if o == nil || o.info == nil || streamResponse == nil { + return + } + for _, choice := range streamResponse.Choices { + for position, toolCall := range choice.Delta.ToolCalls { + toolIndex := position + if toolCall.Index != nil { + toolIndex = *toolCall.Index + } + o.remember(openAIStreamToolCallKey{choiceIndex: choice.Index, toolIndex: toolIndex}, openAIStreamToolCallIdentity{ + name: strings.TrimSpace(toolCall.Function.Name), + callID: strings.TrimSpace(toolCall.ID), + }) + } + if choice.FinishReason != nil && strings.TrimSpace(*choice.FinishReason) != "" { + o.finishChoice(choice.Index, strings.TrimSpace(*choice.FinishReason)) + } + } +} + +func (o *openAIStreamToolCallObserver) remember(key openAIStreamToolCallKey, identity openAIStreamToolCallIdentity) { + if _, blocked := o.ambiguous[key]; blocked { + return + } + if o.pending == nil { + o.pending = make(map[openAIStreamToolCallKey]openAIStreamToolCallIdentity) + } + if pending, ok := o.pending[key]; ok { + merged, compatible := mergeOpenAIStreamToolCallIdentity(pending, identity) + if !compatible { + delete(o.pending, key) + if o.ambiguous == nil { + o.ambiguous = make(map[openAIStreamToolCallKey]struct{}) + } + o.ambiguous[key] = struct{}{} + return + } + o.pending[key] = merged + return + } + if identity.name != "" || identity.callID != "" { + o.pending[key] = identity + } +} + +func (o *openAIStreamToolCallObserver) finishChoice(choiceIndex int, finishReason string) { + completed := finishReason == constant.FinishReasonToolCalls + for key, identity := range o.pending { + if key.choiceIndex != choiceIndex { + continue + } + delete(o.pending, key) + _, blocked := o.ambiguous[key] + delete(o.ambiguous, key) + if !completed || blocked || identity.name == "" { + continue + } + o.info.ObserveCustomToolCall(identity.name, relaycommon.ToolCallIdentity{ + Scope: "openai-chat", + CallID: identity.callID, + Position: fmt.Sprintf("choice:%d:tool:%d", key.choiceIndex, key.toolIndex), + }) + } + for key := range o.ambiguous { + if key.choiceIndex == choiceIndex { + delete(o.ambiguous, key) + } + } +} + +func mergeOpenAIStreamToolCallIdentity(first, second openAIStreamToolCallIdentity) (openAIStreamToolCallIdentity, bool) { + if first.name != "" && second.name != "" && first.name != second.name { + return openAIStreamToolCallIdentity{}, false + } + if first.callID != "" && second.callID != "" && first.callID != second.callID { + return openAIStreamToolCallIdentity{}, false + } + if first.name == "" { + first.name = second.name + } + if first.callID == "" { + first.callID = second.callID + } + return first, true +} + func streamTTSResponse(c *gin.Context, resp *http.Response) { c.Writer.WriteHeaderNow() diff --git a/relay/channel/openai/relay_responses.go b/relay/channel/openai/relay_responses.go index 8ec2ffab..4c4a34ba 100644 --- a/relay/channel/openai/relay_responses.go +++ b/relay/channel/openai/relay_responses.go @@ -56,6 +56,81 @@ func applyResponsesUsage(dst *dto.Usage, src *dto.Usage) { } } +func normalizedResponsesStatus(response *dto.OpenAIResponsesResponse) string { + if response == nil || len(response.Status) == 0 { + return "" + } + var status string + if err := common.Unmarshal(response.Status, &status); err != nil { + return "" + } + return strings.ToLower(strings.TrimSpace(status)) +} + +func responsesTerminalError(response *dto.OpenAIResponsesResponse, statusCode int, skipRetry bool) *types.MaxAPIError { + status := normalizedResponsesStatus(response) + switch status { + case "failed", "cancelled", "canceled": + default: + return nil + } + + if statusCode < http.StatusBadRequest { + statusCode = http.StatusInternalServerError + } + if response != nil { + if oaiErr := response.GetOpenAIError(); oaiErr != nil && (oaiErr.Type != "" || oaiErr.Message != "") { + if skipRetry { + return types.WithOpenAIError(*oaiErr, statusCode, types.ErrOptionWithSkipRetry()) + } + return types.WithOpenAIError(*oaiErr, statusCode) + } + } + err := fmt.Errorf("responses request terminated with status %s", status) + if skipRetry { + return types.NewOpenAIError(err, types.ErrorCodeBadResponse, statusCode, types.ErrOptionWithSkipRetry()) + } + return types.NewOpenAIError(err, types.ErrorCodeBadResponse, statusCode) +} + +func responsesStreamTerminalError(streamResponse *dto.ResponsesStreamResponse, skipRetry bool) *types.MaxAPIError { + if streamResponse != nil && streamResponse.Type == "error" { + message := streamResponse.Message + if strings.TrimSpace(message) == "" { + message = "responses stream error: error" + } + oaiErr := types.OpenAIError{ + Message: message, + Param: streamResponse.Param, + Code: streamResponse.Code, + } + if skipRetry { + return types.WithOpenAIError(oaiErr, http.StatusInternalServerError, types.ErrOptionWithSkipRetry()) + } + return types.WithOpenAIError(oaiErr, http.StatusInternalServerError) + } + if streamResponse != nil && streamResponse.Response != nil { + if terminalErr := responsesTerminalError(streamResponse.Response, http.StatusInternalServerError, skipRetry); terminalErr != nil { + return terminalErr + } + if oaiErr := streamResponse.Response.GetOpenAIError(); oaiErr != nil && (oaiErr.Type != "" || oaiErr.Message != "") { + if skipRetry { + return types.WithOpenAIError(*oaiErr, http.StatusInternalServerError, types.ErrOptionWithSkipRetry()) + } + return types.WithOpenAIError(*oaiErr, http.StatusInternalServerError) + } + } + eventType := "response.error" + if streamResponse != nil && strings.TrimSpace(streamResponse.Type) != "" { + eventType = streamResponse.Type + } + err := fmt.Errorf("responses stream error: %s", eventType) + if skipRetry { + return types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError, types.ErrOptionWithSkipRetry()) + } + return types.NewOpenAIError(err, types.ErrorCodeBadResponse, http.StatusInternalServerError) +} + func OaiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.MaxAPIError) { defer service.CloseResponseBodyGracefully(resp) @@ -69,6 +144,9 @@ func OaiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http if err != nil { return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) } + if terminalErr := responsesTerminalError(&responsesResponse, resp.StatusCode, false); terminalErr != nil { + return nil, terminalErr + } if oaiError := responsesResponse.GetOpenAIError(); oaiError != nil && oaiError.Type != "" { return nil, types.WithOpenAIError(*oaiError, resp.StatusCode) } @@ -99,19 +177,63 @@ func OaiResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http service.IOCopyBytesGracefully(c, resp, responseBody) - if info == nil || info.ResponsesUsageInfo == nil || info.ResponsesUsageInfo.BuiltInTools == nil { - return &usage, nil + observeResponsesOutputs(info, responsesResponse.Output) + return &usage, nil +} + +func observeResponsesOutputs(info *relaycommon.RelayInfo, outputs []dto.ResponsesOutput) { + for outputIndex := range outputs { + observeResponsesToolOutput(info, &outputs[outputIndex], &outputIndex) + } +} + +func observeResponsesToolOutput(info *relaycommon.RelayInfo, output *dto.ResponsesOutput, outputIndex *int) { + if info == nil || output == nil || !isBillableResponsesToolStatus(output.Status) { + return + } + callID := strings.TrimSpace(output.CallId) + if callID == "" { + callID = strings.TrimSpace(output.ID) + } + position := "" + if outputIndex != nil && *outputIndex >= 0 { + position = fmt.Sprintf("output:%d", *outputIndex) + } + identity := relaycommon.ToolCallIdentity{ + Scope: "openai-responses", + CallID: callID, + Position: position, + } + + switch output.Type { + case dto.BuildInCallWebSearchCall: + info.ObserveBuiltInToolCall(responsesWebSearchToolName(info), identity) + case dto.BuildInCallFileSearchCall: + info.ObserveBuiltInToolCall(dto.BuildInToolFileSearch, identity) + case dto.BuildInCallFunctionCall, dto.BuildInCallCustomToolCall, dto.BuildInCallToolUse: + info.ObserveCustomToolCall(output.Name, identity) + } +} + +func isBillableResponsesToolStatus(status string) bool { + switch strings.ToLower(strings.TrimSpace(status)) { + case "", "completed": + return true + default: + return false } - // 解析 Tools 用量 - for _, tool := range responsesResponse.Tools { - buildToolinfo, ok := info.ResponsesUsageInfo.BuiltInTools[common.Interface2String(tool["type"])] - if !ok || buildToolinfo == nil { - logger.LogError(c, fmt.Sprintf("BuiltInTools not found for tool type: %v", tool["type"])) - continue +} + +func responsesWebSearchToolName(info *relaycommon.RelayInfo) string { + if info != nil && info.ResponsesUsageInfo != nil { + if _, ok := info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolWebSearchPreview]; ok { + return dto.BuildInToolWebSearchPreview + } + if _, ok := info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolWebSearch]; ok { + return dto.BuildInToolWebSearch } - buildToolinfo.CallCount++ } - return &usage, nil + return dto.BuildInToolWebSearchPreview } func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.MaxAPIError) { @@ -161,10 +283,16 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp if responsesStreamHasVisibleOutput(&streamResponse) { hasVisiblePayload = true } + var terminalEventErr *types.MaxAPIError switch streamResponse.Type { case "response.completed": 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) c.Set("image_generation_call_quality", streamResponse.Response.GetQuality()) @@ -186,14 +314,21 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp // 函数调用处理 if streamResponse.Item != nil { switch streamResponse.Item.Type { - case dto.BuildInCallWebSearchCall: - if info != nil && info.ResponsesUsageInfo != nil && info.ResponsesUsageInfo.BuiltInTools != nil { - if webSearchTool, exists := info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolWebSearchPreview]; exists && webSearchTool != nil { - webSearchTool.CallCount++ - } - } + case dto.BuildInCallWebSearchCall, dto.BuildInCallFileSearchCall, dto.BuildInCallFunctionCall, dto.BuildInCallCustomToolCall, dto.BuildInCallToolUse: + observeResponsesToolOutput(info, streamResponse.Item, streamResponse.OutputIndex) } } + case "error", "response.error", "response.failed", "response.cancelled", "response.canceled": + terminalEventErr = responsesStreamTerminalError(&streamResponse, streamForwarded) + } + + if terminalEventErr != nil { + if streamForwarded { + sendResponsesStreamData(c, streamResponse, data) + } + streamErr = terminalEventErr + sr.Stop(streamErr) + return } if streamForwarded { @@ -206,7 +341,7 @@ func OaiResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp data: data, }) switch streamResponse.Type { - case "response.completed", "response.error", "response.failed": + case "response.completed": flushPendingEvents() default: if hasVisiblePayload || len(pendingEvents) >= maxPendingResponsesStreamEvents || !shouldBufferForEmptyRetry() { diff --git a/relay/channel/openai/responses_usage_test.go b/relay/channel/openai/responses_usage_test.go index 3fe5fd1c..4d55b77e 100644 --- a/relay/channel/openai/responses_usage_test.go +++ b/relay/channel/openai/responses_usage_test.go @@ -5,11 +5,13 @@ 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/types" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) @@ -145,3 +147,446 @@ func TestOaiResponsesCompactionHandlerPreservesDetailedUsage(t *testing.T) { require.Nil(t, maxAPIError) assertDetailedResponsesUsage(t, usage) } + +func TestOaiResponsesHandlerBillsActualFunctionOutputNotToolDeclaration(t *testing.T) { + setOpenAIToolPricesForTest(t, map[string]float64{"lookup": 5}) + c, info := newOpenAIToolBillingContext("gpt-test") + info.ResponsesUsageInfo = &relaycommon.ResponsesUsageInfo{ + BuiltInTools: map[string]*relaycommon.BuildInToolInfo{ + "function": {ToolName: "function"}, + }, + } + payload := dto.OpenAIResponsesResponse{ + Tools: []map[string]any{{"type": "function", "name": "lookup"}}, + Output: []dto.ResponsesOutput{{ + Type: "function_call", + ID: "fc-1", + CallId: "call-1", + Name: "lookup", + Status: "completed", + }}, + Usage: detailedResponsesUsage(), + } + body, err := common.Marshal(payload) + require.NoError(t, err) + + _, maxAPIError := OaiResponsesHandler(c, info, &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(bytes.NewReader(body)), + }) + require.Nil(t, maxAPIError) + require.True(t, info.CommitToolUsageAttempt()) + require.Equal(t, []relaycommon.ToolUsageItem{{ + Name: "lookup", + CallCount: 1, + PricePer1K: 5, + }}, info.ToolUsageSnapshot().Items) + require.Zero(t, info.ResponsesUsageInfo.BuiltInTools["function"].CallCount) +} + +func TestOaiResponsesHandlerDoesNotBillDeclarationWithoutOutputCall(t *testing.T) { + setOpenAIToolPricesForTest(t, map[string]float64{"lookup": 5}) + c, info := newOpenAIToolBillingContext("gpt-test") + info.ResponsesUsageInfo = &relaycommon.ResponsesUsageInfo{ + BuiltInTools: map[string]*relaycommon.BuildInToolInfo{ + "function": {ToolName: "function"}, + }, + } + payload := dto.OpenAIResponsesResponse{ + Tools: []map[string]any{{"type": "function", "name": "lookup"}}, + Output: []dto.ResponsesOutput{{ + Type: "message", + Role: "assistant", + Content: []dto.ResponsesOutputContent{{ + Type: "output_text", + Text: "no tool needed", + }}, + }}, + Usage: detailedResponsesUsage(), + } + body, err := common.Marshal(payload) + require.NoError(t, err) + + _, maxAPIError := OaiResponsesHandler(c, info, &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(bytes.NewReader(body)), + }) + require.Nil(t, maxAPIError) + require.True(t, info.CommitToolUsageAttempt()) + require.Empty(t, info.ToolUsageSnapshot().Items) + require.Zero(t, info.ResponsesUsageInfo.BuiltInTools["function"].CallCount) +} + +func TestOaiResponsesHandlerBillsOnlyKnownSuccessfulToolStatuses(t *testing.T) { + setOpenAIToolPricesForTest(t, map[string]float64{"lookup": 5}) + + tests := []struct { + name string + status string + billable bool + }{ + {name: "empty status", billable: true}, + {name: "completed", status: "completed", billable: true}, + {name: "in progress", status: "in_progress"}, + {name: "queued", status: "queued"}, + {name: "future status", status: "future_status"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + c, info := newOpenAIToolBillingContext("gpt-test") + payload := dto.OpenAIResponsesResponse{ + Output: []dto.ResponsesOutput{{ + Type: dto.BuildInCallFunctionCall, + ID: "fc-1", + CallId: "call-1", + Name: "lookup", + Status: tt.status, + }}, + Usage: detailedResponsesUsage(), + } + body, err := common.Marshal(payload) + require.NoError(t, err) + + _, maxAPIError := OaiResponsesHandler(c, info, &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(bytes.NewReader(body)), + }) + require.Nil(t, maxAPIError) + require.True(t, info.CommitToolUsageAttempt()) + if tt.billable { + require.Equal(t, []relaycommon.ToolUsageItem{{Name: "lookup", CallCount: 1, PricePer1K: 5}}, info.ToolUsageSnapshot().Items) + } else { + require.Empty(t, info.ToolUsageSnapshot().Items) + } + }) + } +} + +func TestOaiResponsesHandlerRecordsCustomToolCallOutput(t *testing.T) { + setOpenAIToolPricesForTest(t, map[string]float64{"code_exec": 6}) + c, info := newOpenAIToolBillingContext("gpt-test") + payload := dto.OpenAIResponsesResponse{ + Status: []byte(`"completed"`), + Output: []dto.ResponsesOutput{{ + Type: dto.BuildInCallCustomToolCall, + ID: "ct-1", + CallId: "call-1", + Name: "code_exec", + Status: "completed", + }}, + Usage: &dto.Usage{InputTokens: 1, OutputTokens: 1, TotalTokens: 2}, + } + body, err := common.Marshal(payload) + require.NoError(t, err) + + _, maxAPIError := OaiResponsesHandler(c, info, &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(bytes.NewReader(body)), + }) + require.Nil(t, maxAPIError) + require.True(t, info.CommitToolUsageAttempt()) + require.Equal(t, []relaycommon.ToolUsageItem{{ + Name: "code_exec", + CallCount: 1, + PricePer1K: 6, + }}, info.ToolUsageSnapshot().Items) +} + +func TestOaiResponsesStreamHandlerDeduplicatesFunctionItemDoneAndCompleted(t *testing.T) { + setOpenAIToolPricesForTest(t, map[string]float64{"lookup": 5}) + c, info := newOpenAIToolBillingContext("gpt-test") + + itemDone := dto.ResponsesStreamResponse{ + Type: dto.ResponsesOutputTypeItemDone, + OutputIndex: common.GetPointer(0), + Item: &dto.ResponsesOutput{ + Type: "function_call", + ID: "fc-1", + CallId: "call-1", + Name: "lookup", + Status: "completed", + }, + } + completed := dto.ResponsesStreamResponse{ + Type: "response.completed", + Response: &dto.OpenAIResponsesResponse{ + Output: []dto.ResponsesOutput{{ + Type: "function_call", + ID: "fc-1", + CallId: "call-1", + Name: "lookup", + Status: "completed", + }}, + Usage: detailedResponsesUsage(), + }, + } + itemData, err := common.Marshal(itemDone) + require.NoError(t, err) + completedData, err := common.Marshal(completed) + require.NoError(t, err) + body := append([]byte("data: "), itemData...) + body = append(body, []byte("\ndata: ")...) + body = append(body, itemData...) + body = append(body, []byte("\ndata: ")...) + body = append(body, completedData...) + body = append(body, []byte("\ndata: [DONE]\n")...) + + _, maxAPIError := OaiResponsesStreamHandler(c, info, &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(bytes.NewReader(body)), + }) + require.Nil(t, maxAPIError) + require.True(t, info.CommitToolUsageAttempt()) + require.Equal(t, []relaycommon.ToolUsageItem{{ + Name: "lookup", + CallCount: 1, + PricePer1K: 5, + }}, info.ToolUsageSnapshot().Items) +} + +func TestOaiResponsesStreamHandlerCountsBuiltInToolsOnce(t *testing.T) { + setOpenAIToolPricesForTest(t, map[string]float64{ + dto.BuildInToolWebSearchPreview: 5, + dto.BuildInToolFileSearch: 7, + }) + + tests := []struct { + name string + callType string + toolName string + }{ + {name: "web search", callType: dto.BuildInCallWebSearchCall, toolName: dto.BuildInToolWebSearchPreview}, + {name: "file search", callType: dto.BuildInCallFileSearchCall, toolName: dto.BuildInToolFileSearch}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + c, info := newOpenAIToolBillingContext("gpt-test") + info.ResponsesUsageInfo = &relaycommon.ResponsesUsageInfo{ + BuiltInTools: map[string]*relaycommon.BuildInToolInfo{ + tt.toolName: {ToolName: tt.toolName}, + }, + } + output := dto.ResponsesOutput{ + Type: tt.callType, + ID: "tool-1", + Status: "completed", + } + itemDone := dto.ResponsesStreamResponse{ + Type: dto.ResponsesOutputTypeItemDone, + OutputIndex: common.GetPointer(0), + Item: &output, + } + completed := dto.ResponsesStreamResponse{ + Type: "response.completed", + Response: &dto.OpenAIResponsesResponse{ + Output: []dto.ResponsesOutput{output}, + Usage: detailedResponsesUsage(), + }, + } + itemData, err := common.Marshal(itemDone) + require.NoError(t, err) + completedData, err := common.Marshal(completed) + require.NoError(t, err) + body := append([]byte("data: "), itemData...) + body = append(body, []byte("\ndata: ")...) + body = append(body, itemData...) + body = append(body, []byte("\ndata: ")...) + body = append(body, completedData...) + body = append(body, []byte("\ndata: [DONE]\n")...) + + _, maxAPIError := OaiResponsesStreamHandler(c, info, &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(bytes.NewReader(body)), + }) + require.Nil(t, maxAPIError) + require.Equal(t, 1, info.ResponsesUsageInfo.BuiltInTools[tt.toolName].CallCount) + }) + } +} + +func TestOaiResponsesHandlerRejectsFailedTerminalStatusWithoutBillingTools(t *testing.T) { + setOpenAIToolPricesForTest(t, map[string]float64{"lookup": 5}) + c, info := newOpenAIToolBillingContext("gpt-test") + payload := dto.OpenAIResponsesResponse{ + Status: []byte(`"failed"`), + Output: []dto.ResponsesOutput{{ + Type: dto.BuildInCallFunctionCall, + ID: "fc-1", + CallId: "call-1", + Name: "lookup", + Status: "completed", + }}, + } + body, err := common.Marshal(payload) + require.NoError(t, err) + + _, maxAPIError := OaiResponsesHandler(c, info, &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(bytes.NewReader(body)), + }) + require.NotNil(t, maxAPIError) + require.Empty(t, info.ToolUsageSnapshot().Items) +} + +func TestOaiResponsesStreamHandlerRejectsFailedTerminalAfterToolEvent(t *testing.T) { + setOpenAIToolPricesForTest(t, map[string]float64{"lookup": 5}) + c, info := newOpenAIToolBillingContext("gpt-test") + + body := strings.Join([]string{ + `data: {"type":"response.output_item.done","output_index":0,"item":{"type":"function_call","id":"fc-1","call_id":"call-1","name":"lookup","status":"completed"}}`, + `data: {"type":"response.failed","response":{"id":"resp-1","status":"failed","error":{"type":"server_error","code":"server_error","message":"upstream failed"}}}`, + `data: [DONE]`, + "", + }, "\n") + + _, maxAPIError := OaiResponsesStreamHandler(c, info, &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(body)), + }) + require.NotNil(t, maxAPIError) + require.True(t, types.IsSkipRetryError(maxAPIError)) + require.Empty(t, info.ToolUsageSnapshot().Items) +} + +func TestOaiResponsesStreamHandlerRejectsStandardErrorAfterToolEvent(t *testing.T) { + setOpenAIToolPricesForTest(t, map[string]float64{"lookup": 5}) + c, info := newOpenAIToolBillingContext("gpt-test") + + body := strings.Join([]string{ + `data: {"type":"response.output_item.done","output_index":0,"item":{"type":"function_call","id":"fc-1","call_id":"call-1","name":"lookup","status":"completed"}}`, + `data: {"type":"error","code":"server_error","message":"upstream failed","param":"tool"}`, + `data: [DONE]`, + "", + }, "\n") + + _, maxAPIError := OaiResponsesStreamHandler(c, info, &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(body)), + }) + require.NotNil(t, maxAPIError) + require.True(t, types.IsSkipRetryError(maxAPIError)) + openAIError := maxAPIError.ToOpenAIError() + require.Equal(t, "upstream failed", openAIError.Message) + require.Equal(t, "server_error", openAIError.Code) + require.Equal(t, "tool", openAIError.Param) + require.Empty(t, info.ToolUsageSnapshot().Items) +} + +func TestOaiResponsesStreamHandlerAllowsRetryWhenFirstEventIsError(t *testing.T) { + tests := []struct { + name string + event string + }{ + {name: "standard error", event: `{"type":"error","code":"server_error","message":"upstream failed"}`}, + {name: "failed response", event: `{"type":"response.failed","response":{"id":"resp-1","status":"failed","error":{"message":"upstream failed"}}}`}, + {name: "completed event with failed response", event: `{"type":"response.completed","response":{"id":"resp-1","status":"failed","error":{"message":"upstream failed"}}}`}, + } + 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{ + RelayFormat: types.RelayFormatOpenAIResponses, + OriginModelName: "gpt-test", + DisablePing: true, + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-test"}, + } + body := strings.Join([]string{ + "data: " + tt.event, + `data: [DONE]`, + "", + }, "\n") + + _, maxAPIError := OaiResponsesStreamHandler(c, info, &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(body)), + }) + require.NotNil(t, maxAPIError) + require.False(t, types.IsSkipRetryError(maxAPIError)) + require.Empty(t, recorder.Body.String()) + }) + } +} + +func TestOaiResponsesToChatStreamHandlerPreservesMessageOnlyTerminalError(t *testing.T) { + setOpenAIToolPricesForTest(t, map[string]float64{"lookup": 5}) + c, info := newOpenAIToolBillingContext("gpt-test") + info.SendResponseCount = 1 + + body := strings.Join([]string{ + `data: {"type":"response.failed","response":{"id":"resp-1","status":"failed","error":{"message":"upstream failed"}}}`, + `data: [DONE]`, + "", + }, "\n") + _, maxAPIError := OaiResponsesToChatStreamHandler(c, info, &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(body)), + }) + require.NotNil(t, maxAPIError) + require.True(t, types.IsSkipRetryError(maxAPIError)) + require.Equal(t, "upstream failed", maxAPIError.Error()) +} + +func TestOaiResponsesToChatStreamHandlerRejectsStandardErrorEvent(t *testing.T) { + c, info := newOpenAIToolBillingContext("gpt-test") + info.SendResponseCount = 1 + + body := strings.Join([]string{ + `data: {"type":"error","code":"server_error","message":"upstream failed","param":"tool"}`, + `data: [DONE]`, + "", + }, "\n") + _, maxAPIError := OaiResponsesToChatStreamHandler(c, info, &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(body)), + }) + require.NotNil(t, maxAPIError) + require.True(t, types.IsSkipRetryError(maxAPIError)) + openAIError := maxAPIError.ToOpenAIError() + require.Equal(t, "upstream failed", openAIError.Message) + require.Equal(t, "server_error", openAIError.Code) + require.Equal(t, "tool", openAIError.Param) +} + +func TestOaiResponsesToChatStreamHandlerSkipsRetryAfterOpenAIOutput(t *testing.T) { + gin.SetMode(gin.TestMode) + 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"}, + } + body := strings.Join([]string{ + `data: {"type":"response.output_text.delta","delta":"hello"}`, + `data: {"type":"error","code":"server_error","message":"upstream failed"}`, + `data: [DONE]`, + "", + }, "\n") + + _, maxAPIError := OaiResponsesToChatStreamHandler(c, info, &http.Response{ + StatusCode: http.StatusOK, + Header: make(http.Header), + Body: io.NopCloser(strings.NewReader(body)), + }) + require.NotNil(t, maxAPIError) + require.True(t, types.IsSkipRetryError(maxAPIError)) + require.Greater(t, info.SendResponseCount, 0) + require.Contains(t, recorder.Body.String(), "hello") +} diff --git a/relay/channel/openai/responses_via_chat.go b/relay/channel/openai/responses_via_chat.go index 4b090f7f..6913c807 100644 --- a/relay/channel/openai/responses_via_chat.go +++ b/relay/channel/openai/responses_via_chat.go @@ -35,6 +35,7 @@ func OaiChatToResponsesHandler(c *gin.Context, info *relaycommon.RelayInfo, resp if oaiError := chatResp.GetOpenAIError(); oaiError != nil && oaiError.Type != "" { return nil, types.WithOpenAIError(*oaiError, resp.StatusCode) } + observeOpenAITextToolCalls(info, chatResp.Choices) responseID := helper.GetResponseID(c) responsesResp, usage, err := service.ChatCompletionsResponseToResponsesResponse(&chatResp, responseID) @@ -66,6 +67,7 @@ func OaiChatToResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo responseID := helper.GetResponseID(c) state := openaicompat.NewChatToResponsesStreamState(responseID, info.UpstreamModelName) + toolCallObserver := newOpenAIStreamToolCallObserver(info) var streamErr *types.MaxAPIError sendEvent := func(event openaicompat.ChatToResponsesStreamEvent) bool { @@ -104,6 +106,7 @@ func OaiChatToResponsesStreamHandler(c *gin.Context, info *relaycommon.RelayInfo sr.Error(err) return } + toolCallObserver.Observe(&chunk) events, err := openaicompat.ChatCompletionsStreamChunkToResponsesEvents(&chunk, state) if err != nil { diff --git a/relay/channel/openai/tool_billing_test.go b/relay/channel/openai/tool_billing_test.go new file mode 100644 index 00000000..813cfe43 --- /dev/null +++ b/relay/channel/openai/tool_billing_test.go @@ -0,0 +1,335 @@ +package openai + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/MAX-API-Next/MAX-API/setting/config" + "github.com/MAX-API-Next/MAX-API/setting/operation_setting" + + 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 setOpenAIToolPricesForTest(t *testing.T, additions map[string]float64) { + t.Helper() + + setting := config.GlobalConfig.Get("tool_price_setting").(*operation_setting.ToolPriceSetting) + original := make(map[string]float64, len(setting.Prices)) + for name, price := range setting.Prices { + original[name] = price + } + setting.Prices = make(map[string]float64, len(original)+len(additions)) + for name, price := range original { + setting.Prices[name] = price + } + for name, price := range additions { + setting.Prices[name] = price + } + operation_setting.RebuildToolPriceIndex() + + t.Cleanup(func() { + setting.Prices = original + operation_setting.RebuildToolPriceIndex() + }) +} + +func newOpenAIToolBillingContext(model string) (*gin.Context, *relaycommon.RelayInfo) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequestWithContext(context.Background(), http.MethodPost, "/v1/chat/completions", nil) + ledger := relaycommon.NewToolUsageLedger(model) + ledger.BeginAttempt(0) + return c, &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatOpenAI, + OriginModelName: model, + DisablePing: true, + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: model, + }, + ToolUsage: ledger, + } +} + +func TestOpenaiHandlerRecordsActualCustomToolCalls(t *testing.T) { + setOpenAIToolPricesForTest(t, map[string]float64{"lookup": 5}) + c, info := newOpenAIToolBillingContext("gpt-test") + + payload := `{ + "id":"chatcmpl-1", + "object":"chat.completion", + "model":"gpt-test", + "choices":[{ + "index":0, + "message":{"role":"assistant","content":"","tool_calls":[ + {"id":"call-1","type":"function","function":{"name":"lookup","arguments":"{}"}} + ]}, + "finish_reason":"tool_calls" + }], + "usage":{"prompt_tokens":10,"completion_tokens":2,"total_tokens":12} + }` + resp := &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(payload))} + + _, maxErr := OpenaiHandler(c, info, resp) + require.Nil(t, maxErr) + require.True(t, info.CommitToolUsageAttempt()) + require.Equal(t, []relaycommon.ToolUsageItem{{ + Name: "lookup", + CallCount: 1, + PricePer1K: 5, + }}, info.ToolUsageSnapshot().Items) +} + +func TestOpenaiHandlerDoesNotRecordIncompleteCustomToolCalls(t *testing.T) { + setOpenAIToolPricesForTest(t, map[string]float64{"lookup": 5}) + c, info := newOpenAIToolBillingContext("gpt-test") + + payload := `{ + "id":"chatcmpl-1", + "object":"chat.completion", + "model":"gpt-test", + "choices":[{ + "index":0, + "message":{"role":"assistant","content":"","tool_calls":[ + {"id":"call-1","type":"function","function":{"name":"lookup","arguments":"{}"}} + ]}, + "finish_reason":"length" + }], + "usage":{"prompt_tokens":10,"completion_tokens":2,"total_tokens":12} + }` + resp := &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(payload))} + + _, maxErr := OpenaiHandler(c, info, resp) + require.Nil(t, maxErr) + require.True(t, info.CommitToolUsageAttempt()) + require.Empty(t, info.ToolUsageSnapshot().Items) +} + +func TestOaiStreamHandlerDeduplicatesCustomToolCallChunks(t *testing.T) { + setOpenAIToolPricesForTest(t, map[string]float64{ + "lookup": 5, + "get_time": 7, + }) + c, info := newOpenAIToolBillingContext("gpt-test") + + body := strings.Join([]string{ + `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","model":"gpt-test","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call-1","type":"function","function":{"name":"lookup","arguments":""}}]},"finish_reason":null}]}`, + `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","model":"gpt-test","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{}"}}]},"finish_reason":null}]}`, + `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","model":"gpt-test","choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"id":"call-2","type":"function","function":{"name":"get_time","arguments":"{}"}}]},"finish_reason":null}]}`, + `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","model":"gpt-test","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":10,"completion_tokens":2,"total_tokens":12}}`, + `data: [DONE]`, + "", + }, "\n") + resp := &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(body))} + + _, maxErr := OaiStreamHandler(c, info, resp) + require.Nil(t, maxErr) + require.True(t, info.CommitToolUsageAttempt()) + require.Equal(t, []relaycommon.ToolUsageItem{ + {Name: "get_time", CallCount: 1, PricePer1K: 7}, + {Name: "lookup", CallCount: 1, PricePer1K: 5}, + }, info.ToolUsageSnapshot().Items) +} + +func TestOaiStreamHandlerDoesNotRecordIncompleteCustomToolCall(t *testing.T) { + setOpenAIToolPricesForTest(t, map[string]float64{"lookup": 5}) + c, info := newOpenAIToolBillingContext("gpt-test") + + body := strings.Join([]string{ + `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","model":"gpt-test","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call-1","type":"function","function":{"name":"lookup","arguments":""}}]},"finish_reason":null}]}`, + `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","model":"gpt-test","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{}"}}]},"finish_reason":null}]}`, + `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","model":"gpt-test","choices":[{"index":0,"delta":{},"finish_reason":"length"}],"usage":{"prompt_tokens":10,"completion_tokens":2,"total_tokens":12}}`, + `data: [DONE]`, + "", + }, "\n") + resp := &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(body))} + + _, maxErr := OaiStreamHandler(c, info, resp) + require.Nil(t, maxErr) + require.True(t, info.CommitToolUsageAttempt()) + require.Empty(t, info.ToolUsageSnapshot().Items) +} + +func TestOaiStreamHandlerDoesNotRecordUnterminatedCustomToolCall(t *testing.T) { + setOpenAIToolPricesForTest(t, map[string]float64{"lookup": 5}) + c, info := newOpenAIToolBillingContext("gpt-test") + + body := strings.Join([]string{ + `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","model":"gpt-test","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call-1","type":"function","function":{"name":"lookup","arguments":"{}"}}]},"finish_reason":null}]}`, + `data: [DONE]`, + "", + }, "\n") + resp := &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(body))} + + _, maxErr := OaiStreamHandler(c, info, resp) + require.Nil(t, maxErr) + require.True(t, info.CommitToolUsageAttempt()) + require.Empty(t, info.ToolUsageSnapshot().Items) +} + +func TestOpenAIStreamChunkMayAffectToolBilling(t *testing.T) { + tests := []struct { + name string + data string + want bool + }{ + {name: "text delta with null finish", data: `{"choices":[{"delta":{"content":"ok"},"finish_reason":null}]}`, want: false}, + {name: "tool fragment", data: `{"choices":[{"delta":{"tool_calls":[{"index":0}]},"finish_reason":null}]}`, want: true}, + {name: "terminal choice", data: `{"choices":[{"delta":{},"finish_reason" : "length"}]}`, want: true}, + {name: "second choice terminal", data: `{"choices":[{"delta":{},"finish_reason":null},{"delta":{},"finish_reason":"tool_calls"}]}`, want: true}, + {name: "all choices open", data: `{"choices":[{"delta":{},"finish_reason":null},{"delta":{},"finish_reason":null}]}`, want: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.Equal(t, tt.want, openAIStreamChunkMayAffectToolBilling(tt.data)) + }) + } +} + +func TestOaiResponsesToChatHandlerRecordsConvertedToolCallOnce(t *testing.T) { + setOpenAIToolPricesForTest(t, map[string]float64{"lookup": 5}) + c, info := newOpenAIToolBillingContext("gpt-test") + + payload := `{ + "id":"resp-1","object":"response","status":"completed","model":"gpt-test", + "output":[{"type":"function_call","id":"fc-1","call_id":"call-1","name":"lookup","arguments":"{}","status":"completed"}], + "usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2} + }` + resp := &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(payload))} + + _, maxErr := OaiResponsesToChatHandler(c, info, resp) + require.Nil(t, maxErr) + require.True(t, info.CommitToolUsageAttempt()) + require.Equal(t, []relaycommon.ToolUsageItem{{Name: "lookup", CallCount: 1, PricePer1K: 5}}, info.ToolUsageSnapshot().Items) +} + +func TestOaiResponsesToChatStreamHandlerDeduplicatesDoneAndCompleted(t *testing.T) { + setOpenAIToolPricesForTest(t, map[string]float64{"lookup": 5}) + c, info := newOpenAIToolBillingContext("gpt-test") + + body := strings.Join([]string{ + `data: {"type":"response.output_item.done","output_index":0,"item":{"type":"function_call","id":"fc-1","call_id":"call-1","name":"lookup","arguments":"{}","status":"completed"}}`, + `data: {"type":"response.completed","response":{"id":"resp-1","status":"completed","model":"gpt-test","output":[{"type":"function_call","id":"fc-1","call_id":"call-1","name":"lookup","arguments":"{}","status":"completed"}],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}`, + `data: [DONE]`, + "", + }, "\n") + resp := &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(body))} + + _, maxErr := OaiResponsesToChatStreamHandler(c, info, resp) + require.Nil(t, maxErr) + require.True(t, info.CommitToolUsageAttempt()) + require.Equal(t, []relaycommon.ToolUsageItem{{Name: "lookup", CallCount: 1, PricePer1K: 5}}, info.ToolUsageSnapshot().Items) +} + +func TestOaiResponsesToChatStreamHandlerReplaysToolCallAfterText(t *testing.T) { + setOpenAIToolPricesForTest(t, map[string]float64{"lookup": 5}) + recorder := httptest.NewRecorder() + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequestWithContext(context.Background(), http.MethodPost, "/v1/chat/completions", nil) + info := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatOpenAI, + OriginModelName: "gpt-test", + DisablePing: true, + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-test"}, + ToolUsage: relaycommon.NewToolUsageLedger("gpt-test"), + } + info.ToolUsage.BeginAttempt(0) + body := strings.Join([]string{ + `data: {"type":"response.output_text.delta","delta":"hello"}`, + `data: {"type":"response.completed","response":{"id":"resp-1","status":"completed","model":"gpt-test","output":[{"type":"message","id":"msg-1","status":"completed","content":[{"type":"output_text","text":"hello"}]},{"type":"custom_tool_call","id":"ct-1","call_id":"call-1","name":"lookup","input":"payload","status":"completed"}],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}`, + `data: [DONE]`, + "", + }, "\n") + resp := &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(body))} + + _, maxErr := OaiResponsesToChatStreamHandler(c, info, resp) + require.Nil(t, maxErr) + output := recorder.Body.String() + require.Contains(t, output, `"arguments":"payload"`) + require.Contains(t, output, `"name":"lookup"`) + require.Contains(t, output, `"finish_reason":"tool_calls"`) +} + +func TestOaiChatToResponsesHandlerRecordsConvertedToolCallOnce(t *testing.T) { + setOpenAIToolPricesForTest(t, map[string]float64{"lookup": 5}) + c, info := newOpenAIToolBillingContext("gpt-test") + info.RelayFormat = types.RelayFormatOpenAIResponses + + payload := `{ + "id":"chatcmpl-1","object":"chat.completion","model":"gpt-test", + "choices":[{"index":0,"message":{"role":"assistant","content":"","tool_calls":[{"id":"call-1","type":"function","function":{"name":"lookup","arguments":"{}"}}]},"finish_reason":"tool_calls"}], + "usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2} + }` + resp := &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(payload))} + + _, maxErr := OaiChatToResponsesHandler(c, info, resp) + require.Nil(t, maxErr) + require.True(t, info.CommitToolUsageAttempt()) + require.Equal(t, []relaycommon.ToolUsageItem{{Name: "lookup", CallCount: 1, PricePer1K: 5}}, info.ToolUsageSnapshot().Items) +} + +func TestOaiChatToResponsesHandlerDoesNotRecordIncompleteToolCall(t *testing.T) { + setOpenAIToolPricesForTest(t, map[string]float64{"lookup": 5}) + c, info := newOpenAIToolBillingContext("gpt-test") + info.RelayFormat = types.RelayFormatOpenAIResponses + + payload := `{ + "id":"chatcmpl-1","object":"chat.completion","model":"gpt-test", + "choices":[{"index":0,"message":{"role":"assistant","content":"","tool_calls":[{"id":"call-1","type":"function","function":{"name":"lookup","arguments":"{}"}}]},"finish_reason":"length"}], + "usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2} + }` + resp := &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(payload))} + + _, maxErr := OaiChatToResponsesHandler(c, info, resp) + require.Nil(t, maxErr) + require.True(t, info.CommitToolUsageAttempt()) + require.Empty(t, info.ToolUsageSnapshot().Items) +} + +func TestOaiChatToResponsesStreamHandlerDeduplicatesToolFragments(t *testing.T) { + setOpenAIToolPricesForTest(t, map[string]float64{"lookup": 5}) + c, info := newOpenAIToolBillingContext("gpt-test") + info.RelayFormat = types.RelayFormatOpenAIResponses + + body := strings.Join([]string{ + `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","model":"gpt-test","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call-1","type":"function","function":{"name":"lookup","arguments":""}}]},"finish_reason":null}]}`, + `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","model":"gpt-test","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{}"}}]},"finish_reason":null}]}`, + `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","model":"gpt-test","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`, + `data: [DONE]`, + "", + }, "\n") + resp := &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(body))} + + _, maxErr := OaiChatToResponsesStreamHandler(c, info, resp) + require.Nil(t, maxErr) + require.True(t, info.CommitToolUsageAttempt()) + require.Equal(t, []relaycommon.ToolUsageItem{{Name: "lookup", CallCount: 1, PricePer1K: 5}}, info.ToolUsageSnapshot().Items) +} + +func TestOaiChatToResponsesStreamHandlerDoesNotRecordIncompleteToolCall(t *testing.T) { + setOpenAIToolPricesForTest(t, map[string]float64{"lookup": 5}) + c, info := newOpenAIToolBillingContext("gpt-test") + info.RelayFormat = types.RelayFormatOpenAIResponses + + body := strings.Join([]string{ + `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","model":"gpt-test","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call-1","type":"function","function":{"name":"lookup","arguments":""}}]},"finish_reason":null}]}`, + `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","model":"gpt-test","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{}"}}]},"finish_reason":null}]}`, + `data: {"id":"chatcmpl-1","object":"chat.completion.chunk","model":"gpt-test","choices":[{"index":0,"delta":{},"finish_reason":"content_filter"}],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}`, + `data: [DONE]`, + "", + }, "\n") + resp := &http.Response{StatusCode: http.StatusOK, Body: io.NopCloser(strings.NewReader(body))} + + _, maxErr := OaiChatToResponsesStreamHandler(c, info, resp) + require.Nil(t, maxErr) + require.True(t, info.CommitToolUsageAttempt()) + require.Empty(t, info.ToolUsageSnapshot().Items) +} diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go index 57c07c7e..fe97b94e 100644 --- a/relay/common/relay_info.go +++ b/relay/common/relay_info.go @@ -165,6 +165,7 @@ type RelayInfo struct { PriceData types.PriceData TaskBilling *types.TaskBillingResult + ToolUsage *ToolUsageLedger // QuotaClamp is set when a quota conversion saturated at the int32 bound // or fell back from NaN while computing this request's charge. @@ -196,6 +197,19 @@ type RelayInfo struct { } func (info *RelayInfo) InitChannelMeta(c *gin.Context) { + if info.beginToolUsageAttempt() { + if info.ResponsesUsageInfo != nil { + for _, tool := range info.ResponsesUsageInfo.BuiltInTools { + if tool != nil { + tool.CallCount = 0 + } + } + } + c.Set("claude_web_search_requests", 0) + c.Set("image_generation_call", false) + c.Set("image_generation_call_quality", "") + c.Set("image_generation_call_size", "") + } channelType := common.GetContextKeyInt(c, constant.ContextKeyChannelType) paramOverride := common.GetContextKeyStringMap(c, constant.ContextKeyChannelParamOverride) headerOverride := common.GetContextKeyStringMap(c, constant.ContextKeyChannelHeaderOverride) @@ -503,9 +517,11 @@ func genBaseRelayInfo(c *gin.Context, request dto.Request) *RelayInfo { if reqId == "" { reqId = common.GetTimeString() + common.GetRandomString(8) } + originModelName := common.GetContextKeyString(c, constant.ContextKeyOriginalModel) info := &RelayInfo{ Request: request, ReasoningEffort: reasoningEffortFromRequest(request), + ToolUsage: NewToolUsageLedger(originModelName), RequestId: reqId, UserId: common.GetContextKeyInt(c, constant.ContextKeyUserId), @@ -514,7 +530,7 @@ func genBaseRelayInfo(c *gin.Context, request dto.Request) *RelayInfo { UserQuota: common.GetContextKeyInt64(c, constant.ContextKeyUserQuota), UserEmail: common.GetContextKeyString(c, constant.ContextKeyUserEmail), - OriginModelName: common.GetContextKeyString(c, constant.ContextKeyOriginalModel), + OriginModelName: originModelName, TokenId: common.GetContextKeyInt(c, constant.ContextKeyTokenId), TokenKey: common.GetContextKeyString(c, constant.ContextKeyTokenKey), diff --git a/relay/common/tool_usage.go b/relay/common/tool_usage.go new file mode 100644 index 00000000..54c11288 --- /dev/null +++ b/relay/common/tool_usage.go @@ -0,0 +1,391 @@ +package common + +import ( + "sort" + "strings" + "sync" + + "github.com/MAX-API-Next/MAX-API/setting/operation_setting" +) + +type ToolUsageKind string + +const ( + ToolUsageKindCustom ToolUsageKind = "custom" + ToolUsageKindBuiltIn ToolUsageKind = "built_in" +) + +var reservedBillableToolNames = map[string]struct{}{ + "web_search": {}, + "web_search_preview": {}, + "file_search": {}, + "google_search": {}, + "image_generation": {}, +} + +// ToolCallIdentity carries independent aliases for one actual upstream call. +// Scope separates protocols; CallID is preferred when available; Position is +// the stable choice/output/content-block position used by streaming protocols. +type ToolCallIdentity struct { + Scope string + CallID string + Position string +} + +type ToolUsageItem struct { + Name string `json:"name"` + CallCount int `json:"call_count"` + PricePer1K float64 `json:"price_per_1k"` +} + +type ToolUsageSnapshot struct { + Attempt int `json:"attempt"` + ModelName string `json:"model_name"` + PriceVersion uint64 `json:"price_version"` + Items []ToolUsageItem `json:"items,omitempty"` +} + +type toolUsageRecord struct { + name string + kind ToolUsageKind + aliases map[string]struct{} +} + +// ToolUsageLedger is a request-local, attempt-aware ledger. It records only +// actual, stably identifiable upstream calls and prices them using the frozen +// request-start snapshot. Starting a new retry attempt discards observations +// from the failed attempt without touching the durable billing lifecycle. +type ToolUsageLedger struct { + mu sync.Mutex + + modelName string + prices operation_setting.ToolPriceSnapshot + + attempt int + attemptStarted bool + committed bool + nextRecordID uint64 + records map[uint64]*toolUsageRecord + aliases map[string]uint64 + ambiguous map[string]struct{} +} + +func NewToolUsageLedger(modelName string) *ToolUsageLedger { + return &ToolUsageLedger{ + modelName: modelName, + prices: operation_setting.CaptureToolPriceSnapshot(modelName), + } +} + +// PriceFor returns the request-start price for a tool. The underlying price +// index is immutable, so config reloads cannot change an in-flight request's +// billing decision. +func (l *ToolUsageLedger) PriceFor(toolName string) float64 { + if l == nil { + return 0 + } + return l.prices.PriceFor(toolName) +} + +// BeginAttempt selects the active upstream attempt. Re-entering the same +// attempt is a no-op so multiple response stages can share one ledger. +func (l *ToolUsageLedger) BeginAttempt(attempt int) bool { + if l == nil { + return false + } + l.mu.Lock() + defer l.mu.Unlock() + return l.beginAttemptLocked(attempt) +} + +func (l *ToolUsageLedger) beginAttemptLocked(attempt int) bool { + if l.attemptStarted && l.attempt == attempt { + return false + } + l.attempt = attempt + l.attemptStarted = true + l.committed = false + l.nextRecordID = 0 + l.records = make(map[uint64]*toolUsageRecord) + l.aliases = make(map[string]uint64) + l.ambiguous = make(map[string]struct{}) + return true +} + +// CommitAttempt makes the active attempt visible to settlement. Observations +// from an uncommitted or failed attempt never appear in Snapshot. +func (l *ToolUsageLedger) CommitAttempt(attempt int) bool { + if l == nil { + return false + } + l.mu.Lock() + defer l.mu.Unlock() + if !l.attemptStarted || l.attempt != attempt { + return false + } + l.committed = true + return true +} + +func (l *ToolUsageLedger) ObserveCustom(name string, identity ToolCallIdentity) bool { + return l.observe(ToolUsageKindCustom, name, identity) +} + +func (l *ToolUsageLedger) ObserveBuiltIn(name string, identity ToolCallIdentity) bool { + return l.observe(ToolUsageKindBuiltIn, name, identity) +} + +func (l *ToolUsageLedger) observe(kind ToolUsageKind, name string, identity ToolCallIdentity) bool { + if l == nil { + return false + } + name = strings.TrimSpace(name) + if name == "" { + return false + } + if kind == ToolUsageKindCustom { + if _, reserved := reservedBillableToolNames[name]; reserved { + return false + } + } + if l.prices.PriceFor(name) <= 0 { + return false + } + + identity.Scope = strings.TrimSpace(identity.Scope) + identity.CallID = strings.TrimSpace(identity.CallID) + identity.Position = strings.TrimSpace(identity.Position) + if identity.Scope == "" || (identity.CallID == "" && identity.Position == "") { + return false + } + + aliases := make([]string, 0, 2) + if identity.CallID != "" { + aliases = append(aliases, identity.Scope+"\x00id\x00"+identity.CallID) + } + if identity.Position != "" { + aliases = append(aliases, identity.Scope+"\x00position\x00"+identity.Position) + } + + l.mu.Lock() + defer l.mu.Unlock() + if !l.attemptStarted { + l.beginAttemptLocked(0) + } + if l.committed { + return false + } + + matched := make(map[uint64]struct{}) + hasAmbiguousAlias := false + for _, alias := range aliases { + if _, ambiguous := l.ambiguous[alias]; ambiguous { + hasAmbiguousAlias = true + } + if recordID, ok := l.aliases[alias]; ok { + matched[recordID] = struct{}{} + } + } + if hasAmbiguousAlias { + l.invalidateAmbiguousLocked(matched, aliases) + return false + } + + if len(matched) == 0 { + l.nextRecordID++ + recordID := l.nextRecordID + record := &toolUsageRecord{ + name: name, + kind: kind, + aliases: make(map[string]struct{}, len(aliases)), + } + for _, alias := range aliases { + record.aliases[alias] = struct{}{} + l.aliases[alias] = recordID + } + l.records[recordID] = record + return true + } + + var primaryID uint64 + for recordID := range matched { + record := l.records[recordID] + if record == nil || record.name != name || record.kind != kind { + l.invalidateAmbiguousLocked(matched, aliases) + return false + } + if primaryID == 0 || recordID < primaryID { + primaryID = recordID + } + } + primary := l.records[primaryID] + for recordID := range matched { + if recordID == primaryID { + continue + } + duplicate := l.records[recordID] + for alias := range duplicate.aliases { + primary.aliases[alias] = struct{}{} + l.aliases[alias] = primaryID + } + delete(l.records, recordID) + } + for _, alias := range aliases { + primary.aliases[alias] = struct{}{} + l.aliases[alias] = primaryID + } + return false +} + +// invalidateAmbiguousLocked removes every record connected to a conflicting +// identity and tombstones all known aliases. A later chunk therefore cannot +// recreate a charge from only one side of an identity conflict. +func (l *ToolUsageLedger) invalidateAmbiguousLocked(recordIDs map[uint64]struct{}, aliases []string) { + if l.ambiguous == nil { + l.ambiguous = make(map[string]struct{}) + } + for recordID := range recordIDs { + record := l.records[recordID] + if record == nil { + continue + } + for alias := range record.aliases { + delete(l.aliases, alias) + l.ambiguous[alias] = struct{}{} + } + delete(l.records, recordID) + } + for _, alias := range aliases { + delete(l.aliases, alias) + l.ambiguous[alias] = struct{}{} + } +} + +func (l *ToolUsageLedger) Snapshot() ToolUsageSnapshot { + if l == nil { + return ToolUsageSnapshot{} + } + l.mu.Lock() + defer l.mu.Unlock() + + counts := make(map[string]int) + if l.committed { + for _, record := range l.records { + if record.kind != ToolUsageKindCustom { + continue + } + counts[record.name]++ + } + } + items := make([]ToolUsageItem, 0, len(counts)) + for name, count := range counts { + items = append(items, ToolUsageItem{ + Name: name, + CallCount: count, + PricePer1K: l.prices.PriceFor(name), + }) + } + sort.Slice(items, func(i, j int) bool { + if items[i].Name == items[j].Name { + return items[i].PricePer1K < items[j].PricePer1K + } + return items[i].Name < items[j].Name + }) + + return ToolUsageSnapshot{ + Attempt: l.attempt, + ModelName: l.modelName, + PriceVersion: l.prices.Version(), + Items: items, + } +} + +func (l *ToolUsageLedger) currentBuiltInCallCounts() map[string]int { + counts := make(map[string]int) + if l == nil { + return counts + } + l.mu.Lock() + defer l.mu.Unlock() + for _, record := range l.records { + if record != nil && record.kind == ToolUsageKindBuiltIn { + counts[record.name]++ + } + } + return counts +} + +func (info *RelayInfo) beginToolUsageAttempt() bool { + if info == nil { + return false + } + if info.ToolUsage == nil { + info.ToolUsage = NewToolUsageLedger(info.OriginModelName) + } + return info.ToolUsage.BeginAttempt(info.RetryIndex) +} + +func (info *RelayInfo) ObserveCustomToolCall(name string, identity ToolCallIdentity) bool { + if info == nil { + return false + } + if info.ToolUsage == nil { + info.ToolUsage = NewToolUsageLedger(info.OriginModelName) + info.ToolUsage.BeginAttempt(info.RetryIndex) + } + return info.ToolUsage.ObserveCustom(name, identity) +} + +func (info *RelayInfo) ObserveBuiltInToolCall(name string, identity ToolCallIdentity) bool { + if info == nil { + return false + } + if info.ToolUsage == nil { + info.ToolUsage = NewToolUsageLedger(info.OriginModelName) + info.ToolUsage.BeginAttempt(info.RetryIndex) + } + accepted := info.ToolUsage.ObserveBuiltIn(name, identity) + if info.ResponsesUsageInfo != nil { + counts := info.ToolUsage.currentBuiltInCallCounts() + for _, tool := range info.ResponsesUsageInfo.BuiltInTools { + if tool != nil { + tool.CallCount = 0 + } + } + for toolName, count := range counts { + if tool := info.ResponsesUsageInfo.BuiltInTools[toolName]; tool != nil { + tool.CallCount = count + } + } + } + return accepted +} + +// FrozenToolPrice returns the request-scoped price used by both pre-consume +// and settlement paths. Normal relay requests initialize ToolUsage in +// genBaseRelayInfo before any upstream work. The lazy fallback keeps +// hand-built RelayInfo values used by legacy callers/tests compatible while +// still freezing their first observed price for the remainder of the request. +func (info *RelayInfo) FrozenToolPrice(toolName string) float64 { + if info == nil { + return 0 + } + if info.ToolUsage == nil { + info.ToolUsage = NewToolUsageLedger(info.OriginModelName) + } + return info.ToolUsage.PriceFor(toolName) +} + +func (info *RelayInfo) CommitToolUsageAttempt() bool { + if info == nil || info.ToolUsage == nil { + return false + } + return info.ToolUsage.CommitAttempt(info.RetryIndex) +} + +func (info *RelayInfo) ToolUsageSnapshot() ToolUsageSnapshot { + if info == nil || info.ToolUsage == nil { + return ToolUsageSnapshot{} + } + return info.ToolUsage.Snapshot() +} diff --git a/relay/common/tool_usage_test.go b/relay/common/tool_usage_test.go new file mode 100644 index 00000000..73eb1c5f --- /dev/null +++ b/relay/common/tool_usage_test.go @@ -0,0 +1,354 @@ +package common + +import ( + "sync" + "sync/atomic" + "testing" + + "github.com/MAX-API-Next/MAX-API/setting/config" + "github.com/MAX-API-Next/MAX-API/setting/operation_setting" + "github.com/stretchr/testify/require" +) + +func replaceToolPricesForTest(t *testing.T, prices map[string]float64) { + t.Helper() + + setting := config.GlobalConfig.Get("tool_price_setting").(*operation_setting.ToolPriceSetting) + original := make(map[string]float64, len(setting.Prices)) + for name, price := range setting.Prices { + original[name] = price + } + + setting.Prices = make(map[string]float64, len(prices)) + for name, price := range prices { + setting.Prices[name] = price + } + operation_setting.RebuildToolPriceIndex() + + t.Cleanup(func() { + setting.Prices = original + operation_setting.RebuildToolPriceIndex() + }) +} + +func TestToolUsageLedgerRecordsOnlyPricedCustomCalls(t *testing.T) { + replaceToolPricesForTest(t, map[string]float64{ + "lookup": 5, + }) + + ledger := NewToolUsageLedger("gpt-test") + ledger.BeginAttempt(0) + + require.True(t, ledger.ObserveCustom("lookup", ToolCallIdentity{ + Scope: "openai-chat", + CallID: "call-1", + Position: "choice:0:tool:0", + })) + require.False(t, ledger.ObserveCustom("unpriced", ToolCallIdentity{ + Scope: "openai-chat", + CallID: "call-2", + Position: "choice:0:tool:1", + })) + require.True(t, ledger.CommitAttempt(0)) + + snapshot := ledger.Snapshot() + require.Equal(t, "gpt-test", snapshot.ModelName) + require.Len(t, snapshot.Items, 1) + require.Equal(t, ToolUsageItem{ + Name: "lookup", + CallCount: 1, + PricePer1K: 5, + }, snapshot.Items[0]) +} + +func TestToolUsageLedgerSnapshotExcludesBuiltInRecords(t *testing.T) { + replaceToolPricesForTest(t, map[string]float64{ + "lookup": 5, + "web_search": 10, + }) + + ledger := NewToolUsageLedger("gpt-test") + ledger.BeginAttempt(0) + require.True(t, ledger.ObserveBuiltIn("web_search", ToolCallIdentity{ + Scope: "openai-responses", + CallID: "web-1", + Position: "output:0", + })) + require.True(t, ledger.ObserveCustom("lookup", ToolCallIdentity{ + Scope: "openai-responses", + CallID: "custom-1", + Position: "output:1", + })) + require.True(t, ledger.CommitAttempt(0)) + + require.Equal(t, []ToolUsageItem{{ + Name: "lookup", + CallCount: 1, + PricePer1K: 5, + }}, ledger.Snapshot().Items) +} + +func TestRelayInfoBuiltInCountsFollowAttemptLedger(t *testing.T) { + replaceToolPricesForTest(t, map[string]float64{ + "web_search_preview": 5, + "file_search": 7, + }) + + info := &RelayInfo{ + OriginModelName: "gpt-test", + ResponsesUsageInfo: &ResponsesUsageInfo{BuiltInTools: map[string]*BuildInToolInfo{ + "web_search_preview": {ToolName: "web_search_preview"}, + "file_search": {ToolName: "file_search"}, + }}, + ToolUsage: NewToolUsageLedger("gpt-test"), + } + info.ToolUsage.BeginAttempt(0) + identity := ToolCallIdentity{Scope: "openai-responses", CallID: "call-1", Position: "output:0"} + + require.True(t, info.ObserveBuiltInToolCall("web_search_preview", identity)) + require.Equal(t, 1, info.ResponsesUsageInfo.BuiltInTools["web_search_preview"].CallCount) + require.False(t, info.ObserveBuiltInToolCall("web_search_preview", identity)) + require.Equal(t, 1, info.ResponsesUsageInfo.BuiltInTools["web_search_preview"].CallCount) + + // Conflicting identities invalidate the previously accepted record, so the + // projected counter must fail closed instead of retaining a charge. + require.False(t, info.ObserveBuiltInToolCall("file_search", identity)) + require.Zero(t, info.ResponsesUsageInfo.BuiltInTools["web_search_preview"].CallCount) + require.Zero(t, info.ResponsesUsageInfo.BuiltInTools["file_search"].CallCount) + + info.RetryIndex = 1 + info.ToolUsage.BeginAttempt(1) + require.True(t, info.ObserveBuiltInToolCall("file_search", ToolCallIdentity{ + Scope: "openai-responses", + CallID: "call-2", + Position: "output:1", + })) + require.Zero(t, info.ResponsesUsageInfo.BuiltInTools["web_search_preview"].CallCount) + require.Equal(t, 1, info.ResponsesUsageInfo.BuiltInTools["file_search"].CallCount) +} + +func TestToolUsageLedgerDeduplicatesAliasesButCountsDistinctCalls(t *testing.T) { + replaceToolPricesForTest(t, map[string]float64{ + "lookup": 5, + }) + + ledger := NewToolUsageLedger("gpt-test") + ledger.BeginAttempt(0) + + require.True(t, ledger.ObserveCustom("lookup", ToolCallIdentity{ + Scope: "openai-chat", + Position: "choice:0:tool:0", + })) + // Some providers emit the stable call ID in a later chunk. The position + // alias must join that later observation to the already-counted call. + require.False(t, ledger.ObserveCustom("lookup", ToolCallIdentity{ + Scope: "openai-chat", + CallID: "call-1", + Position: "choice:0:tool:0", + })) + require.False(t, ledger.ObserveCustom("lookup", ToolCallIdentity{ + Scope: "openai-chat", + CallID: "call-1", + })) + require.True(t, ledger.ObserveCustom("lookup", ToolCallIdentity{ + Scope: "openai-chat", + CallID: "call-2", + Position: "choice:0:tool:1", + })) + require.True(t, ledger.CommitAttempt(0)) + + snapshot := ledger.Snapshot() + require.Len(t, snapshot.Items, 1) + require.Equal(t, 2, snapshot.Items[0].CallCount) +} + +func TestToolUsageLedgerRejectsObservationsAfterCommit(t *testing.T) { + replaceToolPricesForTest(t, map[string]float64{ + "lookup": 5, + }) + + ledger := NewToolUsageLedger("gpt-test") + ledger.BeginAttempt(0) + require.True(t, ledger.ObserveCustom("lookup", ToolCallIdentity{ + Scope: "openai-chat", + CallID: "call-1", + })) + require.True(t, ledger.CommitAttempt(0)) + require.False(t, ledger.ObserveCustom("lookup", ToolCallIdentity{ + Scope: "openai-chat", + CallID: "call-2", + })) + + require.Equal(t, []ToolUsageItem{{ + Name: "lookup", + CallCount: 1, + PricePer1K: 5, + }}, ledger.Snapshot().Items) +} + +func TestToolUsageLedgerAliasConflictFailsClosed(t *testing.T) { + replaceToolPricesForTest(t, map[string]float64{ + "lookup": 5, + "search": 7, + }) + + ledger := NewToolUsageLedger("gpt-test") + ledger.BeginAttempt(0) + require.True(t, ledger.ObserveCustom("lookup", ToolCallIdentity{ + Scope: "gemini", + CallID: "call-1", + Position: "candidate:0:part:0", + })) + + // The same upstream aliases resolving to different names are ambiguous. + // Billing must discard the conflicted record rather than retain a charge + // that may refer to a conversion artifact or a replayed stream position. + require.False(t, ledger.ObserveCustom("search", ToolCallIdentity{ + Scope: "gemini", + CallID: "call-1", + Position: "candidate:0:part:0", + })) + require.True(t, ledger.CommitAttempt(0)) + require.Empty(t, ledger.Snapshot().Items) +} + +func TestToolUsageLedgerRequiresStableIdentity(t *testing.T) { + replaceToolPricesForTest(t, map[string]float64{ + "lookup": 5, + }) + + ledger := NewToolUsageLedger("gpt-test") + ledger.BeginAttempt(0) + + require.False(t, ledger.ObserveCustom("lookup", ToolCallIdentity{Scope: "openai-chat"})) + require.True(t, ledger.CommitAttempt(0)) + require.Empty(t, ledger.Snapshot().Items) +} + +func TestToolUsageLedgerRejectsReservedNamesForCustomCalls(t *testing.T) { + replaceToolPricesForTest(t, map[string]float64{ + "web_search": 10, + "web_search_preview": 10, + "file_search": 2.5, + "google_search": 14, + "image_generation": 40, + }) + + ledger := NewToolUsageLedger("gpt-test") + ledger.BeginAttempt(0) + + for i, name := range []string{ + "web_search", + "web_search_preview", + "file_search", + "google_search", + "image_generation", + } { + require.False(t, ledger.ObserveCustom(name, ToolCallIdentity{ + Scope: "openai-chat", + Position: "choice:0:tool:" + string(rune('0'+i)), + })) + } + require.True(t, ledger.CommitAttempt(0)) + require.Empty(t, ledger.Snapshot().Items) +} + +func TestToolUsageLedgerRetryKeepsOnlyWinningAttempt(t *testing.T) { + replaceToolPricesForTest(t, map[string]float64{ + "first_attempt": 5, + "second_attempt": 7, + }) + + ledger := NewToolUsageLedger("gpt-test") + ledger.BeginAttempt(0) + require.True(t, ledger.ObserveCustom("first_attempt", ToolCallIdentity{ + Scope: "openai-chat", + CallID: "failed-call", + Position: "choice:0:tool:0", + })) + + ledger.BeginAttempt(1) + require.True(t, ledger.ObserveCustom("second_attempt", ToolCallIdentity{ + Scope: "openai-chat", + CallID: "winning-call", + Position: "choice:0:tool:0", + })) + // Re-entering the same attempt must not discard observations made by an + // earlier stage of the same response pipeline. + ledger.BeginAttempt(1) + require.True(t, ledger.CommitAttempt(1)) + + snapshot := ledger.Snapshot() + require.Equal(t, 1, snapshot.Attempt) + require.Equal(t, []ToolUsageItem{{ + Name: "second_attempt", + CallCount: 1, + PricePer1K: 7, + }}, snapshot.Items) +} + +func TestToolUsageLedgerFreezesPricesAtRequestStart(t *testing.T) { + replaceToolPricesForTest(t, map[string]float64{ + "lookup": 5, + }) + + ledger := NewToolUsageLedger("gpt-test") + ledger.BeginAttempt(0) + + setting := config.GlobalConfig.Get("tool_price_setting").(*operation_setting.ToolPriceSetting) + setting.Prices = map[string]float64{"lookup": 9} + operation_setting.RebuildToolPriceIndex() + + require.True(t, ledger.ObserveCustom("lookup", ToolCallIdentity{ + Scope: "openai-chat", + CallID: "call-1", + Position: "choice:0:tool:0", + })) + require.True(t, ledger.CommitAttempt(0)) + + snapshot := ledger.Snapshot() + require.NotZero(t, snapshot.PriceVersion) + require.Len(t, snapshot.Items, 1) + require.Equal(t, float64(5), snapshot.Items[0].PricePer1K) + + newLedger := NewToolUsageLedger("gpt-test") + newLedger.BeginAttempt(0) + require.True(t, newLedger.ObserveCustom("lookup", ToolCallIdentity{ + Scope: "openai-chat", + CallID: "call-2", + Position: "choice:0:tool:0", + })) + require.True(t, newLedger.CommitAttempt(0)) + require.Equal(t, float64(9), newLedger.Snapshot().Items[0].PricePer1K) +} + +func TestToolUsageLedgerConcurrentDuplicateObservationsCountOnce(t *testing.T) { + replaceToolPricesForTest(t, map[string]float64{"lookup": 5}) + + ledger := NewToolUsageLedger("gpt-test") + ledger.BeginAttempt(0) + var inserted atomic.Int32 + var workers sync.WaitGroup + for range 32 { + workers.Add(1) + go func() { + defer workers.Done() + if ledger.ObserveCustom("lookup", ToolCallIdentity{ + Scope: "openai-chat", + CallID: "call-1", + Position: "choice:0:tool:0", + }) { + inserted.Add(1) + } + }() + } + workers.Wait() + + require.EqualValues(t, 1, inserted.Load()) + require.True(t, ledger.CommitAttempt(0)) + require.Equal(t, []ToolUsageItem{{ + Name: "lookup", + CallCount: 1, + PricePer1K: 5, + }}, ledger.Snapshot().Items) +} diff --git a/service/alpha_search_billing.go b/service/alpha_search_billing.go index e5bcc0b4..daa0dbb4 100644 --- a/service/alpha_search_billing.go +++ b/service/alpha_search_billing.go @@ -8,7 +8,6 @@ import ( "github.com/MAX-API-Next/MAX-API/dto" 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/setting/operation_setting" "github.com/shopspring/decimal" ) @@ -28,7 +27,7 @@ func AlphaSearchPreConsumeQuota(baseQuota int, info *relaycommon.RelayInfo, grou return 0, fmt.Errorf("alpha search group ratio is invalid: %v", groupRatio) } - pricePerThousand := operation_setting.GetToolPriceForModel(dto.BuildInToolWebSearchPreview, info.OriginModelName) + pricePerThousand := info.FrozenToolPrice(dto.BuildInToolWebSearchPreview) if math.IsNaN(pricePerThousand) || math.IsInf(pricePerThousand, 0) || pricePerThousand < 0 { return 0, fmt.Errorf("alpha search tool price is invalid: %v", pricePerThousand) } diff --git a/service/billing_session.go b/service/billing_session.go index ac1f5538..a3367660 100644 --- a/service/billing_session.go +++ b/service/billing_session.go @@ -1070,6 +1070,13 @@ func NewBillingSession(c *gin.Context, relayInfo *relaycommon.RelayInfo, preCons session, apiErr := trySubscription() if apiErr != nil { if apiErr.GetErrorCode() == types.ErrorCodeInsufficientUserQuota { + allowWalletOverflow, policyErr := model.UserActiveSubscriptionsAllowWalletOverflow(relayInfo.UserId) + if policyErr != nil { + return nil, types.NewError(policyErr, types.ErrorCodeQueryDataError, types.ErrOptionWithSkipRetry()) + } + if !allowWalletOverflow { + return nil, apiErr + } return tryWallet(false) } return nil, apiErr diff --git a/service/billing_session_test.go b/service/billing_session_test.go index 184c0bc2..16931616 100644 --- a/service/billing_session_test.go +++ b/service/billing_session_test.go @@ -1098,6 +1098,72 @@ func TestBillingSessionWalletFirstReplayCannotSwitchFromSubscriptionToWallet(t * assert.EqualValues(t, 1, countSubscriptionPreConsumeRecords(t, "subscription-replay-request")) } +func TestNewBillingSessionSubscriptionFirstHonorsWalletOverflowSnapshot(t *testing.T) { + truncate(t) + const userID, tokenID, planID, subscriptionID = 824, 825, 826, 827 + seedUser(t, userID, 100) + seedToken(t, tokenID, userID, "subscription-strict-token", 100) + allow := false + require.NoError(t, model.DB.Create(&model.SubscriptionPlan{ + Id: planID, Title: "strict-subscription-plan", Enabled: true, + TotalAmount: 5, AllowWalletOverflow: &allow, + QuotaResetPeriod: model.SubscriptionResetNever, + }).Error) + now := time.Now() + require.NoError(t, model.DB.Create(&model.UserSubscription{ + Id: subscriptionID, UserId: userID, PlanId: planID, + AmountTotal: 5, AmountUsed: 5, Status: "active", + StartTime: now.Add(-time.Hour).Unix(), EndTime: now.Add(time.Hour).Unix(), + AllowWalletOverflow: false, + }).Error) + + ctx, _ := gin.CreateTestContext(nil) + info := &relaycommon.RelayInfo{ + RequestId: "subscription-strict-fallback", UserId: userID, + TokenId: tokenID, TokenKey: "subscription-strict-token", + OriginModelName: "subscription-strict-model", + UserSetting: dto.UserSetting{BillingPreference: "subscription_first"}, + } + session, apiErr := NewBillingSession(ctx, info, 10) + require.Nil(t, session) + require.NotNil(t, apiErr) + assert.Equal(t, types.ErrorCodeInsufficientUserQuota, apiErr.GetErrorCode()) + assert.EqualValues(t, 100, getUserQuota(t, userID)) +} + +func TestNewBillingSessionSubscriptionFirstAllowsWalletOverflowWhenAllSnapshotsAllow(t *testing.T) { + truncate(t) + const userID, tokenID, planID, subscriptionID = 828, 829, 830, 831 + seedUser(t, userID, 100) + seedToken(t, tokenID, userID, "subscription-allow-token", 100) + allow := true + require.NoError(t, model.DB.Create(&model.SubscriptionPlan{ + Id: planID, Title: "allow-subscription-plan", Enabled: true, + TotalAmount: 5, AllowWalletOverflow: &allow, + QuotaResetPeriod: model.SubscriptionResetNever, + }).Error) + now := time.Now() + require.NoError(t, model.DB.Create(&model.UserSubscription{ + Id: subscriptionID, UserId: userID, PlanId: planID, + AmountTotal: 5, AmountUsed: 5, Status: "active", + StartTime: now.Add(-time.Hour).Unix(), EndTime: now.Add(time.Hour).Unix(), + AllowWalletOverflow: true, + }).Error) + + ctx, _ := gin.CreateTestContext(nil) + info := &relaycommon.RelayInfo{ + RequestId: "subscription-allow-fallback", UserId: userID, + TokenId: tokenID, TokenKey: "subscription-allow-token", + OriginModelName: "subscription-allow-model", + UserSetting: dto.UserSetting{BillingPreference: "subscription_first"}, + } + session, apiErr := NewBillingSession(ctx, info, 10) + require.Nil(t, apiErr) + require.NotNil(t, session) + assert.Equal(t, BillingSourceWallet, session.funding.Source()) + assert.EqualValues(t, 90, getUserQuota(t, userID)) +} + func TestSettleAndRecordConsumeLogsPostSettlementSubscriptionUsage(t *testing.T) { truncate(t) const userID, tokenID, channelID, planID, subscriptionID = 819, 820, 821, 822, 823 @@ -1197,6 +1263,7 @@ func seedBillingSubscription(t *testing.T, planID int, subscriptionID int, userI Id: subscriptionID, UserId: userID, PlanId: planID, AmountTotal: amountTotal, Status: "active", StartTime: now.Add(-time.Hour).Unix(), EndTime: now.Add(time.Hour).Unix(), + AllowWalletOverflow: true, }).Error) } diff --git a/service/openaicompat/responses_compat_test.go b/service/openaicompat/responses_compat_test.go index 58b72316..49b507eb 100644 --- a/service/openaicompat/responses_compat_test.go +++ b/service/openaicompat/responses_compat_test.go @@ -395,6 +395,49 @@ func TestResponsesResponseToChatCompletionsResponsePreservesUsageDetails(t *test assert.Equal(t, usage.CompletionTokenDetails, resp.Usage.CompletionTokenDetails) } +func TestResponsesResponseToChatCompletionsResponsePreservesCustomToolInput(t *testing.T) { + var upstream dto.OpenAIResponsesResponse + require.NoError(t, common.Unmarshal([]byte(`{ + "id":"resp_custom", + "status":"completed", + "model":"gpt-test", + "output":[{ + "type":"custom_tool_call", + "id":"ctc_1", + "call_id":"call_custom", + "name":"apply_patch", + "input":"*** Begin Patch\n*** End Patch" + }] + }`), &upstream)) + + got, _, err := ResponsesResponseToChatCompletionsResponse(&upstream, "chatcmpl_custom") + require.NoError(t, err) + require.Len(t, got.Choices, 1) + toolCalls := got.Choices[0].Message.ParseToolCalls() + require.Len(t, toolCalls, 1) + assert.Equal(t, "call_custom", toolCalls[0].ID) + assert.Equal(t, "apply_patch", toolCalls[0].Function.Name) + assert.Equal(t, "*** Begin Patch\n*** End Patch", toolCalls[0].Function.Arguments) +} + +func TestResponsesStreamOutputItemPreservesCustomToolInput(t *testing.T) { + var event dto.ResponsesStreamResponse + require.NoError(t, common.Unmarshal([]byte(`{ + "type":"response.output_item.done", + "output_index":0, + "item":{ + "type":"custom_tool_call", + "id":"ctc_1", + "call_id":"call_custom", + "name":"apply_patch", + "input":"patch body" + } + }`), &event)) + + require.NotNil(t, event.Item) + assert.Equal(t, "patch body", event.Item.ArgumentsString()) +} + func TestChatCompletionsStreamToResponsesEventsAggregatesUsage(t *testing.T) { state := NewChatToResponsesStreamState("resp_1", "gpt-test") state.Created = 123 diff --git a/service/quota.go b/service/quota.go index d4443755..13d92af7 100644 --- a/service/quota.go +++ b/service/quota.go @@ -286,6 +286,7 @@ func CalcOpenRouterCacheCreateTokens(usage dto.Usage, priceData types.PriceData) } func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage, extraContent string) { + relayInfo.CommitToolUsageAttempt() var tieredUsedVars map[string]bool if snap := relayInfo.TieredBillingSnapshot; snap != nil { @@ -313,6 +314,7 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, u groupRatio := relayInfo.PriceData.GroupRatioInfo.GroupRatio modelPrice := relayInfo.PriceData.ModelPrice usePrice := relayInfo.PriceData.UsePrice + customToolUsage, customToolSurcharge := calculateCustomToolCallSurcharge(relayInfo, groupRatio) quotaInfo := QuotaInfo{ InputDetails: TokenDetails{ @@ -344,15 +346,28 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, u logContent = fmt.Sprintf("模型价格 %.2f,分组倍率 %.2f", modelPrice, groupRatio) } - // record all the consume log even if quota is 0 - shouldUpdateUsageStats := totalTokens != 0 + // Record all consume logs, including tool-only requests with zero tokens. + shouldUpdateUsageStats := totalTokens != 0 || !customToolSurcharge.IsZero() if totalTokens == 0 { - // in this case, must be some error happened - // we cannot just return, because we may have to return the pre-consumed quota + // Token usage is unavailable, so discard only the token-derived base. + // A completed, priced custom tool call remains independently billable. quota = 0 - logContent += "(可能是上游超时)" - logger.LogError(ctx, fmt.Sprintf("total tokens is 0, cannot consume quota, userId %d, channelId %d, "+ - "tokenId %d, model %s, pre-consumed quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, relayInfo.OriginModelName, relayInfo.FinalPreConsumedQuota)) + if customToolSurcharge.IsZero() { + logContent += "(可能是上游超时)" + logger.LogError(ctx, fmt.Sprintf("total tokens is 0, cannot consume quota, userId %d, channelId %d, "+ + "tokenId %d, model %s, pre-consumed quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, relayInfo.OriginModelName, relayInfo.FinalPreConsumedQuota)) + } else { + logContent += "(上游未返回 token 计费信息,仅结算已完成的自定义工具调用)" + } + } + if !customToolSurcharge.IsZero() { + composedQuota, surchargeClamp := common.QuotaFromDecimalChecked(decimal.NewFromInt(int64(quota)).Add(customToolSurcharge)) + quota = composedQuota + noteQuotaClamp(relayInfo, surchargeClamp) + for _, item := range customToolUsage.Items { + cost := customToolItemQuota(item, groupRatio) + logContent += fmt.Sprintf(", Tool %q called %d time(s), cost %s", item.Name, item.CallCount, cost.String()) + } } if fallbackQuota, ok := streamFallbackQuota(relayInfo, quota); ok { quota = fallbackQuota @@ -369,6 +384,11 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, u if tieredResult != nil { InjectTieredBillingInfo(other, relayInfo, tieredResult) } + if len(customToolUsage.Items) > 0 { + other["tool_calls"] = customToolUsage.Items + other["tool_price_version"] = customToolUsage.PriceVersion + other["tool_billing_model"] = customToolUsage.ModelName + } attachQuotaSaturation(ctx, relayInfo, other) settleAndRecordConsume(ctx, relayInfo, shouldUpdateUsageStats, model.RecordConsumeLogParams{ ChannelId: relayInfo.ChannelId, diff --git a/service/text_quota.go b/service/text_quota.go index 7f2049fd..b24c256f 100644 --- a/service/text_quota.go +++ b/service/text_quota.go @@ -55,6 +55,7 @@ type textQuotaSummary struct { FileSearchCallCount int AudioInputPrice float64 ImageGenerationCallPrice float64 + CustomToolUsage relaycommon.ToolUsageSnapshot ToolCallSurchargeQuota decimal.Decimal } @@ -92,16 +93,43 @@ func isLegacyClaudeDerivedOpenAIUsage(relayInfo *relaycommon.RelayInfo, usage *d return usage.ClaudeCacheCreation5mTokens > 0 || usage.ClaudeCacheCreation1hTokens > 0 } +func customToolItemQuota(item relaycommon.ToolUsageItem, groupRatio float64) decimal.Decimal { + return decimal.NewFromFloat(item.PricePer1K). + Mul(decimal.NewFromInt(int64(item.CallCount))). + Div(decimal.NewFromInt(1000)). + Mul(decimal.NewFromFloat(groupRatio)). + Mul(decimal.NewFromFloat(common.QuotaPerUnit)) +} + +func calculateCustomToolCallSurcharge(relayInfo *relaycommon.RelayInfo, groupRatio float64) (relaycommon.ToolUsageSnapshot, decimal.Decimal) { + if relayInfo == nil || relayInfo.PriceData.UsePrice { + return relaycommon.ToolUsageSnapshot{}, decimal.Zero + } + snapshot := relayInfo.ToolUsageSnapshot() + var surcharge decimal.Decimal + for _, item := range snapshot.Items { + if item.CallCount <= 0 || item.PricePer1K <= 0 { + continue + } + surcharge = surcharge.Add(customToolItemQuota(item, groupRatio)) + } + return snapshot, surcharge +} + func calculateTextToolCallSurcharge(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, summary *textQuotaSummary) decimal.Decimal { + if relayInfo == nil || relayInfo.PriceData.UsePrice { + return decimal.Zero + } dGroupRatio := decimal.NewFromFloat(summary.GroupRatio) dQuotaPerUnit := decimal.NewFromFloat(common.QuotaPerUnit) var surcharge decimal.Decimal + 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 = operation_setting.GetToolPriceForModel("web_search_preview", summary.ModelName) + summary.WebSearchPrice = relayInfo.FrozenToolPrice("web_search_preview") surcharge = surcharge.Add(decimal.NewFromFloat(summary.WebSearchPrice). Mul(decimal.NewFromInt(int64(webSearchTool.CallCount))). Div(decimal.NewFromInt(1000)). @@ -110,7 +138,7 @@ func calculateTextToolCallSurcharge(ctx *gin.Context, relayInfo *relaycommon.Rel } } else if strings.HasSuffix(summary.ModelName, "search-preview") { summary.WebSearchCallCount = 1 - summary.WebSearchPrice = operation_setting.GetToolPriceForModel("web_search_preview", summary.ModelName) + summary.WebSearchPrice = relayInfo.FrozenToolPrice("web_search_preview") surcharge = surcharge.Add(decimal.NewFromFloat(summary.WebSearchPrice). Div(decimal.NewFromInt(1000)). Mul(dGroupRatio). @@ -119,7 +147,7 @@ func calculateTextToolCallSurcharge(ctx *gin.Context, relayInfo *relaycommon.Rel summary.ClaudeWebSearchCallCount = ctx.GetInt("claude_web_search_requests") if summary.ClaudeWebSearchCallCount > 0 { - summary.ClaudeWebSearchPrice = operation_setting.GetToolPrice("web_search") + summary.ClaudeWebSearchPrice = relayInfo.FrozenToolPrice("web_search") surcharge = surcharge.Add(decimal.NewFromFloat(summary.ClaudeWebSearchPrice). Div(decimal.NewFromInt(1000)). Mul(dGroupRatio). @@ -130,7 +158,7 @@ func calculateTextToolCallSurcharge(ctx *gin.Context, relayInfo *relaycommon.Rel if relayInfo.ResponsesUsageInfo != nil { if fileSearchTool, exists := relayInfo.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolFileSearch]; exists && fileSearchTool.CallCount > 0 { summary.FileSearchCallCount = fileSearchTool.CallCount - summary.FileSearchPrice = operation_setting.GetToolPrice("file_search") + summary.FileSearchPrice = relayInfo.FrozenToolPrice("file_search") surcharge = surcharge.Add(decimal.NewFromFloat(summary.FileSearchPrice). Mul(decimal.NewFromInt(int64(fileSearchTool.CallCount))). Div(decimal.NewFromInt(1000)). @@ -385,6 +413,7 @@ func streamFallbackQuota(relayInfo *relaycommon.RelayInfo, quota int) (int, bool } func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage, extraContent []string) { + relayInfo.CommitToolUsageAttempt() originUsage := usage billingUsage := effectiveBillingUsage(usage) if usage == nil { @@ -428,6 +457,11 @@ func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us extraContent = append(extraContent, fmt.Sprintf("Image Generation Call 花费 %s", decimal.NewFromFloat(summary.ImageGenerationCallPrice).Mul(decimal.NewFromFloat(summary.GroupRatio)).Mul(decimal.NewFromFloat(common.QuotaPerUnit)).String())) } + for _, item := range summary.CustomToolUsage.Items { + cost := customToolItemQuota(item, summary.GroupRatio) + extraContent = append(extraContent, fmt.Sprintf("Tool %q called %d time(s), cost %s", item.Name, item.CallCount, cost.String())) + } + shouldUpdateUsageStats := summary.hasBillableUsage(relayInfo) if !summary.hasBillableUsage(relayInfo) { extraContent = append(extraContent, "上游没有返回计费信息,无法扣费(可能是上游超时)") @@ -499,6 +533,11 @@ func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us other["image_generation_call"] = true other["image_generation_call_price"] = summary.ImageGenerationCallPrice } + if len(summary.CustomToolUsage.Items) > 0 { + other["tool_calls"] = summary.CustomToolUsage.Items + other["tool_price_version"] = summary.CustomToolUsage.PriceVersion + other["tool_billing_model"] = summary.CustomToolUsage.ModelName + } if summary.CacheCreationTokens > 0 { other["cache_creation_tokens"] = summary.CacheCreationTokens other["cache_creation_ratio"] = summary.CacheCreationRatio diff --git a/service/text_quota_test.go b/service/text_quota_test.go index 1586d6c6..3fadd88f 100644 --- a/service/text_quota_test.go +++ b/service/text_quota_test.go @@ -988,6 +988,136 @@ func TestCalculateTextQuotaSummaryZeroTokensStillBillsToolSurcharge(t *testing.T require.Equal(t, common.QuotaFromDecimal(summary.ToolCallSurchargeQuota), summary.Quota) } +func TestCustomToolItemQuotaPreservesBillingExpression(t *testing.T) { + item := relaycommon.ToolUsageItem{Name: "lookup", CallCount: 3, PricePer1K: 2.5} + groupRatio := 1.25 + expected := decimal.NewFromFloat(item.PricePer1K). + Mul(decimal.NewFromInt(int64(item.CallCount))). + Div(decimal.NewFromInt(1000)). + Mul(decimal.NewFromFloat(groupRatio)). + Mul(decimal.NewFromFloat(common.QuotaPerUnit)) + + require.True(t, expected.Equal(customToolItemQuota(item, groupRatio))) +} + +func TestCalculateTextQuotaSummaryBillsFrozenCustomToolUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + + setServiceToolPricesForTest(t, map[string]float64{"lookup": 5}) + toolPrices := config.GlobalConfig.Get("tool_price_setting").(*operation_setting.ToolPriceSetting) + + ledger := relaycommon.NewToolUsageLedger("gpt-test") + ledger.BeginAttempt(0) + require.True(t, ledger.ObserveCustom("lookup", relaycommon.ToolCallIdentity{ + Scope: "openai-chat", + CallID: "call-1", + Position: "choice:0:tool:0", + })) + require.True(t, ledger.CommitAttempt(0)) + + // A config reload after the request starts must not reprice the call. + toolPrices.Prices["lookup"] = 9 + operation_setting.RebuildToolPriceIndex() + + relayInfo := &relaycommon.RelayInfo{ + OriginModelName: "gpt-test", + PriceData: types.PriceData{ + ModelRatio: 1, + CompletionRatio: 1, + GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 1}, + }, + ToolUsage: ledger, + StartTime: time.Now(), + } + + summary := calculateTextQuotaSummary(ctx, relayInfo, &dto.Usage{}) + expected := decimal.NewFromFloat(5.0 / 1000).Mul(decimal.NewFromFloat(common.QuotaPerUnit)) + + require.Equal(t, 0, summary.TotalTokens) + require.True(t, expected.Equal(summary.ToolCallSurchargeQuota)) + require.Equal(t, common.QuotaFromDecimal(expected), summary.Quota) +} + +func TestCalculateTextQuotaSummaryBillsFrozenBuiltInToolPrices(t *testing.T) { + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set("claude_web_search_requests", 2) + + setServiceToolPricesForTest(t, map[string]float64{ + "web_search_preview:gpt-4o*": 7, + "web_search": 11, + "file_search": 13, + }) + toolPrices := config.GlobalConfig.Get("tool_price_setting").(*operation_setting.ToolPriceSetting) + ledger := relaycommon.NewToolUsageLedger("gpt-4o") + + relayInfo := &relaycommon.RelayInfo{ + OriginModelName: "gpt-4o", + PriceData: types.PriceData{ + ModelRatio: 1, + CompletionRatio: 1, + GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 1}, + }, + ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{ + BuiltInTools: map[string]*relaycommon.BuildInToolInfo{ + dto.BuildInToolWebSearchPreview: {CallCount: 1}, + dto.BuildInToolFileSearch: {CallCount: 2}, + }, + }, + ToolUsage: ledger, + StartTime: time.Now(), + } + + // The request captures its price index before the admin reload below. + toolPrices.Prices["web_search_preview:gpt-4o*"] = 70 + toolPrices.Prices["web_search"] = 110 + toolPrices.Prices["file_search"] = 130 + operation_setting.RebuildToolPriceIndex() + + summary := calculateTextQuotaSummary(ctx, relayInfo, &dto.Usage{}) + require.Equal(t, float64(7), summary.WebSearchPrice) + require.Equal(t, float64(11), summary.ClaudeWebSearchPrice) + require.Equal(t, float64(13), summary.FileSearchPrice) + require.Equal(t, 1, summary.WebSearchCallCount) + require.Equal(t, 2, summary.ClaudeWebSearchCallCount) + require.Equal(t, 2, summary.FileSearchCallCount) + + expected := decimal.NewFromFloat(7 + 11*2 + 13*2). + Div(decimal.NewFromInt(1000)). + Mul(decimal.NewFromFloat(common.QuotaPerUnit)) + require.True(t, expected.Equal(summary.ToolCallSurchargeQuota)) +} + +func TestAlphaSearchPreConsumeQuotaUsesFrozenBuiltInToolPrice(t *testing.T) { + setServiceToolPricesForTest(t, map[string]float64{"web_search_preview": 7}) + toolPrices := config.GlobalConfig.Get("tool_price_setting").(*operation_setting.ToolPriceSetting) + info := &relaycommon.RelayInfo{ + RelayMode: relayconstant.RelayModeAlphaSearch, + OriginModelName: "o1", + ToolUsage: relaycommon.NewToolUsageLedger("o1"), + PriceData: types.PriceData{ + GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 1}, + }, + ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{ + BuiltInTools: map[string]*relaycommon.BuildInToolInfo{ + dto.BuildInToolWebSearchPreview: {CallCount: 1}, + }, + }, + } + + preConsumed, err := AlphaSearchPreConsumeQuota(0, info, 1) + require.NoError(t, err) + + toolPrices.Prices["web_search_preview"] = 70 + operation_setting.RebuildToolPriceIndex() + + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + summary := calculateTextQuotaSummary(ctx, info, &dto.Usage{}) + require.Equal(t, preConsumed, summary.Quota) + require.Equal(t, float64(7), summary.WebSearchPrice) +} + func TestCalculateTextQuotaSummaryAlphaSearchKeepsFixedPriceWhenToolPriceIsZero(t *testing.T) { gin.SetMode(gin.TestMode) ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) diff --git a/service/tool_usage_billing_test.go b/service/tool_usage_billing_test.go new file mode 100644 index 00000000..831d48b1 --- /dev/null +++ b/service/tool_usage_billing_test.go @@ -0,0 +1,218 @@ +package service + +import ( + "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" + relaycommon "github.com/MAX-API-Next/MAX-API/relay/common" + "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/shopspring/decimal" + "github.com/stretchr/testify/require" +) + +func setServiceToolPricesForTest(t *testing.T, additions map[string]float64) { + t.Helper() + + setting := config.GlobalConfig.Get("tool_price_setting").(*operation_setting.ToolPriceSetting) + original := make(map[string]float64, len(setting.Prices)) + for name, price := range setting.Prices { + original[name] = price + } + setting.Prices = make(map[string]float64, len(original)+len(additions)) + for name, price := range original { + setting.Prices[name] = price + } + for name, price := range additions { + setting.Prices[name] = price + } + operation_setting.RebuildToolPriceIndex() + + t.Cleanup(func() { + setting.Prices = original + operation_setting.RebuildToolPriceIndex() + }) +} + +func newToolSettlementInfo(t *testing.T, requestID string, userID, tokenID, channelID int, tokenKey string, initialQuota, preConsumedQuota int) (*gin.Context, *relaycommon.RelayInfo) { + t.Helper() + truncate(t) + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set(common.RequestIdKey, requestID) + ctx.Set("token_quota", int64(initialQuota)) + + seedUser(t, userID, initialQuota) + seedToken(t, tokenID, userID, tokenKey, initialQuota) + seedChannel(t, channelID) + + info := &relaycommon.RelayInfo{ + RequestId: requestID, + UserId: userID, + TokenId: tokenID, + TokenKey: tokenKey, + OriginModelName: "gpt-test", + StartTime: time.Now(), + UserSetting: dto.UserSetting{BillingPreference: "wallet_only"}, + FinalRequestRelayFormat: types.RelayFormatOpenAI, + UsingGroup: "default", + ForcePreConsume: true, + ChannelMeta: &relaycommon.ChannelMeta{ChannelId: channelID}, + PriceData: types.PriceData{ + ModelRatio: 1, + CompletionRatio: 1, + GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 1}, + }, + ToolUsage: relaycommon.NewToolUsageLedger("gpt-test"), + } + session, apiErr := NewBillingSession(ctx, info, preConsumedQuota) + require.Nil(t, apiErr) + require.NotNil(t, session) + info.Billing = session + return ctx, info +} + +func assertSingleToolSettlement(t *testing.T, requestID string, userID, tokenID, channelID, initialQuota, expectedQuota int, expectedTool string) { + t.Helper() + require.EqualValues(t, initialQuota-expectedQuota, getUserQuota(t, userID)) + require.Equal(t, initialQuota-expectedQuota, getTokenRemainQuota(t, tokenID)) + + var user model.User + require.NoError(t, model.DB.Select("used_quota", "request_count").First(&user, userID).Error) + require.EqualValues(t, expectedQuota, user.UsedQuota) + require.Equal(t, 1, user.RequestCount) + var channel model.Channel + require.NoError(t, model.DB.Select("used_quota").First(&channel, channelID).Error) + require.EqualValues(t, expectedQuota, channel.UsedQuota) + require.Equal(t, int64(1), countLogs(t)) + + log := getLastLog(t) + require.Equal(t, expectedQuota, log.Quota) + require.Equal(t, requestID, log.RequestId) + var other map[string]any + require.NoError(t, common.Unmarshal([]byte(log.Other), &other)) + require.Equal(t, "gpt-test", other["tool_billing_model"]) + require.NotZero(t, other["tool_price_version"]) + items, ok := other["tool_calls"].([]any) + require.True(t, ok) + require.Len(t, items, 1) + item, ok := items[0].(map[string]any) + require.True(t, ok) + require.Equal(t, expectedTool, item["name"]) + require.Equal(t, float64(1), item["call_count"]) +} + +func TestPostTextConsumeQuotaZeroTokensKeepsOnlyWinningRetryToolAndSettlesOnce(t *testing.T) { + setServiceToolPricesForTest(t, map[string]float64{"failed_tool": 5, "winning_tool": 7}) + const requestID = "zero-token-tool-retry-request" + const userID, tokenID, channelID = 701, 702, 703 + const initialQuota, preConsumedQuota = 100_000, 100 + ctx, info := newToolSettlementInfo(t, requestID, userID, tokenID, channelID, "zero-token-tool-retry-key", initialQuota, preConsumedQuota) + + info.RetryIndex = 0 + info.ToolUsage.BeginAttempt(0) + require.True(t, info.ObserveCustomToolCall("failed_tool", relaycommon.ToolCallIdentity{Scope: "openai-chat", CallID: "failed-call"})) + info.RetryIndex = 1 + info.ToolUsage.BeginAttempt(1) + require.True(t, info.ObserveCustomToolCall("winning_tool", relaycommon.ToolCallIdentity{Scope: "openai-chat", CallID: "winning-call"})) + + expectedQuota := common.QuotaFromDecimal(decimal.NewFromFloat(7).Div(decimal.NewFromInt(1000)).Mul(decimal.NewFromFloat(common.QuotaPerUnit))) + PostTextConsumeQuota(ctx, info, &dto.Usage{}, nil) + assertSingleToolSettlement(t, requestID, userID, tokenID, channelID, initialQuota, expectedQuota, "winning_tool") + + PostTextConsumeQuota(ctx, info, &dto.Usage{}, nil) + assertSingleToolSettlement(t, requestID, userID, tokenID, channelID, initialQuota, expectedQuota, "winning_tool") +} + +func TestPostAudioConsumeQuotaZeroTokensBillsCustomToolOnce(t *testing.T) { + setServiceToolPricesForTest(t, map[string]float64{"transcribe_lookup": 9}) + const requestID = "zero-token-audio-tool-request" + const userID, tokenID, channelID = 711, 712, 713 + const initialQuota, preConsumedQuota = 100_000, 100 + ctx, info := newToolSettlementInfo(t, requestID, userID, tokenID, channelID, "zero-token-audio-tool-key", initialQuota, preConsumedQuota) + + info.ToolUsage.BeginAttempt(0) + require.True(t, info.ObserveCustomToolCall("transcribe_lookup", relaycommon.ToolCallIdentity{Scope: "openai-responses", CallID: "call-1"})) + + expectedQuota := common.QuotaFromDecimal(decimal.NewFromFloat(9).Div(decimal.NewFromInt(1000)).Mul(decimal.NewFromFloat(common.QuotaPerUnit))) + PostAudioConsumeQuota(ctx, info, &dto.Usage{}, "") + assertSingleToolSettlement(t, requestID, userID, tokenID, channelID, initialQuota, expectedQuota, "transcribe_lookup") + + PostAudioConsumeQuota(ctx, info, &dto.Usage{}, "") + assertSingleToolSettlement(t, requestID, userID, tokenID, channelID, initialQuota, expectedQuota, "transcribe_lookup") +} + +func TestCalculateTextToolCallSurchargeSkipsPerRequestModels(t *testing.T) { + setServiceToolPricesForTest(t, map[string]float64{"lookup": 9}) + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + info := &relaycommon.RelayInfo{ + OriginModelName: "gpt-test", + PriceData: types.PriceData{UsePrice: true}, + ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{BuiltInTools: map[string]*relaycommon.BuildInToolInfo{ + dto.BuildInToolWebSearchPreview: {ToolName: dto.BuildInToolWebSearchPreview, CallCount: 1}, + }}, + ToolUsage: relaycommon.NewToolUsageLedger("gpt-test"), + } + info.ToolUsage.BeginAttempt(0) + require.True(t, info.ObserveCustomToolCall("lookup", relaycommon.ToolCallIdentity{Scope: "openai-chat", CallID: "call-1"})) + require.True(t, info.CommitToolUsageAttempt()) + + summary := textQuotaSummary{ModelName: "gpt-test", GroupRatio: 1} + require.True(t, calculateTextToolCallSurcharge(ctx, info, &summary).IsZero()) + require.Empty(t, summary.CustomToolUsage.Items) +} + +func TestCalculateTextQuotaSummaryBillsProjectedResponsesBuiltInToolOnce(t *testing.T) { + setServiceToolPricesForTest(t, map[string]float64{ + dto.BuildInToolWebSearchPreview: 5, + dto.BuildInToolFileSearch: 7, + }) + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + info := &relaycommon.RelayInfo{ + OriginModelName: "gpt-test", + PriceData: types.PriceData{ + ModelRatio: 1, + CompletionRatio: 1, + GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 1}, + }, + ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{BuiltInTools: map[string]*relaycommon.BuildInToolInfo{ + dto.BuildInToolWebSearchPreview: {ToolName: dto.BuildInToolWebSearchPreview}, + dto.BuildInToolFileSearch: {ToolName: dto.BuildInToolFileSearch}, + }}, + ToolUsage: relaycommon.NewToolUsageLedger("gpt-test"), + StartTime: time.Now(), + } + info.ToolUsage.BeginAttempt(0) + require.True(t, info.ObserveBuiltInToolCall(dto.BuildInToolWebSearchPreview, relaycommon.ToolCallIdentity{ + Scope: "openai-responses", + CallID: "web-1", + Position: "output:0", + })) + require.False(t, info.ObserveBuiltInToolCall(dto.BuildInToolWebSearchPreview, relaycommon.ToolCallIdentity{ + Scope: "openai-responses", + CallID: "web-1", + Position: "output:0", + })) + require.True(t, info.ObserveBuiltInToolCall(dto.BuildInToolFileSearch, relaycommon.ToolCallIdentity{ + Scope: "openai-responses", + CallID: "file-1", + Position: "output:1", + })) + + summary := calculateTextQuotaSummary(ctx, info, &dto.Usage{}) + expected := decimal.NewFromFloat(5 + 7). + Div(decimal.NewFromInt(1000)). + Mul(decimal.NewFromFloat(common.QuotaPerUnit)) + require.Equal(t, 1, summary.WebSearchCallCount) + require.Equal(t, 1, summary.FileSearchCallCount) + require.True(t, expected.Equal(summary.ToolCallSurchargeQuota)) +} diff --git a/setting/operation_setting/tools.go b/setting/operation_setting/tools.go index d9c5c50c..d4558f81 100644 --- a/setting/operation_setting/tools.go +++ b/setting/operation_setting/tools.go @@ -66,11 +66,13 @@ type prefixEntry struct { } type toolPriceIndex struct { + version uint64 defaults map[string]float64 prefixes map[string][]prefixEntry } var currentIndex atomic.Pointer[toolPriceIndex] +var currentToolPriceVersion atomic.Uint64 // RebuildToolPriceIndex rebuilds the lookup index from the current config. // Called on init and after config updates. Not on the billing hot path. @@ -87,6 +89,7 @@ func RebuildToolPriceIndex() { } idx := &toolPriceIndex{ + version: currentToolPriceVersion.Add(1), defaults: make(map[string]float64), prefixes: make(map[string][]prefixEntry), } @@ -114,10 +117,38 @@ func RebuildToolPriceIndex() { currentIndex.Store(idx) } -// GetToolPriceForModel returns the price ($/1K calls) for a tool given a model name. -// Lookup: longest prefix match → tool default → 0. -func GetToolPriceForModel(toolName, modelName string) float64 { - idx := currentIndex.Load() +// ToolPriceSnapshot is an immutable request-scoped view of tool prices. It +// retains the exact index that was active when the request started so config +// reloads cannot change the price of an in-flight tool call. +type ToolPriceSnapshot struct { + modelName string + index *toolPriceIndex +} + +// CaptureToolPriceSnapshot freezes the current price index for one billing +// model. Published indexes are immutable after currentIndex.Store. +func CaptureToolPriceSnapshot(modelName string) ToolPriceSnapshot { + return ToolPriceSnapshot{ + modelName: modelName, + index: currentIndex.Load(), + } +} + +// Version returns the monotonically increasing config generation captured by +// this snapshot. Zero means the index was unavailable during early startup. +func (s ToolPriceSnapshot) Version() uint64 { + if s.index == nil { + return 0 + } + return s.index.version +} + +// PriceFor returns the frozen price ($/1K calls) for a tool. +func (s ToolPriceSnapshot) PriceFor(toolName string) float64 { + return lookupToolPrice(s.index, toolName, s.modelName) +} + +func lookupToolPrice(idx *toolPriceIndex, toolName, modelName string) float64 { if idx == nil { if v, ok := defaultToolPrices[toolName]; ok { return v @@ -139,6 +170,12 @@ func GetToolPriceForModel(toolName, modelName string) float64 { return 0 } +// GetToolPriceForModel returns the price ($/1K calls) for a tool given a model name. +// Lookup: longest prefix match → tool default → 0. +func GetToolPriceForModel(toolName, modelName string) float64 { + return lookupToolPrice(currentIndex.Load(), toolName, modelName) +} + // GetToolPrice is a convenience wrapper when no model name is needed. func GetToolPrice(toolName string) float64 { return GetToolPriceForModel(toolName, "") diff --git a/web/default/src/features/subscriptions/components/subscriptions-columns.tsx b/web/default/src/features/subscriptions/components/subscriptions-columns.tsx index 5c6ea73d..f76a4de2 100644 --- a/web/default/src/features/subscriptions/components/subscriptions-columns.tsx +++ b/web/default/src/features/subscriptions/components/subscriptions-columns.tsx @@ -207,6 +207,18 @@ export function useSubscriptionsColumns(): ColumnDef[] { }, size: 100, }, + { + id: 'allow_wallet_overflow', + meta: { label: t('Wallet fallback'), mobileHidden: true }, + header: ({ column }) => ( + + ), + cell: ({ row }) => + row.original.plan.allow_wallet_overflow === false + ? t('Disable') + : t('Enable'), + size: 110, + }, { id: 'actions', cell: ({ row }) => , diff --git a/web/default/src/features/subscriptions/components/subscriptions-mutate-drawer.tsx b/web/default/src/features/subscriptions/components/subscriptions-mutate-drawer.tsx index c973250b..509ba583 100644 --- a/web/default/src/features/subscriptions/components/subscriptions-mutate-drawer.tsx +++ b/web/default/src/features/subscriptions/components/subscriptions-mutate-drawer.tsx @@ -423,6 +423,47 @@ export function SubscriptionsMutateDrawer({ /> +
+ ( + + {t('Downgrade Group')} + + + + )} + /> +
+
)} /> + + ( + + + {t('Allow wallet fallback after quota exhaustion')} + + + + + + )} + />
diff --git a/web/default/src/features/subscriptions/lib/plan-form.test.ts b/web/default/src/features/subscriptions/lib/plan-form.test.ts new file mode 100644 index 00000000..ceda0d0d --- /dev/null +++ b/web/default/src/features/subscriptions/lib/plan-form.test.ts @@ -0,0 +1,40 @@ +import assert from 'node:assert/strict' +import { describe, test } from 'node:test' +import { + PLAN_FORM_DEFAULTS, + formValuesToPlanPayload, + planToFormValues, +} from './plan-form' + +describe('subscription plan policy fields', () => { + test('defaults wallet overflow to enabled and keeps explicit false', () => { + assert.equal(PLAN_FORM_DEFAULTS.allow_wallet_overflow, true) + + const values = planToFormValues({ + id: 1, + title: 'Strict', + subtitle: '', + price_amount: 1, + currency: 'USD', + duration_unit: 'month', + duration_value: 1, + custom_seconds: 0, + quota_reset_period: 'never', + quota_reset_custom_seconds: 0, + enabled: true, + sort_order: 0, + allow_balance_pay: true, + allow_wallet_overflow: false, + max_purchase_per_user: 0, + total_amount: 100, + upgrade_group: 'svip', + downgrade_group: 'default', + }) + assert.equal(values.allow_wallet_overflow, false) + assert.equal(values.downgrade_group, 'default') + + const payload = formValuesToPlanPayload(values) + assert.equal(payload.plan.allow_wallet_overflow, false) + assert.equal(payload.plan.downgrade_group, 'default') + }) +}) diff --git a/web/default/src/features/subscriptions/lib/plan-form.ts b/web/default/src/features/subscriptions/lib/plan-form.ts index 44bbde06..546935af 100644 --- a/web/default/src/features/subscriptions/lib/plan-form.ts +++ b/web/default/src/features/subscriptions/lib/plan-form.ts @@ -40,9 +40,11 @@ export function getPlanFormSchema(t: TFunction) { enabled: z.boolean(), sort_order: z.coerce.number(), allow_balance_pay: z.boolean(), + allow_wallet_overflow: z.boolean(), max_purchase_per_user: z.coerce.number().min(0), total_amount: z.coerce.number().min(0), upgrade_group: z.string().optional(), + downgrade_group: z.string().optional(), stripe_price_id: z.string().optional(), creem_product_id: z.string().optional(), waffo_pancake_product_id: z.string().optional(), @@ -63,9 +65,11 @@ export const PLAN_FORM_DEFAULTS: PlanFormValues = { enabled: true, sort_order: 0, allow_balance_pay: true, + allow_wallet_overflow: true, max_purchase_per_user: 0, total_amount: 0, upgrade_group: '', + downgrade_group: '', stripe_price_id: '', creem_product_id: '', waffo_pancake_product_id: '', @@ -84,9 +88,11 @@ export function planToFormValues(plan: SubscriptionPlan): PlanFormValues { enabled: plan.enabled !== false, sort_order: Number(plan.sort_order || 0), allow_balance_pay: plan.allow_balance_pay !== false, + allow_wallet_overflow: plan.allow_wallet_overflow !== false, max_purchase_per_user: Number(plan.max_purchase_per_user || 0), total_amount: quotaUnitsToDollars(Number(plan.total_amount || 0)), upgrade_group: plan.upgrade_group || '', + downgrade_group: plan.downgrade_group || '', stripe_price_id: plan.stripe_price_id || '', creem_product_id: plan.creem_product_id || '', waffo_pancake_product_id: plan.waffo_pancake_product_id || '', @@ -110,6 +116,7 @@ export function formValuesToPlanPayload(values: PlanFormValues): PlanPayload { max_purchase_per_user: Number(values.max_purchase_per_user || 0), total_amount: parseQuotaFromDollars(Number(values.total_amount || 0)), upgrade_group: values.upgrade_group || '', + downgrade_group: values.downgrade_group || '', }, } } diff --git a/web/default/src/features/subscriptions/types.ts b/web/default/src/features/subscriptions/types.ts index 8d6a289e..bba1a0cf 100644 --- a/web/default/src/features/subscriptions/types.ts +++ b/web/default/src/features/subscriptions/types.ts @@ -36,9 +36,11 @@ export const subscriptionPlanSchema = z.object({ enabled: z.boolean(), sort_order: z.number(), allow_balance_pay: z.boolean().optional().default(true), + allow_wallet_overflow: z.boolean().optional().default(true), max_purchase_per_user: z.number(), total_amount: z.number(), upgrade_group: z.string().optional(), + downgrade_group: z.string().optional(), stripe_price_id: z.string().optional(), creem_product_id: z.string().optional(), waffo_pancake_product_id: z.string().optional(), @@ -65,6 +67,8 @@ export const userSubscriptionSchema = z.object({ amount_total: z.number(), amount_used: z.number(), next_reset_time: z.number().optional(), + allow_wallet_overflow: z.boolean().optional().default(true), + downgrade_group: z.string().optional(), }) export type UserSubscription = z.infer diff --git a/web/default/src/features/wallet/components/subscription-plans-card.tsx b/web/default/src/features/wallet/components/subscription-plans-card.tsx index f00a1af0..d3cc0315 100644 --- a/web/default/src/features/wallet/components/subscription-plans-card.tsx +++ b/web/default/src/features/wallet/components/subscription-plans-card.tsx @@ -534,6 +534,12 @@ export function SubscriptionPlansCard({ plan.upgrade_group ? `${t('Upgrade Group')}: ${plan.upgrade_group}` : null, + plan.downgrade_group + ? `${t('Downgrade Group')}: ${plan.downgrade_group}` + : null, + plan.allow_wallet_overflow === false + ? t('Wallet fallback disabled') + : null, ].filter(Boolean) as string[] return ( diff --git a/web/default/src/i18n/locales/en.json b/web/default/src/i18n/locales/en.json index c48ff705..6b71b8a4 100644 --- a/web/default/src/i18n/locales/en.json +++ b/web/default/src/i18n/locales/en.json @@ -394,6 +394,7 @@ "Allocated Memory": "Allocated Memory", "Allow accountFilter parameter": "Allow accountFilter parameter", "Allow balance redemption": "Allow balance redemption", + "Allow wallet fallback after quota exhaustion": "Allow wallet fallback after quota exhaustion", "Allow Claude beta query passthrough": "Allow Claude beta query passthrough", "Allow clients to query configured ratios via `/api/ratio`.": "Allow clients to query configured ratios via `/api/ratio`.", "Allow HTTP image requests": "Allow HTTP image requests", @@ -3128,6 +3129,7 @@ "No token found.": "No token found.", "No tools configured": "No tools configured", "No Upgrade": "No Upgrade", + "No Downgrade": "No Downgrade", "No upstream price differences found": "No upstream price differences found", "No upstream ratio differences found": "No upstream ratio differences found", "No uptime data available": "No uptime data available", @@ -5020,6 +5022,9 @@ "Updating all channel balances. This may take a while. Please refresh to see results.": "Updating all channel balances. This may take a while. Please refresh to see results.", "Updating...": "Updating...", "Upgrade Group": "Upgrade Group", + "Downgrade Group": "Downgrade Group", + "Wallet fallback": "Wallet fallback", + "Wallet fallback disabled": "Wallet fallback disabled", "Upload": "Upload", "Upload a single service account JSON file": "Upload a single service account JSON file", "Upload file": "Upload file", diff --git a/web/default/src/i18n/locales/fr.json b/web/default/src/i18n/locales/fr.json index 2567ed69..b610c722 100644 --- a/web/default/src/i18n/locales/fr.json +++ b/web/default/src/i18n/locales/fr.json @@ -394,6 +394,7 @@ "Allocated Memory": "Mémoire allouée", "Allow accountFilter parameter": "Autoriser le paramètre accountFilter", "Allow balance redemption": "Autoriser le paiement avec le solde", + "Allow wallet fallback after quota exhaustion": "Autoriser le portefeuille lorsque le quota est épuisé", "Allow Claude beta query passthrough": "Autoriser le passage des requêtes bêta Claude", "Allow clients to query configured ratios via `/api/ratio`.": "Autoriser les clients à interroger les ratios configurés via `/api/ratio`.", "Allow HTTP image requests": "Autoriser les requêtes d'images HTTP", @@ -3128,6 +3129,7 @@ "No token found.": "Aucun jeton trouvé.", "No tools configured": "Aucun outil configuré", "No Upgrade": "Pas de mise à niveau", + "No Downgrade": "Pas de rétrogradation", "No upstream price differences found": "Aucune différence de prix amont trouvée", "No upstream ratio differences found": "Aucune différence de ratio amont trouvée", "No uptime data available": "Aucune donnée de disponibilité disponible", @@ -5020,6 +5022,9 @@ "Updating all channel balances. This may take a while. Please refresh to see results.": "Mise à jour de tous les soldes des canaux. Cela peut prendre un certain temps. Veuillez actualiser pour voir les résultats.", "Updating...": "Mise à jour...", "Upgrade Group": "Groupe de mise à niveau", + "Downgrade Group": "Groupe de rétrogradation", + "Wallet fallback": "Portefeuille de secours", + "Wallet fallback disabled": "Portefeuille de secours désactivé", "Upload": "Téléverser", "Upload a single service account JSON file": "Télécharger un seul fichier JSON de compte de service", "Upload file": "Téléverser un fichier", diff --git a/web/default/src/i18n/locales/ja.json b/web/default/src/i18n/locales/ja.json index e7ee2add..e3daf6f5 100644 --- a/web/default/src/i18n/locales/ja.json +++ b/web/default/src/i18n/locales/ja.json @@ -394,6 +394,7 @@ "Allocated Memory": "割り当て済みメモリ", "Allow accountFilter parameter": "accountFilter パラメータを許可", "Allow balance redemption": "残高での交換を許可", + "Allow wallet fallback after quota exhaustion": "クォータ枯渇後のウォレットフォールバックを許可", "Allow Claude beta query passthrough": "Claude ベータクエリのパススルーを許可", "Allow clients to query configured ratios via `/api/ratio`.": "クライアントが `/api/ratio` 経由で設定された比率を照会できるようにします。", "Allow HTTP image requests": "HTTP画像リクエストを許可", @@ -3128,6 +3129,7 @@ "No token found.": "トークンが見つかりません。", "No tools configured": "ツールが未設定です", "No Upgrade": "アップグレードなし", + "No Downgrade": "ダウングレードなし", "No upstream price differences found": "上流価格の差異は見つかりませんでした", "No upstream ratio differences found": "アップストリームの比率の差は見つかりません", "No uptime data available": "稼働率データがありません", @@ -5020,6 +5022,9 @@ "Updating all channel balances. This may take a while. Please refresh to see results.": "すべてのチャネル残高を更新中です。これには少し時間がかかる場合があります。結果を確認するには更新してください。", "Updating...": "更新中...", "Upgrade Group": "グループをアップグレード", + "Downgrade Group": "ダウングレードグループ", + "Wallet fallback": "ウォレットフォールバック", + "Wallet fallback disabled": "ウォレットフォールバック無効", "Upload": "アップロード", "Upload a single service account JSON file": "単一のサービスアカウントJSONファイルをアップロードする", "Upload file": "ファイルをアップロード", diff --git a/web/default/src/i18n/locales/ru.json b/web/default/src/i18n/locales/ru.json index 32b4e169..5209cca7 100644 --- a/web/default/src/i18n/locales/ru.json +++ b/web/default/src/i18n/locales/ru.json @@ -394,6 +394,7 @@ "Allocated Memory": "Выделенная память", "Allow accountFilter parameter": "Разрешить параметр accountFilter", "Allow balance redemption": "Разрешить оплату балансом", + "Allow wallet fallback after quota exhaustion": "Разрешить оплату балансом после исчерпания квоты", "Allow Claude beta query passthrough": "Разрешить проброс бета-запросов Claude", "Allow clients to query configured ratios via `/api/ratio`.": "Разрешить клиентам запрашивать настроенные соотношения через `/api/ratio`.", "Allow HTTP image requests": "Разрешить HTTP-запросы изображений", @@ -3128,6 +3129,7 @@ "No token found.": "Токен не найден.", "No tools configured": "Нет настроенных инструментов", "No Upgrade": "Без повышения", + "No Downgrade": "Без понижения", "No upstream price differences found": "Различий в ценах провайдера не найдено", "No upstream ratio differences found": "Различия в коэффициентах вышестоящего потока не найдены", "No uptime data available": "Данные о доступности недоступны", @@ -5020,6 +5022,9 @@ "Updating all channel balances. This may take a while. Please refresh to see results.": "Обновление балансов всех каналов. Это может занять некоторое время. Пожалуйста, обновите страницу, чтобы увидеть результаты.", "Updating...": "Обновление...", "Upgrade Group": "Повысить группу", + "Downgrade Group": "Понижающая группа", + "Wallet fallback": "Резервная оплата с баланса", + "Wallet fallback disabled": "Резервная оплата с баланса отключена", "Upload": "Загрузка", "Upload a single service account JSON file": "Загрузите JSON-файл одного сервисного аккаунта", "Upload file": "Загрузить файл", diff --git a/web/default/src/i18n/locales/vi.json b/web/default/src/i18n/locales/vi.json index 1ac3ccd5..c15d3b9a 100644 --- a/web/default/src/i18n/locales/vi.json +++ b/web/default/src/i18n/locales/vi.json @@ -394,6 +394,7 @@ "Allocated Memory": "Bộ nhớ đã cấp phát", "Allow accountFilter parameter": "Cho phép tham số accountFilter", "Allow balance redemption": "Cho phép thanh toán bằng số dư", + "Allow wallet fallback after quota exhaustion": "Cho phép dùng ví dự phòng khi hết hạn mức", "Allow Claude beta query passthrough": "Cho phép chuyển tiếp truy vấn beta Claude", "Allow clients to query configured ratios via `/api/ratio`.": "Cho phép khách hàng truy vấn các tỷ lệ đã cấu hình thông qua `/api/ratio`.", "Allow HTTP image requests": "Cho phép yêu cầu hình ảnh HTTP", @@ -3128,6 +3129,7 @@ "No token found.": "Không tìm thấy mã thông báo.", "No tools configured": "Chưa cấu hình công cụ nào", "No Upgrade": "Không nâng cấp", + "No Downgrade": "Không hạ cấp", "No upstream price differences found": "Không tìm thấy sự khác biệt về giá upstream", "No upstream ratio differences found": "Không tìm thấy chênh lệch tỷ lệ thượng nguồn", "No uptime data available": "Không có dữ liệu thời gian hoạt động", @@ -5020,6 +5022,9 @@ "Updating all channel balances. This may take a while. Please refresh to see results.": "Đang cập nhật tất cả số dư kênh. Quá trình này có thể mất một chút thời gian. Vui lòng làm mới để xem kết quả.", "Updating...": "Đang cập nhật...", "Upgrade Group": "Nhóm nâng cấp", + "Downgrade Group": "Nhóm hạ cấp", + "Wallet fallback": "Ví dự phòng", + "Wallet fallback disabled": "Đã tắt ví dự phòng", "Upload": "Tải lên", "Upload a single service account JSON file": "Tải lên một tệp JSON tài khoản dịch vụ", "Upload file": "Tải tệp lên", diff --git a/web/default/src/i18n/locales/zh.json b/web/default/src/i18n/locales/zh.json index ac7bb32f..94ccab25 100644 --- a/web/default/src/i18n/locales/zh.json +++ b/web/default/src/i18n/locales/zh.json @@ -394,6 +394,7 @@ "Allocated Memory": "已分配内存", "Allow accountFilter parameter": "允许 accountFilter 参数", "Allow balance redemption": "允许余额兑换", + "Allow wallet fallback after quota exhaustion": "订阅额度用尽后允许钱包兜底", "Allow Claude beta query passthrough": "允许 Claude beta 查询透传", "Allow clients to query configured ratios via `/api/ratio`.": "允许客户端通过 `/api/ratio` 查询配置的比例。", "Allow HTTP image requests": "允许 HTTP 图像请求", @@ -3128,6 +3129,7 @@ "No token found.": "未找到令牌。", "No tools configured": "未配置工具", "No Upgrade": "不升级", + "No Downgrade": "不降级", "No upstream price differences found": "未发现上游价格差异", "No upstream ratio differences found": "未找到上游比例差异", "No uptime data available": "暂无可用率数据", @@ -5020,6 +5022,9 @@ "Updating all channel balances. This may take a while. Please refresh to see results.": "正在更新所有渠道余额。这可能需要一段时间。请刷新以查看结果。", "Updating...": "正在更新...", "Upgrade Group": "升级分组", + "Downgrade Group": "降级分组", + "Wallet fallback": "钱包兜底", + "Wallet fallback disabled": "不允许钱包兜底", "Upload": "上传", "Upload a single service account JSON file": "上传单个服务账号 JSON 文件", "Upload file": "上传文件", diff --git a/web/default/src/i18n/static-keys.ts b/web/default/src/i18n/static-keys.ts index bd254b57..1b0eb001 100644 --- a/web/default/src/i18n/static-keys.ts +++ b/web/default/src/i18n/static-keys.ts @@ -300,6 +300,11 @@ export const STATIC_I18N_KEYS = [ 'Received', 'Payment initiated', 'Payment page opened', + 'Allow wallet fallback after quota exhaustion', + 'Downgrade Group', + 'No Downgrade', + 'Wallet fallback', + 'Wallet fallback disabled', // Upstream model updates 'Upstream Model Updates',