feat: publish sanitized custom.97 source

This commit is contained in:
cneicy committed 2026-10-27 20:43:39 +08:00
1 parent 8e6ed0317b
commit 86c0276156
1022 files changed
+69903 -19207

No files matched your search

+273
View File
@@ -0,0 +1,273 @@
package service
import (
"errors"
"slices"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/service/authz"
)
const (
AccessTokenGroupPersonal = "personal"
AccessTokenGroupAdmin = "admin"
AccessTokenGroupSystem = "system"
AccessTokenMaxPerUser = 20
AccessTokenDefaultExpiryDays = 30
accessTokenActionReveal = "reveal"
accessTokenScopeSeparator = ":"
)
var (
ErrAccessTokenScopeInvalid = errors.New("access token scope is invalid")
ErrAccessTokenScopeForbidden = errors.New("access token scope cannot be granted")
)
// AccessTokenResource mirrors authz.ResourceDefinition so both catalogs share
// one JSON shape and one frontend matrix component.
type AccessTokenResource struct {
Resource string `json:"resource"`
LabelKey string `json:"label_key"`
Actions []authz.ActionDefinition `json:"actions"`
group string
minRole int
}
type AccessTokenCatalogGroup struct {
Group string `json:"group"`
Resources []AccessTokenResource `json:"resources"`
}
func accessTokenView(description string) authz.ActionDefinition {
return authz.ActionDefinition{Action: authz.ActionRead, LabelKey: "View", DescriptionKey: description}
}
func accessTokenEdit(description string) authz.ActionDefinition {
return authz.ActionDefinition{Action: authz.ActionWrite, LabelKey: "Edit", DescriptionKey: description}
}
// accessTokenStaticResources covers dashboard routes that are not guarded by a
// Casbin permission. Casbin resources come from authz.Catalog() at runtime and
// must never be duplicated here.
var accessTokenStaticResources = []AccessTokenResource{
{Resource: "profile", LabelKey: "Profile", group: AccessTokenGroupPersonal, minRole: common.RoleCommonUser, Actions: []authz.ActionDefinition{
accessTokenView("View your profile, groups, and available models."),
accessTokenEdit("Change your display name and personal settings."),
}},
{Resource: "api_key", LabelKey: "API keys", group: AccessTokenGroupPersonal, minRole: common.RoleCommonUser, Actions: []authz.ActionDefinition{
accessTokenView("View API keys and their settings without the full key."),
accessTokenEdit("Create, edit, and delete API keys."),
{Action: accessTokenActionReveal, LabelKey: "Reveal full keys", DescriptionKey: "View complete API keys."},
}},
{Resource: "usage", LabelKey: "Usage", group: AccessTokenGroupPersonal, minRole: common.RoleCommonUser, Actions: []authz.ActionDefinition{
accessTokenView("View your usage logs, statistics, tasks, and audit logs."),
}},
{Resource: "wallet", LabelKey: "Wallet", group: AccessTokenGroupPersonal, minRole: common.RoleCommonUser, Actions: []authz.ActionDefinition{
accessTokenView("View balance, top-up records, subscriptions, and referral information."),
accessTokenEdit("Top up, pay for subscriptions, check in, and transfer referral rewards."),
}},
{Resource: "account_security", LabelKey: "Account security", group: AccessTokenGroupPersonal, minRole: common.RoleCommonUser, Actions: []authz.ActionDefinition{
accessTokenView("View Passkey, two-factor, and sign-in binding status."),
accessTokenEdit("Includes changes that require security verification, such as password, two-factor settings, and account deletion."),
}},
{Resource: "user", LabelKey: "Users", group: AccessTokenGroupAdmin, minRole: common.RoleAdminUser, Actions: []authz.ActionDefinition{
accessTokenView("View users, their sign-in bindings, and permission settings."),
accessTokenEdit("Create, edit, disable, and delete users, and reset their security settings."),
}},
{Resource: "billing", LabelKey: "Billing management", group: AccessTokenGroupAdmin, minRole: common.RoleAdminUser, Actions: []authz.ActionDefinition{
accessTokenView("View all top-up records, subscription plans, and user subscriptions."),
accessTokenEdit("Complete top-up orders and manage subscription plans and user subscriptions."),
}},
{Resource: "model", LabelKey: "Models", group: AccessTokenGroupAdmin, minRole: common.RoleAdminUser, Actions: []authz.ActionDefinition{
accessTokenView("View models and vendors, and preview upstream sync."),
accessTokenEdit("Edit models and vendors, and sync them from upstream."),
}},
{Resource: "deployment", LabelKey: "Deployments", group: AccessTokenGroupAdmin, minRole: common.RoleAdminUser, Actions: []authz.ActionDefinition{
accessTokenView("View deployments and price estimates."),
accessTokenEdit("Create, edit, and delete deployments, and test connections."),
}},
{Resource: "log", LabelKey: "Logs", group: AccessTokenGroupAdmin, minRole: common.RoleAdminUser, Actions: []authz.ActionDefinition{
accessTokenView("View all users' usage logs, statistics, and tasks."),
}},
{Resource: "redemption", LabelKey: "Redemption codes", group: AccessTokenGroupAdmin, minRole: common.RoleAdminUser, Actions: []authz.ActionDefinition{
accessTokenView("View redemption codes."),
accessTokenEdit("Create, edit, and delete redemption codes."),
}},
{Resource: "group", LabelKey: "Groups", group: AccessTokenGroupAdmin, minRole: common.RoleAdminUser, Actions: []authz.ActionDefinition{
accessTokenView("View groups and prefill groups."),
accessTokenEdit("Create, edit, and delete prefill groups."),
}},
{Resource: "option", LabelKey: "System settings", group: AccessTokenGroupSystem, minRole: common.RoleRootUser, Actions: []authz.ActionDefinition{
accessTokenView("View system settings, custom OAuth providers, and price sync sources."),
accessTokenEdit("Change system settings and custom OAuth providers, and sync prices."),
}},
{Resource: "plugin", LabelKey: "Task plugins", group: AccessTokenGroupSystem, minRole: common.RoleRootUser, Actions: []authz.ActionDefinition{
accessTokenView("View installed task plugins and their settings."),
accessTokenEdit("Install, configure, test, and remove task plugins."),
}},
{Resource: "ops", LabelKey: "Operations", group: AccessTokenGroupSystem, minRole: common.RoleRootUser, Actions: []authz.ActionDefinition{
accessTokenView("View performance, system tasks, and system information."),
accessTokenEdit("Run and manage maintenance tasks and performance settings."),
}},
}
// AccessTokenScopeOf is the token scope key for a Casbin permission.
func AccessTokenScopeOf(permission authz.Permission) string {
return permission.Resource + accessTokenScopeSeparator + permission.Action
}
// AccessTokenCatalog lists the scopes this user may grant right now, grouped
// for display. Resources without a grantable action and empty groups are
// omitted.
func AccessTokenCatalog(userID, role int) []AccessTokenCatalogGroup {
resources := map[string][]AccessTokenResource{}
for _, resource := range accessTokenStaticResources {
if role < resource.minRole {
continue
}
resource.Actions = slices.Clone(resource.Actions)
resources[resource.group] = append(resources[resource.group], resource)
}
if role >= common.RoleAdminUser {
for _, definition := range authz.Catalog() {
resource := AccessTokenResource{Resource: definition.Resource, LabelKey: definition.LabelKey}
for _, action := range definition.Actions {
if authz.Can(userID, role, authz.Permission{Resource: definition.Resource, Action: action.Action}) {
resource.Actions = append(resource.Actions, action)
}
}
if len(resource.Actions) > 0 {
resources[AccessTokenGroupAdmin] = append(resources[AccessTokenGroupAdmin], resource)
}
}
}
groups := make([]AccessTokenCatalogGroup, 0, 3)
for _, group := range []string{AccessTokenGroupPersonal, AccessTokenGroupAdmin, AccessTokenGroupSystem} {
if len(resources[group]) > 0 {
groups = append(groups, AccessTokenCatalogGroup{Group: group, Resources: resources[group]})
}
}
return groups
}
// AccessTokenScopeDictionary lists every resource a token scope can name,
// whatever the viewer may grant, so stored grants keep their labels after the
// granting user loses a permission or when another user reads them.
func AccessTokenScopeDictionary() []AccessTokenResource {
definitions := authz.Catalog()
resources := make([]AccessTokenResource, 0, len(accessTokenStaticResources)+len(definitions))
for _, resource := range accessTokenStaticResources {
resource.Actions = slices.Clone(resource.Actions)
resources = append(resources, resource)
}
for _, definition := range definitions {
resources = append(resources, AccessTokenResource{Resource: definition.Resource, LabelKey: definition.LabelKey, Actions: definition.Actions})
}
return resources
}
// NormalizeAccessTokenScopes validates a requested grant against the catalog
// the user sees and returns it trimmed, de-duplicated and sorted.
func NormalizeAccessTokenScopes(userID, role int, scopes []string) ([]string, error) {
normalized, ok := NormalizeAccessTokenScopeList(scopes)
if !ok {
return nil, ErrAccessTokenScopeInvalid
}
grantable := make(map[string]bool)
for _, group := range AccessTokenCatalog(userID, role) {
for _, resource := range group.Resources {
for _, action := range resource.Actions {
grantable[resource.Resource+accessTokenScopeSeparator+action.Action] = true
}
}
}
for _, scope := range normalized {
if grantable[scope] {
continue
}
if !knownAccessTokenScope(scope) {
return nil, ErrAccessTokenScopeInvalid
}
return nil, ErrAccessTokenScopeForbidden
}
return normalized, nil
}
func knownAccessTokenScope(scope string) bool {
resourceName, action, ok := strings.Cut(scope, accessTokenScopeSeparator)
if !ok {
return false
}
for _, resource := range accessTokenStaticResources {
if resource.Resource == resourceName && slices.ContainsFunc(resource.Actions, func(definition authz.ActionDefinition) bool { return definition.Action == action }) {
return true
}
}
return slices.Contains(authz.AllPermissions(), authz.Permission{Resource: resourceName, Action: action})
}
func AccessTokenScopeGranted(scopes []string, scope string) bool {
return scope != "" && slices.Contains(scopes, scope)
}
// accessTokenVerificationScopes maps every verification scope to the token
// scope a PAT needs to obtain its proof. An empty value means PATs can never
// obtain it; scopes missing from this table are denied too.
var accessTokenVerificationScopes = map[string]string{
VerificationScopeChannelKeyRead: AccessTokenScopeOf(authz.ChannelSecretView),
VerificationScopeAdminUserCreate: "user:write",
VerificationScopeAdminUserUpdate: "user:write",
VerificationScopeAdminUserDelete: "user:write",
VerificationScopeAdminUserManage: "user:write",
VerificationScopeAdminUserPasskeyReset: "user:write",
VerificationScopeAdminUserTwoFADisable: "user:write",
VerificationScopeAdminUserBindingClear: "user:write",
VerificationScopePasskeyRegister: "account_security:write",
VerificationScopePasskeyDelete: "account_security:write",
VerificationScopeTwoFASetup: "account_security:write",
VerificationScopeTwoFADisable: "account_security:write",
VerificationScopeTwoFABackupCodes: "account_security:write",
VerificationScopePasswordSet: "account_security:write",
VerificationScopePasswordChange: "account_security:write",
VerificationScopeAccountBind: "account_security:write",
VerificationScopeAccountUnbind: "account_security:write",
VerificationScopeAccountDelete: "account_security:write",
VerificationScopeAccessTokenGenerate: "",
VerificationScopeAccessTokenUpdate: "",
VerificationScopeAccessTokenRevoke: "",
VerificationScopeLogin: "",
}
// AccessTokenVerificationScope returns the token scope a PAT needs for a
// verification scope and whether the scope is listed at all.
func AccessTokenVerificationScope(verificationScope string) (scope string, listed bool) {
scope, listed = accessTokenVerificationScopes[verificationScope]
return scope, listed
}
// requireAccessTokenVerificationScope rejects a step-up request from a PAT
// whose grant does not cover the operation. Browser sessions pass through.
func requireAccessTokenVerificationScope(identity AuthIdentity, verificationScope string) error {
tokenID, ok := model.ParseAccessTokenSessionID(identity.SessionID)
if !ok {
return nil
}
required, _ := AccessTokenVerificationScope(verificationScope)
if required == "" {
return ErrVerificationForbidden
}
token, err := model.GetUserAccessToken(identity.UserID, tokenID)
if errors.Is(err, model.ErrAccessTokenNotFound) {
return ErrAuthTokenInvalid
}
if err != nil {
return err
}
if !AccessTokenScopeGranted(token.GetScopes(), required) {
return ErrVerificationForbidden
}
return nil
}
+19 -3
View File
@@ -5,9 +5,11 @@ import (
"html"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/oauth"
"github.com/QuantumNous/new-api/setting/system_setting"
"github.com/gin-gonic/gin"
)
func UnbindAccountOAuth(identity AuthIdentity, providerID int) error {
@@ -29,13 +31,27 @@ func UnbindAccountOAuth(identity AuthIdentity, providerID int) error {
return model.UnbindUserOAuthForSession(identity, providerID, enabled)
}
// AccountMailLanguage is the language of mail about a user's account: the
// language the user saved, else the one the request stated.
func AccountMailLanguage(c *gin.Context, userId int) string {
if lang := model.GetUserLanguage(userId); lang != "" {
return lang
}
return i18n.StatedLang(c)
}
// NotifyAccountSecurityChange never includes credentials or tokens. The caller
// records delivery failure independently from the already-committed change.
func NotifyAccountSecurityChange(email, event string) error {
// eventKey names the change; the notice is written in lang (see
// AccountMailLanguage).
func NotifyAccountSecurityChange(lang, email, eventKey string, eventArgs ...map[string]any) error {
if email == "" {
return nil
}
subject := common.SystemName + " — Account security notification"
content := fmt.Sprintf("<p>Your account security settings have changed: %s.</p><p>If you did not make this change, open your account security settings, revoke other login sessions, and contact your administrator.</p>", html.EscapeString(event))
event := html.EscapeString(i18n.Translate(lang, eventKey, eventArgs...))
subject := i18n.Translate(lang, i18n.MsgEmailSecurityNoticeSubject, map[string]any{"SystemName": common.SystemName})
content := fmt.Sprintf("<p>%s</p><p>%s</p>",
i18n.Translate(lang, i18n.MsgEmailSecurityNoticeChanged, map[string]any{"Event": event}),
i18n.Translate(lang, i18n.MsgEmailSecurityNoticeAdvice))
return common.SendEmail(subject, email, content)
}
-3
View File
@@ -286,9 +286,6 @@ func aggregateAstraSampling(channels []model.Channel, tasks []*model.SystemTask,
break
}
}
if len(s.History) > 12 {
s.History = s.History[:12]
}
}
return groups
}
+2 -2
View File
@@ -90,14 +90,14 @@ func TestAstraSamplingAggregationPrivacyAndFailures(t *testing.T) {
func TestAstraSamplingExcludesInconclusiveFromRateWithinSixHours(t *testing.T) {
now := time.Now()
channel := model.Channel{Id: 1, Group: "example-variable", Models: fingerprint.Model}
channel := model.Channel{Id: 1, Group: "gpt-s-unstable", Models: fingerprint.Model}
rounds := []AstraSamplingRound{
{ID: "recent-match", ChannelID: 1, StartedAt: now.Add(-time.Hour).Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "match", ValidSamples: 3, Closest: fingerprint.Model}},
{ID: "recent-luna", ChannelID: 1, StartedAt: now.Add(-2 * time.Hour).Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "inconclusive", ValidSamples: 2, Closest: "gpt-5.6-luna", Gap: .9}},
{ID: "old-luna", ChannelID: 1, StartedAt: now.Add(-7 * time.Hour).Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "inconclusive", ValidSamples: 2, Closest: "gpt-5.6-luna", Gap: .9}},
}
state, _ := common.Marshal(AstraSamplingState{Rounds: rounds})
summary := AggregateAstraSampling([]model.Channel{channel}, []*model.SystemTask{{Status: model.SystemTaskStatusSucceeded, State: string(state)}}, now)["example-variable"]
summary := AggregateAstraSampling([]model.Channel{channel}, []*model.SystemTask{{Status: model.SystemTaskStatusSucceeded, State: string(state)}}, now)["gpt-s-unstable"]
if summary.Rate == nil || *summary.Rate != 0 || summary.Matches != 1 || summary.Mismatches != 0 || summary.Inconclusive != 1 {
t.Fatalf("%+v", summary)
}
+12 -8
View File
@@ -1,7 +1,6 @@
package service
import (
"fmt"
"time"
"github.com/QuantumNous/new-api/common"
@@ -10,8 +9,9 @@ import (
const authArtifactCleanupInterval = time.Hour
// StartAuthArtifactCleanup removes expired dashboard Sessions and old
// one-time authentication flows. Only the master instance performs cleanup.
// StartAuthArtifactCleanup removes expired dashboard Sessions, old one-time
// authentication flows and access tokens that expired beyond the revoked
// session retention. Only the master instance performs cleanup.
func StartAuthArtifactCleanup() {
if !common.IsMasterNode {
return
@@ -30,9 +30,9 @@ func cleanupAuthArtifacts() {
now := time.Now()
count, err := model.CountUserSessionsCreatedSince(0, now.Add(-time.Hour).Unix())
if err != nil {
common.SysError("failed to count hourly user session issuance: " + err.Error())
common.SysError(common.LogText("failed to count hourly user session issuance: %s", err.Error()))
} else if count > int64(common.UserSessionHourlyAlertThreshold) {
common.SysError(fmt.Sprintf(
common.SysError(common.LogText(
"hourly user session issuance exceeded alert threshold: count=%d threshold=%d window_seconds=%d",
count,
common.UserSessionHourlyAlertThreshold,
@@ -40,12 +40,16 @@ func cleanupAuthArtifacts() {
))
}
if err := model.DeleteExpiredUserSessions(now.Unix()); err != nil {
common.SysError("failed to delete expired user sessions: " + err.Error())
common.SysError(common.LogText("failed to delete expired user sessions: %s", err.Error()))
}
if err := model.DeleteOldRevokedUserSessions(now.Unix()); err != nil {
common.SysError("failed to delete old revoked user sessions: " + err.Error())
common.SysError(common.LogText("failed to delete old revoked user sessions: %s", err.Error()))
}
if err := model.DeleteExpiredAuthFlows(now); err != nil {
common.SysError("failed to delete expired authentication flows: " + err.Error())
common.SysError(common.LogText("failed to delete expired authentication flows: %s", err.Error()))
}
accessTokenExpiredBefore := now.Unix() - int64(common.UserSessionRevokedRetentionDays)*24*60*60
if err := model.DeleteExpiredUserAccessTokens(accessTokenExpiredBefore); err != nil {
common.SysError(common.LogText("failed to delete expired access tokens: %s", err.Error()))
}
}
+18 -1
View File
@@ -149,6 +149,16 @@ func ValidateLoginSession(identity AuthIdentity) (*model.UserSession, *model.Use
return session, user, nil
}
// ValidateStepUpIdentity checks that the identity requesting or consuming a
// proof is still live: a browser session or a scoped access token.
func ValidateStepUpIdentity(identity AuthIdentity) error {
if _, ok := model.ParseAccessTokenSessionID(identity.SessionID); ok {
return model.ValidateAccessTokenIdentity(identity)
}
_, _, err := ValidateLoginSession(identity)
return err
}
// ValidateSessionReference validates a server-side flow bound to an existing
// dashboard session without requiring an access token on the callback request.
func ValidateSessionReference(userID int, sid string) (AuthIdentity, error) {
@@ -174,7 +184,8 @@ func ValidateSessionReference(userID int, sid string) (AuthIdentity, error) {
// AdvanceCurrentSessionSecurity increments the user's global auth version,
// preserves only the current browser session at a new session version and
// returns a replacement access token. Call after a successful 2FA/passkey
// security-setting mutation that did not already advance AuthVersion.
// security-setting mutation that did not already advance AuthVersion. When the
// change was made through an access token it returns a nil bundle.
func AdvanceCurrentSessionSecurity(identity AuthIdentity, reason string) (*AuthBundle, error) {
nextUserAuthVersion, err := model.BumpUserAuthVersion(identity.UserID)
if err != nil {
@@ -198,6 +209,12 @@ func AdvanceCurrentSessionToUserVersion(identity AuthIdentity, reason string) (*
}
func advanceCurrentSessionToVersion(identity AuthIdentity, nextUserAuthVersion int64, reason string) (*AuthBundle, error) {
// An access token has no browser session to keep, so every browser session
// signs in again at the new auth version. The token itself stays valid.
if _, ok := model.ParseAccessTokenSessionID(identity.SessionID); ok {
_, err := model.RevokeAllUserSessions(identity.UserID, reason)
return nil, err
}
session, err := model.AdvanceUserSessionAuthVersion(
identity.UserID,
identity.SessionID,
+14 -1
View File
@@ -32,7 +32,7 @@ func setupAuthSessionTestDB(t *testing.T) *model.User {
sqlDB, err := db.DB()
require.NoError(t, err)
sqlDB.SetMaxOpenConns(1)
require.NoError(t, db.AutoMigrate(&model.User{}, &model.UserSession{}, &model.AuthFlow{}))
require.NoError(t, db.AutoMigrate(&model.User{}, &model.UserSession{}, &model.AuthFlow{}, &model.UserAccessToken{}))
model.DB = db
common.RedisEnabled = false
common.UserSessionActiveLimit = common.DefaultUserSessionActiveLimit
@@ -276,6 +276,16 @@ func TestCleanupAuthArtifactsRemovesOnlyExpiredRecords(t *testing.T) {
ExpiresAt: now.Add(time.Minute),
}).Error)
retention := time.Duration(common.UserSessionRevokedRetentionDays) * 24 * time.Hour
for _, token := range []model.UserAccessToken{
{Name: "past retention", TokenHash: "past-retention", ExpiresAt: now.Add(-retention - time.Hour).Unix()},
{Name: "within retention", TokenHash: "within-retention", ExpiresAt: now.Add(-retention + time.Hour).Unix()},
{Name: "never expires", TokenHash: "never-expires"},
} {
token.UserId = 1
require.NoError(t, model.DB.Create(&token).Error)
}
cleanupAuthArtifacts()
var sessionCount int64
@@ -285,6 +295,9 @@ func TestCleanupAuthArtifactsRemovesOnlyExpiredRecords(t *testing.T) {
require.NoError(t, model.DB.Find(&flows).Error)
require.Len(t, flows, 1)
assert.Equal(t, "recent-flow", flows[0].TokenHash)
var kept []string
require.NoError(t, model.DB.Model(&model.UserAccessToken{}).Order("id").Pluck("token_hash", &kept).Error)
assert.Equal(t, []string{"within-retention", "never-expires"}, kept)
}
func TestCleanupAuthArtifactsContinuesWithRevokedCleanupAfterExpiredBatchFailure(t *testing.T) {
+22
View File
@@ -6,6 +6,7 @@ import (
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/golang-jwt/jwt/v5"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
@@ -160,3 +161,24 @@ func TestSecurityProofBindsIdentityAndOperation(t *testing.T) {
_, err = ConsumeOperationProof(expired, identity, operation)
assert.ErrorIs(t, err, ErrAuthTokenExpired)
}
func TestSecurityProofBindsAccessTokenSession(t *testing.T) {
setupAuthSessionTestDB(t)
useTestSessionSecret(t)
identity := AuthIdentity{UserID: 42, SessionID: model.AccessTokenSessionID(7), UserAuthVersion: 3, SessionVersion: model.AccessTokenSessionVersion}
binding, err := BindVerificationOperation(VerificationOperation{Scope: VerificationScopeTwoFADisable})
require.NoError(t, err)
proof, _, err := IssueSecurityProof(identity, "password", binding)
require.NoError(t, err)
claims, err := verifySecurityProof(proof, identity, binding)
require.NoError(t, err)
assert.Equal(t, identity.SessionID, claims.SessionID)
for _, sessionID := range []string{model.AccessTokenSessionID(8), "session-1"} {
other := identity
other.SessionID = sessionID
_, err = verifySecurityProof(proof, other, binding)
assert.ErrorIs(t, err, ErrAuthTokenInvalid, sessionID)
}
}
+1 -1
View File
@@ -44,7 +44,7 @@ func (a *gormAdapter) LoadPolicy(m casbinmodel.Model) error {
// the stored rule for review and deny this permission in memory.
// Skipping a per-user rule would expose the admin role baseline.
effect = EffectDeny
common.SysLog(fmt.Sprintf("authorization policy %d has an unsupported legacy scope; loaded as deny; review the stored policy before granting access", rule.Id))
common.SysLog(common.LogText("authorization policy %d has an unsupported legacy scope; loaded as deny; review the stored policy before granting access", rule.Id))
}
if err := persist.LoadPolicyArray([]string{rule.Ptype, rule.V0, rule.V1, rule.V2, effect}, m); err != nil {
return err
+1 -1
View File
@@ -88,7 +88,7 @@ func StartPolicySync(frequency int) {
for {
time.Sleep(time.Duration(frequency) * time.Second)
if err := ReloadPolicy(); err != nil {
common.SysError("failed to reload authz policy: " + err.Error())
common.SysError(common.LogText("failed to reload authz policy: %s", err.Error()))
}
}
}
+17 -4
View File
@@ -34,8 +34,8 @@ func automaticModelMetrics(modelName string) perfmetrics.QueryResult {
}
type automaticCandidate struct {
group string
cost, ttft, success, ratio float64
group string
cost, ttft, throughput, success, ratio float64
}
func resolveAutomaticPool(c *gin.Context, modelName string) []model.TokenGroupEntry {
@@ -59,7 +59,7 @@ func resolveAutomaticPool(c *gin.Context, modelName string) []model.TokenGroupEn
seen := map[string]bool{}
for _, ability := range abilities {
group := ability.Group
if _, ok := allowed[group]; !ok || seen[group] || !IsUserSelectableGroup(userGroup, group) {
if _, ok := allowed[group]; !ok || seen[group] || !IsUserSelectableGroup(userGroup, group) || RequestRoutingPolicy(c).ExcludesGroup(group) {
continue
}
channel, err := model.CacheGetChannel(ability.ChannelId)
@@ -92,6 +92,14 @@ func resolveAutomaticPool(c *gin.Context, modelName string) []model.TokenGroupEn
}
ratio := model.EffectiveGroupRatio(c, userGroup, group).GroupRatio
item := automaticCandidate{group: group, ttft: math.Inf(1), success: -1, ratio: ratio}
if RequestRoutingPolicy(c).AutoPreference == "throughput" {
metric, err := GroupModelPerformance(c.Request.Context(), group, modelName)
if err != nil {
common.SysError("automatic group performance query failed: " + err.Error())
} else if metric.Throughput != nil {
item.throughput = metric.Throughput.Value
}
}
if metric, ok := metrics[group]; ok && len(metric.Series) > 0 {
item.success = metric.SuccessRate
@@ -125,6 +133,11 @@ func resolveAutomaticPool(c *gin.Context, modelName string) []model.TokenGroupEn
func sortAutomaticCandidates(items []automaticCandidate, preference string) {
slices.SortFunc(items, func(a, b automaticCandidate) int {
if preference == "throughput" {
if n := cmp.Compare(b.throughput, a.throughput); n != 0 {
return n
}
}
if preference == "speed" {
if n := cmp.Compare(a.ttft, b.ttft); n != 0 {
return n
@@ -163,7 +176,7 @@ func AutomaticPoolModels(c *gin.Context) []string {
}
names := map[string]bool{}
for _, ability := range abilities {
if !IsUserSelectableGroup(userGroup, ability.Group) {
if !IsUserSelectableGroup(userGroup, ability.Group) || RequestRoutingPolicy(c).ExcludesGroup(ability.Group) {
continue
}
channel, err := model.CacheGetChannel(ability.ChannelId)
+85 -2
View File
@@ -5,10 +5,13 @@ import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
"github.com/QuantumNous/new-api/setting"
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/driver/mysql"
"gorm.io/driver/postgres"
@@ -45,9 +48,9 @@ func TestAutomaticPoolModelsDatabaseMatrix(t *testing.T) {
model.DB = db
common.RedisEnabled = false
common.MemoryCacheEnabled = false
require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}))
require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.Token{}))
t.Cleanup(func() {
require.NoError(t, db.Migrator().DropTable(&model.Ability{}, &model.Channel{}))
require.NoError(t, db.Migrator().DropTable(&model.Token{}, &model.Ability{}, &model.Channel{}))
model.DB = oldDB
common.RedisEnabled = oldRedis
common.MemoryCacheEnabled = oldMemory
@@ -92,6 +95,86 @@ func TestAutomaticPoolModelsDatabaseMatrix(t *testing.T) {
c.Request = httptest.NewRequest("GET", "/v1/models", nil)
common.SetContextKey(c, constant.ContextKeyUserGroup, "auto-a")
require.Equal(t, []string{"gpt-6-astra", "model-a", "new-model", "shared"}, AutomaticPoolModels(c))
cfg := model_setting.GetGlobalSettings()
oldImageGroup, oldImageModels := cfg.ImageGenerationGroup, cfg.ImageGenerationModels
cfg.ImageGenerationGroup, cfg.ImageGenerationModels = "auto-b", []string{"blacklist-image"}
t.Cleanup(func() {
cfg.ImageGenerationGroup, cfg.ImageGenerationModels = oldImageGroup, oldImageModels
automaticMetricCache.Purge()
_ = poolCooldowns().Purge()
})
createChannelSelectAutoGroupsChannel(t, db, 571, "auto-a", "blacklist-model")
createChannelSelectAutoGroupsChannel(t, db, 572, "auto-b", "blacklist-model")
createChannelSelectAutoGroupsChannel(t, db, 573, "auto-a", "blacklist-image")
createChannelSelectAutoGroupsChannel(t, db, 574, "auto-b", "blacklist-image")
common.MemoryCacheEnabled = true
model.InitChannelCache()
automaticMetricCache.Set("blacklist-model", perfmetrics.QueryResult{})
automaticMetricCache.Set("blacklist-image", perfmetrics.QueryResult{})
makeContext := func(policy *model.TokenRoutingPolicy, tokenID int) *gin.Context {
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
ctx.Request = httptest.NewRequest("POST", "/v1/responses", nil)
ctx.Set("token_id", tokenID)
ctx.Set("id", 7)
common.SetContextKey(ctx, constant.ContextKeyUserGroup, "auto-a")
ctx.Set(model.TokenRoutingContextKey, policy)
return ctx
}
for _, preference := range []string{"cost", "speed", "stable"} {
for _, blocked := range [][]string{nil, {"auto-a"}, {"auto-a", "auto-b"}} {
policy := &model.TokenRoutingPolicy{Version: 1, Mode: "smart", AutoPreference: preference, ExcludedGroups: blocked}
body, err := common.Marshal(policy)
require.NoError(t, err)
token := model.Token{Key: fmt.Sprintf("blacklist-%s-%d", preference, len(blocked)), UserId: 7, RoutingPolicy: string(body)}
require.NoError(t, token.Insert())
var stored model.Token
require.NoError(t, db.First(&stored, token.Id).Error)
loaded, err := stored.GetRoutingPolicy()
require.NoError(t, err)
assert.Equal(t, blocked, loaded.ExcludedGroups)
ctx := makeContext(loaded, stored.Id)
pool := ResolveRequestPool(ctx, "blacklist-model")
require.Len(t, pool, 2-len(blocked))
assert.Equal(t, max(0, 2*len(pool)-1), PoolRetryBudget(ctx, "blacklist-model"))
for _, entry := range pool {
assert.False(t, loaded.ExcludesGroup(entry.Group))
}
channel, group, selectErr := SelectChannelForRequest(ctx, "blacklist-model", &RetryParam{Ctx: ctx, TokenGroup: "auto", ModelName: "blacklist-model"})
if len(pool) == 0 {
assert.Nil(t, channel)
require.NotNil(t, selectErr)
assert.True(t, selectErr.NoAvailableChannel)
} else {
require.Nil(t, selectErr)
require.NotNil(t, channel)
assert.False(t, loaded.ExcludesGroup(group))
RecordPoolRetryableFailure(ctx, "blacklist-model", group)
retryChannel, retryGroup, err := selectTokenRoutingPool(&RetryParam{Ctx: ctx, TokenGroup: "auto", ModelName: "blacklist-model", Retry: common.GetPointer(1)})
require.NoError(t, err)
require.NotNil(t, retryChannel)
assert.False(t, loaded.ExcludesGroup(retryGroup))
}
loaded.ExcludedGroups = nil
body, err = common.Marshal(loaded)
require.NoError(t, err)
stored.RoutingPolicy = string(body)
require.NoError(t, stored.Update())
require.NoError(t, db.First(&stored, token.Id).Error)
cleared, err := stored.GetRoutingPolicy()
require.NoError(t, err)
assert.Empty(t, cleared.ExcludedGroups)
assert.Len(t, ResolveRequestPool(makeContext(cleared, stored.Id), "blacklist-model"), 2)
}
}
policy := &model.TokenRoutingPolicy{Version: 1, Mode: "smart", ExcludedGroups: []string{"auto-b"}}
ctx := makeContext(policy, 900)
assert.Empty(t, ResolveRequestPool(ctx, "blacklist-image"), "the dedicated image group cannot bypass the blacklist or fall back")
assert.NotContains(t, AutomaticPoolModels(ctx), "gpt-6-astra")
assert.Contains(t, AutomaticPoolModels(ctx), "shared")
assert.Contains(t, AutomaticPoolModels(makeContext(&model.TokenRoutingPolicy{Version: 1, Mode: "smart"}, 901)), "gpt-6-astra", "another key is unaffected")
manual := &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "auto-b", Priority: 0}}, ExcludedGroups: []string{"auto-b"}}
assert.Equal(t, manual.Groups, ResolveRequestPool(makeContext(manual, 902), "blacklist-model"))
})
}
}
+4 -3
View File
@@ -4,6 +4,7 @@ import (
"fmt"
"net/http"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/logger"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/types"
@@ -54,19 +55,19 @@ func SettleBilling(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, actualQuo
delta := actualQuota - preConsumed
if delta > 0 {
logger.LogInfo(ctx, fmt.Sprintf("预扣费后补扣费:%s(实际消耗:%s,预扣费:%s)",
logger.LogInfo(ctx, common.LogText("charging the difference after pre-consume: %s (actual %s, pre-consumed %s)",
logger.FormatQuota(delta),
logger.FormatQuota(actualQuota),
logger.FormatQuota(preConsumed),
))
} else if delta < 0 {
logger.LogInfo(ctx, fmt.Sprintf("预扣费后返还扣费:%s(实际消耗:%s,预扣费:%s)",
logger.LogInfo(ctx, common.LogText("refunding the difference after pre-consume: %s (actual %s, pre-consumed %s)",
logger.FormatQuota(-delta),
logger.FormatQuota(actualQuota),
logger.FormatQuota(preConsumed),
))
} else {
logger.LogInfo(ctx, fmt.Sprintf("预扣费与实际消耗一致,无需调整:%s(按次计费)",
logger.LogInfo(ctx, common.LogText("pre-consumed quota matches the actual quota, no adjustment: %s (per-request billing)",
logger.FormatQuota(actualQuota),
))
}
+32 -26
View File
@@ -3,16 +3,19 @@ package service
import (
"errors"
"fmt"
"math"
"net/http"
"strings"
"sync"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/setting/operation_setting"
"github.com/bytedance/gopkg/util/gopool"
"github.com/gin-gonic/gin"
@@ -27,13 +30,14 @@ import (
type BillingSession struct {
relayInfo *relaycommon.RelayInfo
funding FundingSource
preConsumedQuota int // 实际预扣额度(信任用户可能为 0)
tokenConsumed int // 令牌额度实际扣减量
extraReserved int // 发送前补充预扣的额度(订阅退款时需要单独回滚)
trusted bool // 是否命中信任额度旁路
fundingSettled bool // funding.Settle 已成功,资金来源已提交
settled bool // Settle 全部完成(资金 + 令牌)
refunded bool // Refund 已调用
preConsumedQuota int // 实际预扣额度(信任用户可能为 0)
tokenConsumed int // 令牌额度实际扣减量
extraReserved int // 发送前补充预扣的额度(订阅退款时需要单独回滚)
trusted bool // 是否命中信任额度旁路
fundingSettled bool // funding.Settle 已成功,资金来源已提交
settled bool // Settle 全部完成(资金 + 令牌)
refunded bool // Refund 已调用
lang string // 请求声明的语言,为空表示未声明
mu sync.Mutex
}
@@ -68,7 +72,7 @@ func (s *BillingSession) Settle(actualQuota int) error {
}
if tokenErr != nil {
// 资金来源已提交,令牌调整失败只能记录日志;标记 settled 防止 Refund 误退资金
common.SysLog(fmt.Sprintf("error adjusting token quota after funding settled (userId=%d, tokenId=%d, delta=%d): %s",
common.SysLog(common.LogText("error adjusting token quota after funding settled (userId=%d, tokenId=%d, delta=%d): %s",
s.relayInfo.UserId, s.relayInfo.TokenId, delta, tokenErr.Error()))
}
}
@@ -90,7 +94,7 @@ func (s *BillingSession) Refund(c *gin.Context) {
s.refunded = true
s.mu.Unlock()
logger.LogInfo(c, fmt.Sprintf("用户 %d 请求失败, 返还预扣费(token_quota=%s, funding=%s)",
logger.LogInfo(c, common.LogText("user %d request failed, refunding the pre-consumed quota (token_quota=%s, funding=%s)",
s.relayInfo.UserId,
logger.FormatQuota(s.tokenConsumed),
s.funding.Source(),
@@ -108,17 +112,17 @@ func (s *BillingSession) Refund(c *gin.Context) {
gopool.Go(func() {
// 1) 退还资金来源
if err := funding.Refund(); err != nil {
common.SysLog("error refunding billing source: " + err.Error())
common.SysLog(common.LogText("error refunding billing source: %s", err.Error()))
}
if extraReserved > 0 && funding.Source() == BillingSourceSubscription && subscriptionId > 0 {
if err := model.PostConsumeUserSubscriptionDelta(subscriptionId, -int64(extraReserved)); err != nil {
common.SysLog("error refunding subscription extra reserved quota: " + err.Error())
common.SysLog(common.LogText("error refunding subscription extra reserved quota: %s", err.Error()))
}
}
// 2) 退还令牌额度
if tokenConsumed > 0 && !isPlayground {
if err := model.IncreaseTokenQuota(tokenId, tokenKey, tokenConsumed); err != nil {
common.SysLog("error refunding token quota: " + err.Error())
common.SysLog(common.LogText("error refunding token quota: %s", err.Error()))
}
}
})
@@ -200,9 +204,9 @@ func (s *BillingSession) preConsume(c *gin.Context, quota int) *types.NewAPIErro
if s.shouldTrust(c) {
s.trusted = true
effectiveQuota = 0
logger.LogInfo(c, fmt.Sprintf("用户 %d 额度充足, 信任且不需要预扣费 (funding=%s)", s.relayInfo.UserId, s.funding.Source()))
logger.LogInfo(c, common.LogText("user %d has enough quota, trusted without pre-consume (funding=%s)", s.relayInfo.UserId, s.funding.Source()))
} else if effectiveQuota > 0 {
logger.LogInfo(c, fmt.Sprintf("用户 %d 需要预扣费 %s (funding=%s)", s.relayInfo.UserId, logger.FormatQuota(effectiveQuota), s.funding.Source()))
logger.LogInfo(c, common.LogText("user %d needs pre-consume %s (funding=%s)", s.relayInfo.UserId, logger.FormatQuota(effectiveQuota), s.funding.Source()))
}
// ---- 1) 预扣令牌额度 ----
@@ -218,7 +222,7 @@ func (s *BillingSession) preConsume(c *gin.Context, quota int) *types.NewAPIErro
// 预扣费失败,回滚令牌额度
if s.tokenConsumed > 0 && !s.relayInfo.IsPlayground {
if rollbackErr := model.IncreaseTokenQuota(s.relayInfo.TokenId, s.relayInfo.TokenKey, s.tokenConsumed); rollbackErr != nil {
common.SysLog(fmt.Sprintf("error rolling back token quota (userId=%d, tokenId=%d, amount=%d, fundingErr=%s): %s",
common.SysLog(common.LogText("error rolling back token quota (userId=%d, tokenId=%d, amount=%d, fundingErr=%s): %s",
s.relayInfo.UserId, s.relayInfo.TokenId, s.tokenConsumed, err.Error(), rollbackErr.Error()))
}
s.tokenConsumed = 0
@@ -230,13 +234,13 @@ func (s *BillingSession) preConsume(c *gin.Context, quota int) *types.NewAPIErro
userQuota = 0
}
return types.NewErrorWithStatusCode(
fmt.Errorf("用户额度不足, 剩余额度: %s", logger.FormatQuota(userQuota)),
errors.New(i18n.T(c, i18n.MsgQuotaUserInsufficient, map[string]any{"Remaining": logger.FormatQuota(userQuota)})),
types.ErrorCodeInsufficientUserQuota, http.StatusForbidden,
types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
}
errMsg := err.Error()
if strings.Contains(errMsg, "no active subscription") || strings.Contains(errMsg, "subscription quota insufficient") {
return types.NewErrorWithStatusCode(fmt.Errorf("订阅额度不足或未配置订阅: %s", errMsg), types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
return types.NewErrorWithStatusCode(errors.New(i18n.T(c, i18n.MsgQuotaSubscriptionInsufficient, map[string]any{"Error": errMsg})), types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
}
return types.NewError(err, types.ErrorCodeUpdateDataError, types.ErrOptionWithSkipRetry())
}
@@ -275,7 +279,7 @@ func (s *BillingSession) reserveFunding(delta int, requireAvailableQuota bool) e
case *SubscriptionFunding:
if err := model.PostConsumeUserSubscriptionDelta(funding.subscriptionId, int64(delta)); err != nil {
return types.NewErrorWithStatusCode(
fmt.Errorf("订阅额度不足或未配置订阅: %s", err.Error()),
errors.New(i18n.Translate(s.lang, i18n.MsgQuotaSubscriptionInsufficient, map[string]any{"Error": err.Error()})),
types.ErrorCodeInsufficientUserQuota,
http.StatusForbidden,
types.ErrOptionWithSkipRetry(),
@@ -292,13 +296,13 @@ func (s *BillingSession) rollbackFundingReserve(delta int) {
switch funding := s.funding.(type) {
case *WalletFunding:
if err := model.IncreaseUserQuota(funding.userId, delta, false); err != nil {
common.SysLog("error rolling back wallet funding reserve: " + err.Error())
common.SysLog(common.LogText("error rolling back wallet funding reserve: %s", err.Error()))
} else {
funding.consumed -= delta
}
case *SubscriptionFunding:
if err := model.PostConsumeUserSubscriptionDelta(funding.subscriptionId, -int64(delta)); err != nil {
common.SysLog("error rolling back subscription funding reserve: " + err.Error())
common.SysLog(common.LogText("error rolling back subscription funding reserve: %s", err.Error()))
}
}
}
@@ -320,8 +324,8 @@ func (s *BillingSession) shouldTrust(c *gin.Context) bool {
return false
}
trustQuota := common.GetTrustQuota()
if trustQuota <= 0 {
trustQuota := operation_setting.GetQuotaSetting().TrustQuotaUSD * common.QuotaPerUnit
if trustQuota <= 0 || math.IsNaN(trustQuota) || math.IsInf(trustQuota, 0) {
return false
}
@@ -329,7 +333,7 @@ func (s *BillingSession) shouldTrust(c *gin.Context) bool {
tokenTrusted := s.relayInfo.TokenUnlimited
if !tokenTrusted {
tokenQuota := c.GetInt("token_quota")
tokenTrusted = tokenQuota > trustQuota
tokenTrusted = float64(tokenQuota) > trustQuota
}
if !tokenTrusted {
return false
@@ -337,7 +341,7 @@ func (s *BillingSession) shouldTrust(c *gin.Context) bool {
switch s.funding.Source() {
case BillingSourceWallet:
return s.relayInfo.UserQuota > trustQuota
return float64(s.relayInfo.UserQuota) > trustQuota
case BillingSourceSubscription:
// 订阅不能启用信任旁路。原因:
// 1. PreConsumeUserSubscription 要求 amount>0 来创建预扣记录并锁定订阅
@@ -389,13 +393,13 @@ func NewBillingSession(c *gin.Context, relayInfo *relaycommon.RelayInfo, preCons
}
if userQuota <= 0 {
return nil, types.NewErrorWithStatusCode(
fmt.Errorf("用户额度不足, 剩余额度: %s", logger.FormatQuota(userQuota)),
errors.New(i18n.T(c, i18n.MsgQuotaUserInsufficient, map[string]any{"Remaining": logger.FormatQuota(userQuota)})),
types.ErrorCodeInsufficientUserQuota, http.StatusForbidden,
types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
}
if userQuota-preConsumedQuota < 0 {
return nil, types.NewErrorWithStatusCode(
fmt.Errorf("预扣费额度失败, 用户剩余额度: %s, 需要预扣费额度: %s", logger.FormatQuota(userQuota), logger.FormatQuota(preConsumedQuota)),
errors.New(i18n.T(c, i18n.MsgQuotaPreConsumeInsufficient, map[string]any{"Remaining": logger.FormatQuota(userQuota), "Required": logger.FormatQuota(preConsumedQuota)})),
types.ErrorCodeInsufficientUserQuota, http.StatusForbidden,
types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
}
@@ -404,6 +408,7 @@ func NewBillingSession(c *gin.Context, relayInfo *relaycommon.RelayInfo, preCons
session := &BillingSession{
relayInfo: relayInfo,
funding: &WalletFunding{userId: relayInfo.UserId},
lang: i18n.StatedLang(c),
}
if apiErr := session.preConsume(c, preConsumedQuota); apiErr != nil {
return nil, apiErr
@@ -418,6 +423,7 @@ func NewBillingSession(c *gin.Context, relayInfo *relaycommon.RelayInfo, preCons
}
session := &BillingSession{
relayInfo: relayInfo,
lang: i18n.StatedLang(c),
funding: &SubscriptionFunding{
requestId: relayInfo.RequestId,
userId: relayInfo.UserId,
+54
View File
@@ -1,7 +1,10 @@
package service
import (
"slices"
"github.com/QuantumNous/new-api/model"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
)
@@ -23,6 +26,57 @@ func effectiveBillingUsage(usage *dto.Usage) *dto.Usage {
return usage
}
type claudeCacheReadAdjustment struct {
Rule string `json:"rule"`
Model string `json:"model"`
OriginalInputTokens int `json:"original_input_tokens"`
InputTokens int `json:"input_tokens"`
CacheReadTokens int `json:"cache_read_tokens"`
Status string `json:"status"`
}
// normalizeClaudeCacheReadUsage runs once after canonical upstream usage is
// selected. Only the billing copy changes; downstream responses and the
// original upstream snapshot retain their protocol values.
func normalizeClaudeCacheReadUsage(info *relaycommon.RelayInfo, usage *dto.Usage) (*dto.Usage, *claudeCacheReadAdjustment) {
if info == nil || info.ChannelMeta == nil || usage == nil ||
usage.UsageSemantic != dto.BillingUsageSemanticAnthropic ||
usage.BillingUsage == nil || usage.BillingUsage.Estimated || usage.BillingUsage.ClaudeUsage == nil {
return usage, nil
}
upstreamModel := info.GetUpstreamModelName()
if info.ResponseModel != nil && info.ResponseModel.UpstreamModel != "" {
// Some handlers replace UpstreamModelName with the response model.
// Match the name sent upstream, never the provider's response declaration.
upstreamModel = info.ResponseModel.UpstreamModel
}
if !slices.Contains(info.ChannelOtherSettings.ClaudeInputIncludesCacheReadModels, upstreamModel) {
return usage, nil
}
cached := usage.PromptTokensDetails.CachedTokens
adjustment := &claudeCacheReadAdjustment{
Rule: "claude_input_includes_cache_read_models", Model: upstreamModel,
OriginalInputTokens: usage.PromptTokens, InputTokens: usage.PromptTokens,
CacheReadTokens: cached, Status: "invalid_usage",
}
if usage.PromptTokens < 0 || cached < 0 || cached > usage.PromptTokens ||
usage.CompletionTokens < 0 || usage.InputTokens < cached || usage.TotalTokens < cached {
return usage, adjustment
}
if cached == 0 {
return usage, nil
}
adjusted := *usage
adjusted.PromptTokens -= cached
adjusted.InputTokens -= cached
adjusted.TotalTokens -= cached
adjusted.BillingUsage = dto.CloneBillingUsage(usage.BillingUsage)
adjusted.BillingUsage.ClaudeUsage.InputTokens = adjusted.PromptTokens
adjustment.InputTokens = adjusted.PromptTokens
adjustment.Status = "applied"
return &adjusted, adjustment
}
func usageBillingPathForLog(isLocalCountTokens bool, usage *dto.Usage) string {
effectiveUsage, ok := usageFromBillingUsage(usage)
if !ok {
-19
View File
@@ -1,19 +0,0 @@
package service
import "strings"
const CandyOutputTokens = 4096
const CandyPrompt = `在一个黑色的袋子里放有三种口味的糖果,每种糖果有两种不同的形状(圆形和五角星形,不同的形状靠手感可以分辨)。现已知不同口味的糖和不同形状的数量统计如下表。参赛者需要在活动前决定摸出的糖果数目,那么,最少取出多少个糖果才能保证手中同时拥有不同形状的苹果味和桃子味的糖?(同时手中有圆形苹果味匹配五角星桃子味糖果,或者有圆形桃子味匹配五角星苹果味糖果都满足要求)
苹果味 桃子味 西瓜味
圆形 7 9 8
五角星形 7 6 4
只输出最终答案的阿拉伯数字,不要输出解释或其他内容。`
func CandyQuality(answer string) string {
if strings.TrimSpace(answer) == "21" {
return "normal"
}
return "degraded"
}
+12 -9
View File
@@ -5,6 +5,7 @@ import (
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
@@ -18,7 +19,7 @@ func formatNotifyType(channelId int, status int) string {
func shouldCloseActiveWebSocketsAfterDisable(channelId int) bool {
channel, err := model.GetChannelById(channelId, true)
if err != nil {
common.SysLog(fmt.Sprintf("failed to check channel status before closing active websockets: channel_id=%d, error=%v", channelId, err))
common.SysLog(common.LogText("failed to check channel status before closing active websockets: channel_id=%d, error=%v", channelId, err))
return true
}
return channel.Status != common.ChannelStatusEnabled
@@ -26,11 +27,11 @@ func shouldCloseActiveWebSocketsAfterDisable(channelId int) bool {
// disable & notify
func DisableChannel(channelError types.ChannelError, reason string) {
common.SysLog(fmt.Sprintf("通道「%s」(#%d)发生错误,准备禁用,原因:%s", channelError.ChannelName, channelError.ChannelId, common.LocalLogPreview(reason)))
common.SysLog(common.LogText("channel %q (#%d) failed, disabling it, reason: %s", channelError.ChannelName, channelError.ChannelId, common.LocalLogPreview(reason)))
// 检查是否启用自动禁用功能
if !channelError.AutoBan {
common.SysLog(fmt.Sprintf("通道「%s」(#%d)未启用自动禁用功能,跳过禁用操作", channelError.ChannelName, channelError.ChannelId))
common.SysLog(common.LogText("channel %q (#%d) has automatic disabling turned off, skipping", channelError.ChannelName, channelError.ChannelId))
return
}
@@ -39,18 +40,20 @@ func DisableChannel(channelError types.ChannelError, reason string) {
if shouldCloseActiveWebSocketsAfterDisable(channelError.ChannelId) {
CloseActiveWebSocketsForChannel(channelError.ChannelId, ChannelDisabledCloseReason)
}
subject := fmt.Sprintf("通道「%s」(#%d)已被禁用", channelError.ChannelName, channelError.ChannelId)
content := fmt.Sprintf("通道「%s」(#%d)已被禁用,原因:%s", channelError.ChannelName, channelError.ChannelId, reason)
NotifyRootUser(formatNotifyType(channelError.ChannelId, common.ChannelStatusAutoDisabled), subject, content)
params := map[string]any{"Name": channelError.ChannelName, "Id": channelError.ChannelId, "Reason": reason}
NotifyChannelWatchers(formatNotifyType(channelError.ChannelId, common.ChannelStatusAutoDisabled), func(lang string) (string, string) {
return i18n.Translate(lang, i18n.MsgChannelNotifyDisabledSubject, params), i18n.Translate(lang, i18n.MsgChannelNotifyDisabledContent, params)
})
}
}
func EnableChannel(channelId int, usingKey string, channelName string) {
success := model.UpdateChannelStatus(channelId, usingKey, common.ChannelStatusEnabled, "")
if success {
subject := fmt.Sprintf("通道「%s」(#%d)已被启用", channelName, channelId)
content := fmt.Sprintf("通道「%s」(#%d)已被启用", channelName, channelId)
NotifyRootUser(formatNotifyType(channelId, common.ChannelStatusEnabled), subject, content)
params := map[string]any{"Name": channelName, "Id": channelId}
NotifyChannelWatchers(formatNotifyType(channelId, common.ChannelStatusEnabled), func(lang string) (string, string) {
return i18n.Translate(lang, i18n.MsgChannelNotifyEnabledSubject, params), i18n.Translate(lang, i18n.MsgChannelNotifyEnabledContent, params)
})
}
}
+10 -11
View File
@@ -1,7 +1,6 @@
package service
import (
"fmt"
"hash/fnv"
"maps"
"regexp"
@@ -146,7 +145,7 @@ func GetChannelAffinityCacheStats() ChannelAffinityCacheStats {
keys, err := cache.Keys()
if err != nil {
common.SysError(fmt.Sprintf("channel affinity cache list keys failed: err=%v", err))
common.SysError(common.LogText("channel affinity cache list keys failed: err=%v", err))
keys = nil
}
total := len(keys)
@@ -202,12 +201,12 @@ func ClearChannelAffinityCacheAll() int {
cache := getChannelAffinityCache()
keys, err := cache.Keys()
if err != nil {
common.SysError(fmt.Sprintf("channel affinity cache list keys failed: err=%v", err))
common.SysError(common.LogText("channel affinity cache list keys failed: err=%v", err))
keys = nil
}
if len(keys) > 0 {
if _, err := cache.DeleteMany(keys); err != nil {
common.SysError(fmt.Sprintf("channel affinity cache delete many failed: err=%v", err))
common.SysError(common.LogText("channel affinity cache delete many failed: err=%v", err))
}
}
return len(keys)
@@ -216,12 +215,12 @@ func ClearChannelAffinityCacheAll() int {
func ClearChannelAffinityCacheByRuleName(ruleName string) (int, error) {
ruleName = strings.TrimSpace(ruleName)
if ruleName == "" {
return 0, fmt.Errorf("rule_name 不能为空")
return 0, common.NewMessage("rule_name cannot be empty")
}
setting := operation_setting.GetChannelAffinitySetting()
if setting == nil {
return 0, fmt.Errorf("channel_affinity_setting 未初始化")
return 0, common.NewMessage("Channel affinity settings are not initialized")
}
var matchedRule *operation_setting.ChannelAffinityRule
@@ -234,10 +233,10 @@ func ClearChannelAffinityCacheByRuleName(ruleName string) (int, error) {
break
}
if matchedRule == nil {
return 0, fmt.Errorf("未知规则名称")
return 0, common.NewMessage("Unknown rule name")
}
if !matchedRule.IncludeRuleName {
return 0, fmt.Errorf("该规则未启用 include_rule_name,无法按规则清空缓存")
return 0, common.NewMessage("This rule does not have include_rule_name enabled, so its cache cannot be cleared by rule")
}
cache := getChannelAffinityCache()
@@ -620,7 +619,7 @@ func GetPreferredChannelByAffinity(c *gin.Context, modelName string, usingGroup
cache := getChannelAffinityCache()
channelID, found, err := cache.Get(cacheKeySuffix)
if err != nil {
common.SysError(fmt.Sprintf("channel affinity cache get failed: key=%s, err=%v", cacheKeyFull, err))
common.SysError(common.LogText("channel affinity cache get failed: key=%s, err=%v", cacheKeyFull, err))
return 0, false
}
if found {
@@ -672,7 +671,7 @@ func ClearCurrentChannelAffinityCache(c *gin.Context) bool {
cache := getChannelAffinityCache()
deleted, err := cache.DeleteMany([]string{cacheKey})
if err != nil {
common.SysError(fmt.Sprintf("channel affinity cache delete current failed: err=%v", err))
common.SysError(common.LogText("channel affinity cache delete current failed: err=%v", err))
return false
}
c.Set(ginKeyChannelAffinitySkipRetry, false)
@@ -764,7 +763,7 @@ func RecordChannelAffinity(c *gin.Context, channelID int) {
}
cache := getChannelAffinityCache()
if err := cache.SetWithTTL(cacheKey, channelID, time.Duration(ttlSeconds)*time.Second); err != nil {
common.SysError(fmt.Sprintf("channel affinity cache set failed: key=%s, err=%v", cacheKey, err))
common.SysError(common.LogText("channel affinity cache set failed: key=%s, err=%v", cacheKey, err))
}
}
+1 -1
View File
@@ -467,7 +467,7 @@ func SelectChannelForRequest(c *gin.Context, modelName string, retry *RetryParam
constraints := GetChannelConstraints(c)
if pin, found, overridden := constraints.ResolvedPin(); found {
for _, lost := range overridden {
logger.LogWarn(c, fmt.Sprintf(
logger.LogWarn(c, common.LogText(
"channel pin overridden: winning_source=%s winning_channel_id=%d overridden_source=%s overridden_channel_id=%d",
pin.Source, pin.ChannelId, lost.Source, lost.ChannelId,
))
+13 -5
View File
@@ -19,18 +19,26 @@ func FetchCodexChannelModels(channel *model.Channel) ([]string, error) {
return nil, fmt.Errorf("codex channel does not support multi-key model discovery")
}
client, err := NewProxyHttpClient(channel.GetSetting().Proxy)
if err != nil {
return nil, err
}
setting := channel.GetSetting()
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
defer cancel()
clientVersion, err := GetLatestCodexClientVersion(ctx, client)
// The GitHub release lookup is not an upstream request; it always verifies
// certificates, whatever the channel's TLS setting says.
releaseClient, err := GetHttpClientWithProxy(setting.Proxy)
if err != nil {
return nil, err
}
clientVersion, err := GetLatestCodexClientVersion(ctx, releaseClient)
if err != nil {
return nil, fmt.Errorf("failed to get Codex client version: %w", err)
}
client, err := GetHttpClientWithProxySettings(setting.Proxy, setting)
if err != nil {
return nil, err
}
baseURL := channel.GetBaseURL()
if baseURL == "" {
baseURL = constant.GetChannelBaseURL(constant.ChannelTypeCodex)
+5 -6
View File
@@ -2,7 +2,6 @@ package service
import (
"context"
"fmt"
"strings"
"sync"
"sync/atomic"
@@ -39,7 +38,7 @@ func StartCodexCredentialAutoRefreshTask() {
}
gopool.Go(func() {
logger.LogInfo(context.Background(), fmt.Sprintf("codex credential auto-refresh task started: tick=%s threshold=%s", codexCredentialRefreshTickInterval, codexCredentialRefreshThreshold))
logger.LogInfo(context.Background(), common.LogText("codex credential auto-refresh task started: tick=%s threshold=%s", codexCredentialRefreshTickInterval, codexCredentialRefreshThreshold))
ticker := time.NewTicker(codexCredentialRefreshTickInterval)
defer ticker.Stop()
@@ -79,7 +78,7 @@ func runCodexCredentialAutoRefreshOnce() {
Offset(offset).
Find(&channels).Error
if err != nil {
logger.LogError(ctx, fmt.Sprintf("codex credential auto-refresh: query channels failed: %v", err))
logger.LogError(ctx, common.LogText("codex credential auto-refresh: query channels failed: %v", err))
return
}
if len(channels) == 0 {
@@ -121,12 +120,12 @@ func runCodexCredentialAutoRefreshOnce() {
newKey, _, err := RefreshCodexChannelCredential(refreshCtx, ch.Id, CodexCredentialRefreshOptions{ResetCaches: false})
cancel()
if err != nil {
logger.LogWarn(ctx, fmt.Sprintf("codex credential auto-refresh: channel_id=%d name=%s refresh failed: %v", ch.Id, ch.Name, err))
logger.LogWarn(ctx, common.LogText("codex credential auto-refresh: channel_id=%d name=%s refresh failed: %v", ch.Id, ch.Name, err))
continue
}
refreshed++
logger.LogInfo(ctx, fmt.Sprintf("codex credential auto-refresh: channel_id=%d name=%s refreshed, expires_at=%s", ch.Id, ch.Name, newKey.Expired))
logger.LogInfo(ctx, common.LogText("codex credential auto-refresh: channel_id=%d name=%s refreshed, expires_at=%s", ch.Id, ch.Name, newKey.Expired))
}
}
@@ -134,7 +133,7 @@ func runCodexCredentialAutoRefreshOnce() {
func() {
defer func() {
if r := recover(); r != nil {
logger.LogWarn(ctx, fmt.Sprintf("codex credential auto-refresh: InitChannelCache panic: %v", r))
logger.LogWarn(ctx, common.LogText("codex credential auto-refresh: InitChannelCache panic: %v", r))
}
}()
model.InitChannelCache()
+5 -3
View File
@@ -37,15 +37,17 @@ func DetectionPricing(userGroup, group, kind string) (DetectionPrice, int, error
return DetectionPrice{}, 0, errors.New("invalid group ratio")
}
price := DetectionPrice{Snapshot: billingexpr.BillingSnapshot{BillingMode: "tiered_expr", ModelName: fingerprint.Model, ExprString: expression, ExprHash: billingexpr.ExprHashString(expression), GroupRatio: ratio, QuotaPerUnit: common.QuotaPerUnit, ExprVersion: 1}}
effort := "low"
output, calls := fingerprint.OutputTokens, fingerprint.ProbeCount
if kind == "benchmark" {
output, calls = benchmark.OutputTokens, benchmark.ProbeCount
} else if kind == "candy" {
output, calls = CandyOutputTokens, 1
} else if kind == "vision" {
output, calls = VisionOutputTokens, 1
effort = "medium"
} else if kind != "fingerprint" {
return price, 0, errors.New("invalid detection type")
}
price.Request.Body = []byte(fmt.Sprintf(`{"model":"gpt-6-astra","reasoning":{"effort":"low"},"max_output_tokens":%d}`, output))
price.Request.Body = []byte(fmt.Sprintf(`{"model":"gpt-6-astra","reasoning":{"effort":"%s"},"max_output_tokens":%d}`, effort, output))
cost, err := billingexpr.DetectionCostBound(expression, 32768, float64(output))
if err != nil {
return price, 0, err
+10 -3
View File
@@ -83,16 +83,17 @@ func ReconcileAutomaticDetectionSchedules(ctx context.Context, channels []model.
}
} else {
var runs []model.DetectionRun
if err := model.DB.WithContext(ctx).Select("id", "channel_id", "status", "quality", "settled", "finished_at").Where("channel_id IN ? AND user_id = 0 AND kind = ? AND (finished_at >= ? OR settled = ?)", ids, kind, cutoff, false).Find(&runs).Error; err != nil {
if err := model.DB.WithContext(ctx).Select("id", "channel_id", "status", "quality", "settled", "finished_at", "kind", "error", "result").Where("channel_id IN ? AND user_id = 0 AND request_key NOT LIKE 'chain:%' AND kind = ? AND (finished_at >= ? OR settled = ?)", ids, kind, cutoff, false).Find(&runs).Error; err != nil {
return nil, nil, err
}
for _, run := range runs {
run = NormalizeVisionRun(run)
if !run.Settled {
active[run.ChannelID] = true
continue
}
outcome := "unknown"
if run.Status == "succeeded" && (run.Quality == "normal" || run.Quality == "degraded") {
if run.Status == "succeeded" && (run.Quality == "normal" || run.Quality == "degraded" || run.Quality == "incomplete") {
outcome = run.Quality
}
results = append(results, automaticDetectionResult{run.ChannelID, run.ID, run.FinishedAt, outcome})
@@ -115,12 +116,18 @@ func ReconcileAutomaticDetectionSchedules(ctx context.Context, channels []model.
}
seed.LastFinishedAt, seed.LastResultID = result.at, result.id
seed.IntervalMinutes = 15
if result.outcome == "normal" || result.outcome == "degraded" {
if result.outcome == "normal" || result.outcome == "degraded" || result.outcome == "incomplete" {
seed.LastValid = result.outcome
seed.IntervalMinutes = 30
}
seed.NextDueAt = result.at + int64(seed.IntervalMinutes)*60
}
if kind == "ciallo" {
seed.IntervalMinutes = 45
if seed.LastFinishedAt > 0 {
seed.NextDueAt = seed.LastFinishedAt + 45*60
}
}
if err := model.EnsureDetectionSchedule(ctx, seed); err != nil {
return nil, nil, err
}
+45 -30
View File
@@ -11,13 +11,14 @@ import (
)
type DetectionHistory struct {
Closest string `json:"closest,omitempty"`
Weight float64 `json:"relative_weight,omitempty"`
ID string `json:"id"`
At int64 `json:"at"`
Status string `json:"status"`
Quality string `json:"quality"`
UserInitiated bool `json:"user_initiated"`
FirstTokenMS *float64 `json:"first_token_ms,omitempty"`
Closest string `json:"closest,omitempty"`
Weight float64 `json:"relative_weight,omitempty"`
ID string `json:"id"`
At int64 `json:"at"`
Status string `json:"status"`
Quality string `json:"quality"`
UserInitiated bool `json:"user_initiated"`
}
type BenchmarkSummary struct {
LatestValid *DetectionHistory `json:"-"`
@@ -27,6 +28,7 @@ type BenchmarkSummary struct {
WindowHours int `json:"window_hours"`
Normal int `json:"normal"`
Degraded int `json:"degraded"`
Incomplete int `json:"incomplete"`
Unknown int `json:"unknown"`
Failed int `json:"failed"`
Running int `json:"running"`
@@ -38,7 +40,8 @@ type ChannelDetectionEvidence struct {
}
type GroupDetectionSummary struct {
ProbeHistory []ProbeHistoryPoint
Candy *BenchmarkSummary
Ciallo *BenchmarkSummary
Vision *BenchmarkSummary
Benchmark *BenchmarkSummary
Status int
At int64
@@ -78,7 +81,8 @@ func GroupDetectionSummaries(ctx context.Context, now time.Time, fingerprints ma
groups := map[string][]int{}
byChannel := map[int][]string{}
benchmarks := map[string]*BenchmarkSummary{}
candies := map[string]*BenchmarkSummary{}
visions := map[string]*BenchmarkSummary{}
ciallos := map[string]*BenchmarkSummary{}
evidence := map[int]ChannelDetectionEvidence{}
observe := func(id int, at int64, success bool) {
if at > evidence[id].At {
@@ -97,13 +101,14 @@ func GroupDetectionSummaries(ctx context.Context, now time.Time, fingerprints ma
}
}
}
candyRanges, err := DetectionScheduleRanges(ctx, channels, "candy")
visionRanges, err := DetectionScheduleRanges(ctx, channels, "vision")
if err != nil {
return nil, err
}
for group, summary := range benchmarks {
r := candyRanges[group]
candies[group] = &BenchmarkSummary{Channels: summary.Channels, IntervalMinutes: r[0], IntervalMaxMinutes: r[1], WindowHours: 6, History: []DetectionHistory{}}
ciallos[group] = &BenchmarkSummary{Channels: summary.Channels, IntervalMinutes: 45, IntervalMaxMinutes: 45, WindowHours: 6, History: []DetectionHistory{}}
r := visionRanges[group]
visions[group] = &BenchmarkSummary{Channels: summary.Channels, IntervalMinutes: r[0], IntervalMaxMinutes: r[1], WindowHours: 6, History: []DetectionHistory{}}
}
var tasks []model.SystemTask
ranges, err := DetectionScheduleRanges(ctx, channels, "benchmark")
@@ -153,21 +158,28 @@ func GroupDetectionSummaries(ctx context.Context, now time.Time, fingerprints ma
}
var calls []model.DetectionCall
if len(ids) > 0 {
if err = model.DB.WithContext(ctx).Select("run_id", "finished_at", "available").Where("run_id IN ? AND state = ?", ids, "finished").Find(&calls).Error; err != nil {
if err = model.DB.WithContext(ctx).Select("run_id", "finished_at", "available", "first_token_ms").Where("run_id IN ? AND state = ?", ids, "finished").Find(&calls).Error; err != nil {
return nil, err
}
}
runChannels := map[string]int{}
firstTokenTotals := map[string]float64{}
firstTokenCounts := map[string]int{}
for _, run := range runs {
runChannels[run.ID] = run.ChannelID
}
for _, call := range calls {
if call.FirstTokenMS != nil && *call.FirstTokenMS >= 0 && call.Available != nil && *call.Available {
firstTokenTotals[call.RunID] += float64(*call.FirstTokenMS)
firstTokenCounts[call.RunID]++
}
if call.Available != nil {
observe(runChannels[call.RunID], call.FinishedAt, *call.Available)
}
}
for _, run := range runs {
if run.Kind != "fingerprint" && run.Kind != "benchmark" && run.Kind != "candy" {
run = NormalizeVisionRun(run)
if run.Kind != "fingerprint" && run.Kind != "benchmark" && run.Kind != "vision" && run.Kind != "ciallo" {
continue
}
for _, group := range byChannel[run.ChannelID] {
@@ -208,19 +220,26 @@ func GroupDetectionSummaries(ctx context.Context, now time.Time, fingerprints ma
continue
}
s := benchmarks[group]
if run.Kind == "candy" {
s = candies[group]
if run.Kind == "vision" {
s = visions[group]
}
if run.Kind == "ciallo" {
s = ciallos[group]
}
if s == nil {
continue
}
history := DetectionHistory{ID: run.ID, At: run.CreatedAt, Status: run.Status, Quality: run.Quality, UserInitiated: run.UserID > 0}
if count := firstTokenCounts[run.ID]; count > 0 {
average := firstTokenTotals[run.ID] / float64(count)
history.FirstTokenMS = &average
}
var metadata fingerprint.Result
if common.UnmarshalJsonStr(run.Result, &metadata) == nil {
history.Closest, history.Weight = metadata.Closest, metadata.Weight
}
s.History = append(s.History, history)
if run.Status == "succeeded" && (run.Quality == "normal" || run.Quality == "degraded") && (s.LatestValid == nil || history.At > s.LatestValid.At) {
if run.Status == "succeeded" && (run.Quality == "normal" || run.Quality == "degraded" || run.Quality == "incomplete") && (s.LatestValid == nil || history.At > s.LatestValid.At) {
s.LatestValid = &history
}
switch run.Status {
@@ -234,6 +253,8 @@ func GroupDetectionSummaries(ctx context.Context, now time.Time, fingerprints ma
s.Normal++
case "degraded":
s.Degraded++
case "incomplete":
s.Incomplete++
default:
s.Unknown++
}
@@ -246,9 +267,6 @@ func GroupDetectionSummaries(ctx context.Context, now time.Time, fingerprints ma
s.Rate = &rate
}
sort.SliceStable(s.History, func(i, j int) bool { return s.History[i].At > s.History[j].At })
if len(s.History) > 12 {
s.History = s.History[:12]
}
}
results := map[string]GroupDetectionSummary{}
for group, ids := range groups {
@@ -256,18 +274,15 @@ func GroupDetectionSummaries(ctx context.Context, now time.Time, fingerprints ma
s := benchmarks[group]
if s != nil {
sort.SliceStable(s.History, func(i, j int) bool { return s.History[i].At > s.History[j].At })
if len(s.History) > 12 {
s.History = s.History[:12]
}
}
candy := candies[group]
if candy != nil {
sort.SliceStable(candy.History, func(i, j int) bool { return candy.History[i].At > candy.History[j].At })
if len(candy.History) > 12 {
candy.History = candy.History[:12]
}
vision := visions[group]
if vision != nil {
sort.SliceStable(vision.History, func(i, j int) bool { return vision.History[i].At > vision.History[j].At })
}
results[group] = GroupDetectionSummary{Candy: candy, Benchmark: s, Status: status, At: at, ProbeHistory: GroupProbeHistory(ids, tasks)}
if s := ciallos[group]; s != nil {
sort.SliceStable(s.History, func(i, j int) bool { return s.History[i].At > s.History[j].At })
}
results[group] = GroupDetectionSummary{Ciallo: ciallos[group], Vision: vision, Benchmark: s, Status: status, At: at, ProbeHistory: GroupProbeHistory(ids, tasks)}
}
return results, nil
}
+2 -2
View File
@@ -51,7 +51,7 @@ func DoWorkerRequest(req *WorkerRequest) (*http.Response, error) {
func DoDownloadRequest(originUrl string, reason ...string) (resp *http.Response, err error) {
if system_setting.EnableWorker() {
common.SysLog(fmt.Sprintf("downloading file from worker: %s, reason: %s", originUrl, strings.Join(reason, ", ")))
common.SysLog(common.LogText("downloading file from worker: %s, reason: %s", originUrl, strings.Join(reason, ", ")))
req := &WorkerRequest{
URL: originUrl,
Key: system_setting.WorkerValidKey,
@@ -63,7 +63,7 @@ func DoDownloadRequest(originUrl string, reason ...string) (resp *http.Response,
return nil, fmt.Errorf("request reject: %v", err)
}
common.SysLog(fmt.Sprintf("downloading from origin: %s, reason: %s", common.MaskSensitiveInfo(originUrl), strings.Join(reason, ", ")))
common.SysLog(common.LogText("downloading from origin: %s, reason: %s", common.MaskSensitiveInfo(originUrl), strings.Join(reason, ", ")))
return GetSSRFProtectedHTTPClient().Get(originUrl)
}
}
+17 -9
View File
@@ -11,6 +11,7 @@ import (
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/model"
)
@@ -53,7 +54,7 @@ func ValidateAccountEmail(email string) (string, error) {
return email, nil
}
func StartEmailBinding(identity AuthIdentity, authorization *model.AuthFlowAuthorization, email string) (*EmailBindingData, error) {
func StartEmailBinding(identity AuthIdentity, authorization *model.AuthFlowAuthorization, email string, lang string) (*EmailBindingData, error) {
email, err := ValidateAccountEmail(email)
if err != nil {
return nil, err
@@ -80,19 +81,19 @@ func StartEmailBinding(identity AuthIdentity, authorization *model.AuthFlowAutho
if err != nil {
return nil, err
}
if err := sendEmailBindingCodes(state, codes); err != nil {
if err := sendEmailBindingCodes(lang, state, codes); err != nil {
// A partially delivered pair must not leave a usable change request.
_, _ = model.ConsumeAuthFlow(token, model.AuthFlowMatch{Purpose: model.AuthFlowPurposeEmailBinding, UserId: identity.UserID, SessionId: identity.SessionID})
return nil, ErrEmailBindingDelivery
}
data := emailBindingData(token, flow, &state)
if state.CurrentEmail != "" && !requireOld {
data.NotificationWarning = NotifyAccountSecurityChange(state.CurrentEmail, "A change of your email address was requested") != nil
data.NotificationWarning = NotifyAccountSecurityChange(lang, state.CurrentEmail, i18n.MsgSecurityEventEmailChangeRequested) != nil
}
return data, nil
}
func ResendAccountEmailBinding(identity AuthIdentity, token string) (*EmailBindingData, error) {
func ResendAccountEmailBinding(identity AuthIdentity, token string, lang string) (*EmailBindingData, error) {
_, state, err := model.GetEmailBinding(identity, token)
if err != nil {
return nil, err
@@ -118,7 +119,7 @@ func ResendAccountEmailBinding(identity AuthIdentity, token string) (*EmailBindi
if err != nil {
return nil, err
}
if err := sendEmailBindingCodes(*state, codes); err != nil {
if err := sendEmailBindingCodes(lang, *state, codes); err != nil {
_, _ = model.ConsumeAuthFlow(token, model.AuthFlowMatch{Purpose: model.AuthFlowPurposeEmailBinding, UserId: identity.UserID, SessionId: identity.SessionID})
return nil, ErrEmailBindingDelivery
}
@@ -177,14 +178,21 @@ func generateEmailBindingCodes(requireOld bool) (emailBindingCodes, error) {
return codes, nil
}
func sendEmailBindingCodes(state model.EmailBindingState, codes emailBindingCodes) error {
subject := common.SystemName + " — Confirm your email address"
content := fmt.Sprintf("<p>Confirm linking this email address to your account.</p><p>Verification code: <strong>%s</strong></p><p>This code expires in 10 minutes. If you did not request this change, do not share this code.</p>", html.EscapeString(codes.New))
func sendEmailBindingCodes(lang string, state model.EmailBindingState, codes emailBindingCodes) error {
subject := i18n.Translate(lang, i18n.MsgEmailBindingSubject, map[string]any{"SystemName": common.SystemName})
minutes := map[string]any{"Minutes": int(model.EmailBindingTTL / time.Minute)}
content := fmt.Sprintf("<p>%s</p><p>%s</p><p>%s</p>",
i18n.Translate(lang, i18n.MsgEmailBindingNewIntro),
i18n.Translate(lang, i18n.MsgEmailBindingCode, map[string]any{"Code": html.EscapeString(codes.New)}),
i18n.Translate(lang, i18n.MsgEmailBindingNewValidity, minutes))
if err := common.SendEmail(subject, state.Email, content); err != nil {
return err
}
if codes.Old != "" {
content = fmt.Sprintf("<p>A change to your account email address was requested. Confirm replacing your current address.</p><p>Verification code: <strong>%s</strong></p><p>This code expires in 10 minutes. If you did not request this change, do not share this code and contact your administrator.</p>", html.EscapeString(codes.Old))
content = fmt.Sprintf("<p>%s</p><p>%s</p><p>%s</p>",
i18n.Translate(lang, i18n.MsgEmailBindingOldIntro),
i18n.Translate(lang, i18n.MsgEmailBindingCode, map[string]any{"Code": html.EscapeString(codes.Old)}),
i18n.Translate(lang, i18n.MsgEmailBindingOldValidity, minutes))
return common.SendEmail(subject, state.CurrentEmail, content)
}
return nil
+9 -7
View File
@@ -64,8 +64,8 @@ func ClaudeErrorWrapper(err error, code string, statusCode int) *dto.ClaudeError
lowerText := strings.ToLower(text)
if !strings.HasPrefix(lowerText, "get file base64 from url") {
if strings.Contains(lowerText, "post") || strings.Contains(lowerText, "dial") || strings.Contains(lowerText, "http") {
common.SysLog(fmt.Sprintf("error: %s", text))
text = "请求上游地址失败"
common.SysLog(common.LogText("error: %s", text))
text = "request to the upstream address failed"
}
}
claudeError := types.ClaudeError{
@@ -107,7 +107,7 @@ func RelayErrorHandler(ctx context.Context, resp *http.Response, showBodyWhenFai
if showBodyWhenFail {
newApiErr.Err = buildErrWithBody("")
} else {
logger.LogError(ctx, fmt.Sprintf("bad response status code %d, body: %s", resp.StatusCode, responseBodyPreview))
logger.LogError(ctx, common.LogText("bad response status code %d, body: %s", resp.StatusCode, responseBodyPreview))
newApiErr.Err = fmt.Errorf("bad response status code %d", resp.StatusCode)
}
return
@@ -128,7 +128,7 @@ func RelayErrorHandler(ctx context.Context, resp *http.Response, showBodyWhenFai
if message == "" {
// The body parsed as JSON but carried no usable error message; log the
// raw body so the upstream failure remains diagnosable.
logger.LogError(ctx, fmt.Sprintf("bad response status code %d with empty error message, body: %s", resp.StatusCode, responseBodyPreview))
logger.LogError(ctx, common.LogText("bad response status code %d with empty error message, body: %s", resp.StatusCode, responseBodyPreview))
}
newApiErr = types.NewOpenAIError(errors.New(message), types.ErrorCodeBadResponseStatusCode, resp.StatusCode)
if showBodyWhenFail {
@@ -197,13 +197,15 @@ func TaskErrorWrapperLocal(err error, code string, statusCode int) *taskdto.Task
return openaiErr
}
func TaskErrorWrapper(err error, code string, statusCode int) *taskdto.TaskError {
// TaskErrorWrapper masks the upstream addresses of the task channel in network
// errors; see common.MaskSensitiveInfo.
func TaskErrorWrapper(err error, code string, statusCode int, upstreamURLs ...string) *taskdto.TaskError {
text := err.Error()
lowerText := strings.ToLower(text)
if strings.Contains(lowerText, "post") || strings.Contains(lowerText, "dial") || strings.Contains(lowerText, "http") {
common.SysLog(fmt.Sprintf("error: %s", text))
common.SysLog(common.LogText("error: %s", text))
//text = "请求上游地址失败"
text = common.MaskSensitiveInfo(text)
text = common.MaskSensitiveInfo(text, upstreamURLs...)
}
//避免暴露内部错误
taskError := &taskdto.TaskError{
+2 -2
View File
@@ -23,13 +23,13 @@ import (
func GetFileTypeFromUrl(c *gin.Context, url string, reason ...string) (string, error) {
response, err := DoDownloadRequest(url, []string{"get_mime_type", strings.Join(reason, ", ")}...)
if err != nil {
common.SysLog(fmt.Sprintf("fail to get file type from url: %s, error: %s", url, err.Error()))
common.SysLog(common.LogText("fail to get file type from url: %s, error: %s", url, err.Error()))
return "", err
}
defer response.Body.Close()
if response.StatusCode != 200 {
logger.LogError(c, fmt.Sprintf("failed to download file from %s, status code: %d", url, response.StatusCode))
logger.LogError(c, common.LogText("failed to download file from %s, status code: %d", url, response.StatusCode))
return "", fmt.Errorf("failed to download file, status code: %d", response.StatusCode)
}
+1 -1
View File
@@ -198,7 +198,7 @@ func loadFromURL(c *gin.Context, url string, reason ...string) (*types.CachedFil
diskPath, err := writeToDiskCache(base64Data)
if err != nil {
// 磁盘缓存失败,回退到内存
logger.LogWarn(c, fmt.Sprintf("Failed to write to disk cache, falling back to memory: %v", err))
logger.LogWarn(c, common.LogText("Failed to write to disk cache, falling back to memory: %v", err))
cachedData = types.NewMemoryCachedData(base64Data, mimeType, int64(len(fileBytes)))
} else {
cachedData = types.NewDiskCachedData(diskPath, mimeType, int64(len(fileBytes)))
+3 -2
View File
@@ -32,9 +32,10 @@ func GetUserUsableGroups(userGroup string) map[string]string {
}
}
}
// 如果userGroup不在UserUsableGroups中,返回UserUsableGroups + userGroup
// 如果userGroup不在UserUsableGroups中,返回UserUsableGroups + userGroup。
// 描述是英文源文本,由 web 控制台翻译;管理员配置的描述原样显示。
if _, ok := groupsCopy[userGroup]; !ok {
groupsCopy[userGroup] = "用户分组"
groupsCopy[userGroup] = "User group"
}
}
return groupsCopy
+32 -20
View File
@@ -10,9 +10,34 @@ import (
var groupDegradationSnapshot struct {
sync.Mutex
updated time.Time
latest map[string]AstraSamplingHistory
benchmark map[string]DetectionHistory
updated time.Time
latest map[string]AstraSamplingHistory
vision map[string]DetectionHistory
ciallo map[string]DetectionHistory
}
func requestGroupQuality(c *gin.Context) (map[string]AstraSamplingHistory, map[string]DetectionHistory, map[string]DetectionHistory) {
if c.GetBool("pool_degradation_unavailable") {
return nil, nil, nil
}
latest, found := c.Get("pool_group_fingerprints")
vision, _ := c.Get("pool_group_visions")
ciallo, _ := c.Get("pool_group_ciallos")
if !found {
rounds, qualities, ciallos, err := cachedGroupQuality(c.Request.Context())
if err != nil {
c.Set("pool_degradation_unavailable", true)
return nil, nil, nil
}
latest, vision, ciallo = rounds, qualities, ciallos
c.Set("pool_group_fingerprints", latest)
c.Set("pool_group_visions", vision)
c.Set("pool_group_ciallos", ciallo)
}
rounds, _ := latest.(map[string]AstraSamplingHistory)
visions, _ := vision.(map[string]DetectionHistory)
ciallos, _ := ciallo.(map[string]DetectionHistory)
return rounds, visions, ciallos
}
func requestGroupDegraded(c *gin.Context, modelName, group string) bool {
@@ -27,22 +52,9 @@ func requestGroupDegraded(c *gin.Context, modelName, group string) bool {
if !skip {
return false
}
if c.GetBool("pool_degradation_unavailable") {
return true
}
latest, found := c.Get("pool_group_fingerprints")
benchmark, _ := c.Get("pool_group_benchmarks")
if !found {
rounds, qualities, err := cachedGroupQuality(c.Request.Context())
if err != nil {
c.Set("pool_degradation_unavailable", true)
return true
}
latest, benchmark = rounds, qualities
c.Set("pool_group_fingerprints", latest)
c.Set("pool_group_benchmarks", benchmark)
}
rounds, _ := latest.(map[string]AstraSamplingHistory)
rounds, visions, ciallos := requestGroupQuality(c)
h, ok := rounds[group]
return (ok && h.ValidSamples > 0 && h.Closest != "" && h.Closest != fingerprint.Model)
return c.GetBool("pool_degradation_unavailable") ||
(ok && h.ValidSamples > 0 && h.Closest != "" && h.Closest != fingerprint.Model) ||
visions[group].Quality == "degraded" || ciallos[group].Quality == "degraded"
}
+109 -2
View File
@@ -3,6 +3,7 @@ package service
import (
"fmt"
"os"
"strings"
"testing"
"time"
@@ -56,6 +57,8 @@ func TestGroupDegradationDatabaseMatrix(t *testing.T) {
common.SetDatabaseTypes(oldType, oldLogType)
groupDegradationSnapshot.updated = time.Time{}
groupDegradationSnapshot.latest = nil
groupDegradationSnapshot.vision = nil
groupDegradationSnapshot.ciallo = nil
conn, err := db.DB()
require.NoError(t, err)
require.NoError(t, conn.Close())
@@ -110,15 +113,76 @@ func TestGroupDegradationDatabaseMatrix(t *testing.T) {
require.NoError(t, db.Model(&pelican).Update("quality", "normal").Error)
groupDegradationSnapshot.updated = time.Time{}
assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
vision := model.DetectionRun{ID: "vision-degraded", RequestKey: "vision-degraded", ChannelID: 1, Kind: "vision", CreatedAt: now - 5, Settled: true, Status: "succeeded", Quality: "degraded", Result: `{"closest":"石头"}`}
require.NoError(t, db.Create(&vision).Error)
InvalidateGroupQualityCache()
assert.True(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"), "fingerprint recovery must not clear a failed vision test")
recent, err := RecentDegradedGroups(ctx().Request.Context())
require.NoError(t, err)
assert.True(t, recent["default"])
failedVision := model.DetectionRun{ID: "vision-error", RequestKey: "vision-error", ChannelID: 1, Kind: "vision", CreatedAt: now - 1, Settled: true, Status: "failed", Quality: "normal"}
require.NoError(t, db.Create(&failedVision).Error)
InvalidateGroupQualityCache()
assert.True(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"), "transport errors cannot clear the latest valid vision result")
require.NoError(t, db.Model(&failedVision).Update("quality", "degraded").Error)
require.NoError(t, db.Model(&vision).Updates(map[string]any{"quality": "normal", "result": `{"closest":"这是音箱"}`}).Error)
InvalidateGroupQualityCache()
assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
recent, err = RecentDegradedGroups(ctx().Request.Context())
require.NoError(t, err)
assert.False(t, recent["default"])
require.NoError(t, db.Model(&vision).Updates(map[string]any{"quality": "degraded", "created_at": now - 8*3600, "result": `{"closest":"石头"}`}).Error)
InvalidateGroupQualityCache()
assert.True(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"), "routing keeps the latest valid result")
recent, err = RecentDegradedGroups(ctx().Request.Context())
require.NoError(t, err)
assert.False(t, recent["default"], "recent tag expires after the existing six-hour window")
legacyAnswer := `{"closest":"我这边没收到可查看的图片或文件,暂时无法判断是什么。请重新上传图片"}`
require.NoError(t, db.Model(&vision).Updates(map[string]any{"result": legacyAnswer, "status": "failed", "error": "vision_keyword_missing"}).Error)
InvalidateGroupQualityCache()
unavailable, err := VisionUnavailableGroups(ctx().Request.Context())
require.NoError(t, err)
assert.True(t, unavailable["default"], "legacy answers remain effective beyond the recent window and later transport failures")
assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
recent, err = RecentDegradedGroups(ctx().Request.Context())
require.NoError(t, err)
assert.False(t, recent["default"])
require.NoError(t, db.Model(&vision).Update("created_at", now-2).Error)
summaries, err := GroupDetectionSummaries(ctx().Request.Context(), time.Unix(now, 0), nil, nil)
require.NoError(t, err)
summary := summaries["default"].Vision
require.NotNil(t, summary)
assert.Equal(t, 1, summary.Incomplete)
assert.Equal(t, "incomplete", summary.LatestValid.Quality)
var unchanged model.DetectionRun
require.NoError(t, db.First(&unchanged, "id = ?", vision.ID).Error)
assert.Equal(t, "degraded", unchanged.Quality, "reading history must not rewrite stored results")
ciallo := model.DetectionRun{ID: "ciallo-degraded", RequestKey: "ciallo-degraded", ChannelID: 1, Kind: "ciallo", CreatedAt: now, Settled: true, Status: "succeeded", Quality: "degraded"}
require.NoError(t, db.Create(&ciallo).Error)
InvalidateGroupQualityCache()
assert.True(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"), "incomplete vision must not clear Ciallo degradation")
require.NoError(t, db.Model(&vision).Updates(map[string]any{"status": "succeeded", "result": `{"closest":"这是音箱"}`, "created_at": now}).Error)
InvalidateGroupQualityCache()
unavailable, err = VisionUnavailableGroups(ctx().Request.Context())
require.NoError(t, err)
assert.False(t, unavailable["default"])
assert.True(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"), "vision recovery must not clear Ciallo degradation")
require.NoError(t, db.Model(&ciallo).Update("quality", "normal").Error)
InvalidateGroupQualityCache()
assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
require.NoError(t, db.Delete(&manual).Error)
require.NoError(t, model.SaveStatusPageGroup(ctx().Request.Context(), "default", nil))
groupDegradationSnapshot.updated = time.Time{}
assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
unavailable, err = VisionUnavailableGroups(ctx().Request.Context())
require.NoError(t, err)
assert.False(t, unavailable["default"], "administrator reset removes the vision conclusion")
})
}
}
func TestGroupDegradationUsesLatestValidBeforeHistoryTruncation(t *testing.T) {
func TestGroupDegradationRetainsFullHistoryForNotifications(t *testing.T) {
now := time.Now()
channels := []model.Channel{{Id: 1, Group: "default", Models: fingerprint.Model, Status: 1}}
rounds := []AstraSamplingRound{{ID: "valid", ChannelID: 1, StartedAt: now.Unix() - 100, FinishedAt: now.Unix() - 90, Result: fingerprint.Result{Engine: fingerprint.Version, Status: "mismatch", Closest: "gpt-5.5", ValidSamples: 1}}}
@@ -131,7 +195,7 @@ func TestGroupDegradationUsesLatestValidBeforeHistoryTruncation(t *testing.T) {
summary := AggregateAstraSampling(channels, tasks, now)["default"]
require.NotNil(t, summary.LatestValid)
assert.Equal(t, "gpt-5.5", summary.LatestValid.Closest)
assert.Len(t, summary.History, 12)
assert.Len(t, summary.History, 16)
rounds = append(rounds, AstraSamplingRound{ID: "recovered", ChannelID: 1, StartedAt: now.Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "match", Closest: fingerprint.Model, ValidSamples: 2}})
body, err = common.Marshal(AstraSamplingState{Rounds: rounds})
require.NoError(t, err)
@@ -183,6 +247,15 @@ func TestGroupDegradationOverridesExistingPoolAffinity(t *testing.T) {
RecordPoolAffinity(first)
bound, _ := PreferredPoolChannel(ctx(false), fingerprint.Model)
require.NotNil(t, bound)
image := ctx(false)
image.Request = httptest.NewRequest("POST", "/v1/responses", strings.NewReader(`{"input":[{"role":"user","content":[{"type":"input_image","image_url":"data:image/png;base64,eA=="}]}]}`))
image.Request.Header.Set("Session-Id", "degradation-session")
image.Set("pool_group_visions", map[string]DetectionHistory{"default": {Quality: "incomplete"}})
selected, group, err = SelectChannelForRequest(image, fingerprint.Model, &RetryParam{Ctx: image, TokenGroup: "default", ModelName: fingerprint.Model})
require.Nil(t, err)
require.NotNil(t, selected)
assert.Equal(t, 6401, selected.Id, "incomplete vision must preserve the existing session affinity")
assert.Equal(t, "default", group)
next := ctx(true)
selected, group, err = SelectChannelForRequest(next, fingerprint.Model, &RetryParam{Ctx: next, TokenGroup: "default", ModelName: fingerprint.Model})
require.Nil(t, err)
@@ -191,6 +264,40 @@ func TestGroupDegradationOverridesExistingPoolAffinity(t *testing.T) {
assert.Equal(t, "vip", group)
}
func TestVisionIncompleteRoutingAcrossModelsAndPolicies(t *testing.T) {
db := setupChannelSelectAutoGroupsTest(t)
for i, group := range []string{"default", "vip"} {
createChannelSelectAutoGroupsChannel(t, db, 6601+i, group, fingerprint.Model)
require.NoError(t, db.Model(&model.Channel{}).Where("id = ?", 6601+i).Update("models", fingerprint.Model+",gpt-6-sol").Error)
require.NoError(t, db.Create(&model.Ability{Group: group, Model: "gpt-6-sol", ChannelId: 6601 + i, Enabled: true, Weight: 100}).Error)
}
model.InitChannelCache()
ResetGroupAvailability()
t.Cleanup(ResetGroupAvailability)
for _, mode := range []string{"smart", "custom"} {
for _, allow := range []bool{false, true} {
for _, modelName := range []string{fingerprint.Model, "gpt-6-sol"} {
for _, image := range []bool{false, true} {
c := setupGroupPoolContext(t)
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: mode, AllowDegraded: allow, SkipDegradedGroups: !allow, Groups: []model.TokenGroupEntry{{Group: "default"}, {Group: "vip", Priority: 1}}})
if image {
c.Request = httptest.NewRequest("POST", "/v1/responses", strings.NewReader(`{"input":[{"role":"user","content":[{"type":"input_image","image_url":"data:image/png;base64,eA=="}]}]}`))
}
c.Set("pool_group_fingerprints", map[string]AstraSamplingHistory{})
visions := map[string]DetectionHistory{"default": {Quality: "incomplete"}, "vip": {Quality: "incomplete"}}
c.Set("pool_group_visions", visions)
channel, _, err := selectTokenRoutingPool(&RetryParam{Ctx: c, TokenGroup: "default", ModelName: modelName})
require.NoError(t, err)
require.NotNil(t, channel)
assert.Equal(t, 6601, channel.Id, "mode=%s allow=%v model=%s image=%v", mode, allow, modelName, image)
assert.False(t, poolGroupSkipped(c, modelName, "default"), "incomplete vision must not exclude retries or cooldown recovery")
}
}
}
}
}
func TestGroupDegradationSkipsPoolGroupsAndRespectsUnknownRecoveryAndOptOut(t *testing.T) {
db := setupChannelSelectAutoGroupsTest(t)
createChannelSelectAutoGroupsChannel(t, db, 6301, "default", fingerprint.Model)
+29
View File
@@ -0,0 +1,29 @@
package service
import (
"context"
"fmt"
"sync"
"time"
"github.com/QuantumNous/new-api/model"
"github.com/samber/hot"
)
var groupModelPerformanceCache = hot.NewHotCache[string, model.GroupModelPerformance](hot.LRU, 2000).WithTTL(time.Minute).Build()
var groupModelPerformanceMu sync.Mutex
// Share one snapshot between status display and automatic throughput ranking.
func GroupModelPerformance(ctx context.Context, group, modelName string) (model.GroupModelPerformance, error) {
key := fmt.Sprintf("%p/%p/%q/%q", model.DB, model.LOG_DB, group, modelName)
groupModelPerformanceMu.Lock()
defer groupModelPerformanceMu.Unlock()
if value, found, err := groupModelPerformanceCache.Get(key); err == nil && found {
return value, nil
}
value, err := model.GetGroupModelPerformance(ctx, group, modelName)
if err == nil {
groupModelPerformanceCache.Set(key, value)
}
return value, err
}
+58 -14
View File
@@ -11,20 +11,21 @@ import (
)
// Routing evidence has no statistical window. Each detector recovers separately.
func readGroupQuality(ctx context.Context, now time.Time) (map[string]AstraSamplingHistory, map[string]DetectionHistory, error) {
func readGroupQuality(ctx context.Context, now time.Time) (map[string]AstraSamplingHistory, map[string]DetectionHistory, map[string]DetectionHistory, error) {
settings, err := model.GetStatusPageGroups(ctx)
if err != nil {
return nil, nil, err
return nil, nil, nil, err
}
channels, err := samplingChannels(ctx, false)
if err != nil {
return nil, nil, err
return nil, nil, nil, err
}
channelGroups := map[int][]string{}
for _, channel := range channels {
channelGroups[channel.Id] = channel.GetGroups()
}
latest, benchmarks := map[string]AstraSamplingHistory{}, map[string]DetectionHistory{}
latest, visions := map[string]AstraSamplingHistory{}, map[string]DetectionHistory{}
ciallos := map[string]DetectionHistory{}
observe := func(channelID int, h AstraSamplingHistory) {
if h.At <= 0 || h.At > now.Unix() || h.ValidSamples <= 0 || h.Closest == "" || (h.Status != "match" && h.Status != "mismatch") {
return
@@ -37,7 +38,7 @@ func readGroupQuality(ctx context.Context, now time.Time) (map[string]AstraSampl
}
var tasks []model.SystemTask
if err = model.DB.WithContext(ctx).Select("state").Where("type = ?", AstraSamplingTaskType).Find(&tasks).Error; err != nil {
return nil, nil, err
return nil, nil, nil, err
}
for _, task := range tasks {
var state AstraSamplingState
@@ -52,10 +53,20 @@ func readGroupQuality(ctx context.Context, now time.Time) (map[string]AstraSampl
}
}
var runs []model.DetectionRun
if err = model.DB.WithContext(ctx).Select("id", "channel_id", "kind", "created_at", "quality", "result").Where("status = ? AND settled = ? AND created_at <= ?", "succeeded", true, now.Unix()).Find(&runs).Error; err != nil {
return nil, nil, err
if err = model.DB.WithContext(ctx).Select("id", "channel_id", "kind", "created_at", "quality", "result", "status", "error").Where("(status = ? OR (kind IN ? AND status = ?)) AND settled = ? AND created_at <= ?", "succeeded", []string{"vision", "ciallo"}, "failed", true, now.Unix()).Find(&runs).Error; err != nil {
return nil, nil, nil, err
}
for _, run := range runs {
run = NormalizeVisionRun(run)
if run.Status != "succeeded" {
continue
}
if run.Kind == "ciallo" {
var result fingerprint.Result
if common.UnmarshalJsonStr(run.Result, &result) == nil {
run.Quality = CialloQuality(result.Closest)
}
}
if run.Kind == "fingerprint" {
var result fingerprint.Result
if common.UnmarshalJsonStr(run.Result, &result) != nil || result.Engine != fingerprint.Version {
@@ -64,30 +75,42 @@ func readGroupQuality(ctx context.Context, now time.Time) (map[string]AstraSampl
observe(run.ChannelID, AstraSamplingHistory{ID: run.ID, At: run.CreatedAt, Status: AstraSamplingResultStatus(result), ValidSamples: result.ValidSamples, Closest: result.Closest})
}
if (run.Kind == "vision" || run.Kind == "ciallo") && (run.Quality == "normal" || run.Quality == "degraded" || (run.Kind == "vision" && run.Quality == "incomplete")) {
target := visions
if run.Kind == "ciallo" {
target = ciallos
}
for _, group := range channelGroups[run.ChannelID] {
if run.CreatedAt > settings[group].ResetBefore && run.CreatedAt > target[group].At {
target[group] = DetectionHistory{ID: run.ID, At: run.CreatedAt, Status: run.Status, Quality: run.Quality}
}
}
}
}
return latest, benchmarks, nil
return latest, visions, ciallos, nil
}
func cachedGroupQuality(ctx context.Context) (map[string]AstraSamplingHistory, map[string]DetectionHistory, error) {
func cachedGroupQuality(ctx context.Context) (map[string]AstraSamplingHistory, map[string]DetectionHistory, map[string]DetectionHistory, error) {
groupDegradationSnapshot.Lock()
defer groupDegradationSnapshot.Unlock()
now := time.Now()
if now.Sub(groupDegradationSnapshot.updated) >= 15*time.Second {
ctx, cancel := context.WithTimeout(ctx, 2*time.Second)
latest, benchmarks, err := readGroupQuality(ctx, now)
latest, visions, ciallos, err := readGroupQuality(ctx, now)
cancel()
if err == nil {
groupDegradationSnapshot.latest, groupDegradationSnapshot.benchmark = latest, benchmarks
groupDegradationSnapshot.latest, groupDegradationSnapshot.vision = latest, visions
groupDegradationSnapshot.ciallo = ciallos
} else if groupDegradationSnapshot.latest == nil {
return nil, nil, errors.New("group quality temporarily unavailable")
return nil, nil, nil, errors.New("group quality temporarily unavailable")
}
groupDegradationSnapshot.updated = now
}
return groupDegradationSnapshot.latest, groupDegradationSnapshot.benchmark, nil
return groupDegradationSnapshot.latest, groupDegradationSnapshot.vision, groupDegradationSnapshot.ciallo, nil
}
func RecentDegradedGroups(ctx context.Context) (map[string]bool, error) {
latest, _, err := cachedGroupQuality(ctx)
latest, visions, ciallos, err := cachedGroupQuality(ctx)
if err != nil {
return nil, err
}
@@ -98,6 +121,27 @@ func RecentDegradedGroups(ctx context.Context) (map[string]bool, error) {
result[group] = true
}
}
for _, detector := range []map[string]DetectionHistory{visions, ciallos} {
for group, h := range detector {
if h.At >= cutoff && h.Quality == "degraded" {
result[group] = true
}
}
}
return result, nil
}
func VisionUnavailableGroups(ctx context.Context) (map[string]bool, error) {
_, visions, _, err := cachedGroupQuality(ctx)
if err != nil {
return nil, err
}
result := map[string]bool{}
for group, h := range visions {
if h.Quality == "incomplete" {
result[group] = true
}
}
return result, nil
}
+2 -1
View File
@@ -27,7 +27,8 @@ func TestQualitySwitchesAndFailedReadCache(t *testing.T) {
InvalidateGroupQualityCache()
t.Cleanup(func() {
groupDegradationSnapshot.latest = nil
groupDegradationSnapshot.benchmark = nil
groupDegradationSnapshot.vision = nil
groupDegradationSnapshot.ciallo = nil
InvalidateGroupQualityCache()
})
ctx := func(allow, child bool) *gin.Context {
+2 -2
View File
@@ -19,7 +19,7 @@ func CloseResponseBodyGracefully(httpResponse *http.Response) {
}
err := httpResponse.Body.Close()
if err != nil {
common.SysError("failed to close response body: " + err.Error())
common.SysError(common.LogText("failed to close response body: %s", err.Error()))
}
}
@@ -73,7 +73,7 @@ func IOCopyBytesGracefully(c *gin.Context, src *http.Response, data []byte) {
_, err := io.Copy(c.Writer, body)
if err != nil {
logger.LogError(c, fmt.Sprintf("failed to copy response body: %s", err.Error()))
logger.LogError(c, common.LogText("failed to copy response body: %s", err.Error()))
}
c.Writer.Flush()
}
+5 -2
View File
@@ -116,7 +116,7 @@ func newRelayHTTPTransport() *http.Transport {
}
transport.ForceAttemptHTTP2 = true
if common.TLSInsecureSkipVerify {
transport.TLSClientConfig = common.InsecureTLSConfig
transport.TLSClientConfig = common.InsecureTLSConfig.Clone()
}
return transport
}
@@ -176,7 +176,7 @@ func warnLegacyProxyURLOnce(config *proxyURLConfig) {
}
logger.LogWarn(
context.Background(),
fmt.Sprintf(
common.LogText(
"legacy proxy URL suffix ignored at runtime: scheme=%s host=%s; update the channel proxy setting",
config.parsedURL.Scheme,
config.parsedURL.Host,
@@ -333,6 +333,9 @@ func newTransportFactory(proxyURL *url.URL, tlsConfig *tls.Config) (func() *http
}
func newHTTPClientFromPolicy(policy HTTPTransportPolicy, proxyURL *url.URL, tlsConfig *tls.Config) (*http.Client, error) {
if policy.TLSInsecureSkipVerify {
tlsConfig = common.InsecureTLSConfig
}
factory, err := newTransportFactory(proxyURL, tlsConfig)
if err != nil {
return nil, err
+80
View File
@@ -6,6 +6,7 @@ import (
"encoding/pem"
"fmt"
"io"
"net"
"net/http"
"net/http/httptest"
"sync"
@@ -297,6 +298,75 @@ func TestHTTPClientCachePolicyAndCompatibility(t *testing.T) {
proxyHTTP1, err := GetHttpClientWithProxySettings(proxyA, dto.ChannelSettings{HTTPProtocol: dto.HTTPProtocolHTTP1})
require.NoError(t, err)
assert.NotSame(t, clientA, proxyHTTP1)
insecure, err := GetHttpClientWithProxySettings("", dto.ChannelSettings{TLSInsecureSkipVerify: true})
require.NoError(t, err)
assert.NotSame(t, aware, insecure, "skip-verify channels must not share the verifying client")
insecureHTTP1, err := GetHttpClientWithProxySettings("", dto.ChannelSettings{HTTPProtocol: dto.HTTPProtocolHTTP1, TLSInsecureSkipVerify: true})
require.NoError(t, err)
assert.NotSame(t, http1, insecureHTTP1)
proxyInsecure, err := GetHttpClientWithProxySettings(proxyA, dto.ChannelSettings{TLSInsecureSkipVerify: true})
require.NoError(t, err)
assert.NotSame(t, clientA, proxyInsecure)
}
func TestChannelTLSInsecureSkipVerifyAcceptsUntrustedUpstream(t *testing.T) {
initDefaultHTTPClientFixture(t)
upstream := startHTTP2TLSServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
}))
connectProxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodConnect {
w.WriteHeader(http.StatusMethodNotAllowed)
return
}
target, err := net.Dial("tcp", r.Host)
if err != nil {
w.WriteHeader(http.StatusBadGateway)
return
}
defer target.Close()
conn, _, err := http.NewResponseController(w).Hijack()
if err != nil {
return
}
defer conn.Close()
_, _ = conn.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n"))
go func() { _, _ = io.Copy(target, conn) }()
_, _ = io.Copy(conn, target)
}))
t.Cleanup(connectProxy.Close)
skip := dto.ChannelSettings{TLSInsecureSkipVerify: true}
testCases := []struct {
name string
proxy string
settings dto.ChannelSettings
wantProto int // zero means certificate verification must fail
}{
{name: "default direct verifies", settings: dto.ChannelSettings{}},
{name: "default proxy verifies", proxy: connectProxy.URL, settings: dto.ChannelSettings{}},
{name: "http1 verifies", settings: dto.ChannelSettings{HTTPProtocol: dto.HTTPProtocolHTTP1}},
{name: "skip direct", settings: skip, wantProto: 2},
{name: "skip proxy", proxy: connectProxy.URL, settings: skip, wantProto: 2},
{name: "skip http1", settings: dto.ChannelSettings{HTTPProtocol: dto.HTTPProtocolHTTP1, TLSInsecureSkipVerify: true}, wantProto: 1},
{name: "skip sharded", settings: dto.ChannelSettings{HTTP2ConnectionShards: 3, TLSInsecureSkipVerify: true}, wantProto: 2},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
client, err := GetHttpClientWithProxySettings(tc.proxy, tc.settings)
require.NoError(t, err)
resp, err := client.Get(upstream.URL)
if tc.wantProto == 0 {
var verifyErr *tls.CertificateVerificationError
require.ErrorAs(t, err, &verifyErr)
return
}
require.NoError(t, err)
assert.Equal(t, tc.wantProto, resp.ProtoMajor)
drainClose(t, resp)
})
}
}
func TestHTTPClientCacheConcurrentGetOrCreate(t *testing.T) {
@@ -366,6 +436,8 @@ func TestInvalidateProxyClientClosesAllPolicyVariants(t *testing.T) {
require.NoError(t, err)
shardedClient, err := GetHttpClientWithProxySettings(proxyURL, dto.ChannelSettings{HTTP2ConnectionShards: 2})
require.NoError(t, err)
insecureClient, err := GetHttpClientWithProxySettings(proxyURL, dto.ChannelSettings{TLSInsecureSkipVerify: true})
require.NoError(t, err)
InvalidateProxyClient(proxyURL)
@@ -375,10 +447,13 @@ func TestInvalidateProxyClientClosesAllPolicyVariants(t *testing.T) {
require.NoError(t, err)
afterSharded, err := GetHttpClientWithProxySettings(proxyURL, dto.ChannelSettings{HTTP2ConnectionShards: 2})
require.NoError(t, err)
afterInsecure, err := GetHttpClientWithProxySettings(proxyURL, dto.ChannelSettings{TLSInsecureSkipVerify: true})
require.NoError(t, err)
assert.NotSame(t, defaultClient, afterDefault)
assert.NotSame(t, http1Client, afterHTTP1)
assert.NotSame(t, shardedClient, afterSharded)
assert.NotSame(t, insecureClient, afterInsecure)
}
func TestResetProxyClientCacheKeepsDefaultPointerAndRecreatesVariants(t *testing.T) {
@@ -390,6 +465,8 @@ func TestResetProxyClientCacheKeepsDefaultPointerAndRecreatesVariants(t *testing
require.NoError(t, err)
proxyClient, err := GetHttpClientWithProxy("http://reset-proxy.example:8080")
require.NoError(t, err)
insecureClient, err := GetHttpClientWithProxySettings("", dto.ChannelSettings{TLSInsecureSkipVerify: true})
require.NoError(t, err)
ResetProxyClientCache()
@@ -404,9 +481,12 @@ func TestResetProxyClientCacheKeepsDefaultPointerAndRecreatesVariants(t *testing
require.NoError(t, err)
afterProxy, err := GetHttpClientWithProxy("http://reset-proxy.example:8080")
require.NoError(t, err)
afterInsecure, err := GetHttpClientWithProxySettings("", dto.ChannelSettings{TLSInsecureSkipVerify: true})
require.NoError(t, err)
assert.NotSame(t, http1Client, afterHTTP1)
assert.NotSame(t, shardedClient, afterSharded)
assert.NotSame(t, proxyClient, afterProxy)
assert.NotSame(t, insecureClient, afterInsecure)
}
func TestResetProxyClientCacheClosesDefaultIdlePool(t *testing.T) {
+7 -4
View File
@@ -6,6 +6,7 @@ import (
"strings"
"sync"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/relaykit/dto"
)
@@ -13,8 +14,9 @@ import (
// HTTPTransportPolicy is the runtime-normalized outbound HTTP transport policy
// for a channel. Unknown or out-of-range stored values are clamped safely.
type HTTPTransportPolicy struct {
Protocol string // dto.HTTPProtocolAuto or dto.HTTPProtocolHTTP1
Shards int // 1..dto.MaxHTTP2ConnectionShards
Protocol string // dto.HTTPProtocolAuto or dto.HTTPProtocolHTTP1
Shards int // 1..dto.MaxHTTP2ConnectionShards
TLSInsecureSkipVerify bool
}
var httpTransportPolicyWarnings sync.Map
@@ -30,6 +32,7 @@ func defaultHTTPTransportPolicy() HTTPTransportPolicy {
// Invalid stored values never panic; they clamp to defaults and warn once per bad value.
func NormalizeHTTPTransportPolicy(settings dto.ChannelSettings) HTTPTransportPolicy {
policy := defaultHTTPTransportPolicy()
policy.TLSInsecureSkipVerify = settings.TLSInsecureSkipVerify
protocol := strings.ToLower(strings.TrimSpace(settings.HTTPProtocol))
switch protocol {
@@ -78,12 +81,12 @@ func warnHTTPTransportPolicyOnce(field, value string) {
}
logger.LogWarn(
context.Background(),
fmt.Sprintf("invalid channel http transport setting clamped: %s=%q", field, value),
common.LogText("invalid channel http transport setting clamped: %s=%q", field, value),
)
}
func (p HTTPTransportPolicy) cacheKeyPart() string {
return fmt.Sprintf("%s|%d", p.Protocol, p.Shards)
return fmt.Sprintf("%s|%d|%t", p.Protocol, p.Shards, p.TLSInsecureSkipVerify)
}
func (p HTTPTransportPolicy) String() string {
+7 -10
View File
@@ -2,7 +2,6 @@ package service
import (
"crypto/tls"
"fmt"
"net/http"
"strings"
"sync"
@@ -69,15 +68,13 @@ func (s *shardedRoundTripper) RoundTrip(req *http.Request) (*http.Response, erro
}
logger.LogDebug(
req.Context(),
fmt.Sprintf(
"http transport: host=%s protocol=%s shard=%d/%d policy=%s negotiated=%s",
host,
s.policy.Protocol,
idx,
s.n,
s.policy.cacheKeyPart(),
proto,
),
"http transport: host=%s protocol=%s shard=%d/%d policy=%s negotiated=%s",
host,
s.policy.Protocol,
idx,
s.n,
s.policy.cacheKeyPart(),
proto,
)
}
return resp, err
+5 -5
View File
@@ -30,7 +30,7 @@ func DecodeBase64ImageData(base64String string) (image.Config, string, string, e
// 将base64字符串解码为字节切片
decodedData, err := base64.StdEncoding.DecodeString(base64String)
if err != nil {
fmt.Println("Error: Failed to decode base64 string")
fmt.Println(common.LogText("Error: Failed to decode base64 string"))
return image.Config{}, "", "", fmt.Errorf("failed to decode base64 string: %s", err.Error())
}
@@ -119,7 +119,7 @@ func GetImageFromUrl(url string) (mimeType string, data string, err error) {
func DecodeUrlImageData(imageUrl string) (image.Config, string, error) {
response, err := DoDownloadRequest(imageUrl)
if err != nil {
common.SysLog(fmt.Sprintf("fail to get image from url: %s", err.Error()))
common.SysLog(common.LogText("fail to get image from url: %s", err.Error()))
return image.Config{}, "", err
}
defer response.Body.Close()
@@ -137,7 +137,7 @@ func DecodeUrlImageData(imageUrl string) (image.Config, string, error) {
var readData []byte
for _, limit := range []int64{1024 * 8, 1024 * 24, 1024 * 64} {
common.SysLog(fmt.Sprintf("try to decode image config with limit: %d", limit))
common.SysLog(common.LogText("try to decode image config with limit: %d", limit))
// 从response.Body读取更多的数据直到达到当前的限制
additionalData := make([]byte, limit-int64(len(readData)))
@@ -170,13 +170,13 @@ func getImageConfig(reader io.Reader) (image.Config, string, error) {
if err == nil {
return config, format, nil
}
common.SysLog(fmt.Sprintf("fail to decode image config(gif, jpg, png): %s", err.Error()))
common.SysLog(common.LogText("fail to decode image config(gif, jpg, png): %s", err.Error()))
config, err = webp.DecodeConfig(bytes.NewReader(data))
if err == nil {
return config, "webp", nil
}
common.SysLog(fmt.Sprintf("fail to decode image config(webp): %s", err.Error()))
common.SysLog(common.LogText("fail to decode image config(webp): %s", err.Error()))
// Try HEIF/HEIC: parse ISOBMFF ispe box for dimensions
if heifMime := detectHEIF(data); heifMime != "" {
+7 -11
View File
@@ -4,10 +4,8 @@ import (
"errors"
"fmt"
"net/http"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/pkg/billingexpr"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
@@ -18,7 +16,7 @@ import (
// PrepareImageBillingForRequest reserves the effective outbound image quantity
// before each attempt, including channel retries and parameter overrides. The
// client request body stays frozen; only the independent quantity is refreshed.
func PrepareImageBillingForRequest(c *gin.Context, info *relaycommon.RelayInfo, count int, promptExtend bool) *types.NewAPIError {
func PrepareImageBillingForRequest(c *gin.Context, info *relaycommon.RelayInfo, count int) *types.NewAPIError {
if count < 1 || count > dto.MaxImageN {
return types.NewErrorWithStatusCode(fmt.Errorf("image_count must be an integer between 1 and %d", dto.MaxImageN), types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
@@ -42,6 +40,9 @@ func PrepareImageBillingForRequest(c *gin.Context, info *relaycommon.RelayInfo,
return types.NewErrorWithStatusCode(runErr, types.ErrorCodeModelPriceError, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
}
beforeGroup := cost / 1_000_000 * snap.QuotaPerUnit
if trace.BillingUnit != billingexpr.BillingUnitRequest && snap.PreConsumeMultiplier != 0 {
beforeGroup *= snap.PreConsumeMultiplier
}
quota, err = billingexpr.QuotaRoundStrict(beforeGroup * info.PriceData.GroupRatioInfo.GroupRatio)
if err == nil {
snap.EstimatedImageCount = trace.ImageCount
@@ -54,17 +55,12 @@ func PrepareImageBillingForRequest(c *gin.Context, info *relaycommon.RelayInfo,
}
} else {
quantity := 1
if info.PriceData.UsePrice || info.ChannelType == constant.ChannelTypeAli {
if info.PriceData.UsePrice {
quantity = count
}
// Overwrite per-attempt ratios so a failed Ali attempt cannot leak its
// quantity or prompt-extension surcharge into another channel.
// Overwrite the per-attempt ratio so a failed attempt cannot leak its
// quantity into another channel.
info.PriceData.AddOtherRatio("n", float64(quantity))
extensionRatio := 1.0
if info.ChannelType == constant.ChannelTypeAli && strings.Contains(info.UpstreamModelName, "z-image") && promptExtend {
extensionRatio = common.ZImagePromptExtendMultiplier
}
info.PriceData.AddOtherRatio("prompt_extend", extensionRatio)
base := info.ImageQuotaBeforeGroup
if info.PriceData.UsePrice {
base = info.PriceData.ModelPrice * common.QuotaPerUnit
+6 -6
View File
@@ -18,11 +18,11 @@ func TestDedicatedImageGenerationPoolSelection(t *testing.T) {
cfg := model_setting.GetGlobalSettings()
oldGroup, oldModels := cfg.ImageGenerationGroup, cfg.ImageGenerationModels
t.Cleanup(func() { cfg.ImageGenerationGroup, cfg.ImageGenerationModels = oldGroup, oldModels })
cfg.ImageGenerationGroup, cfg.ImageGenerationModels = "example-images", []string{"gpt-image-2"}
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP","example-images":"Images"}`))
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":2,"example-images":1.25}`))
cfg.ImageGenerationGroup, cfg.ImageGenerationModels = "gpt-images", []string{"gpt-image-2"}
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP","gpt-images":"Images"}`))
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":2,"gpt-images":1.25}`))
createChannelSelectAutoGroupsChannel(t, db, 8001, "default", "gpt-image-2")
createChannelSelectAutoGroupsChannel(t, db, 8002, "example-images", "gpt-image-2")
createChannelSelectAutoGroupsChannel(t, db, 8002, "gpt-images", "gpt-image-2")
model.InitChannelCache()
ResetGroupAvailability()
t.Cleanup(ResetGroupAvailability)
@@ -33,12 +33,12 @@ func TestDedicatedImageGenerationPoolSelection(t *testing.T) {
common.SetContextKey(c, constant.ContextKeyUserGroup, "default")
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: mode, Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}}})
entries := ResolveRequestPool(c, "gpt-image-2")
require.Equal(t, []model.TokenGroupEntry{{Group: "example-images", Priority: 0}}, entries)
require.Equal(t, []model.TokenGroupEntry{{Group: "gpt-images", Priority: 0}}, entries)
channel, group, err := CacheGetRandomSatisfiedChannel(&RetryParam{Ctx: c, TokenGroup: "default", ModelName: "gpt-image-2", Retry: common.GetPointer(0)})
require.NoError(t, err)
require.NotNil(t, channel)
require.Equal(t, 8002, channel.Id)
require.Equal(t, "example-images", group)
require.Equal(t, "gpt-images", group)
})
}
c, _ := gin.CreateTestContext(httptest.NewRecorder())
+66
View File
@@ -0,0 +1,66 @@
package service
import (
"context"
"fmt"
"sync"
"time"
"github.com/QuantumNous/new-api/common"
)
const logExportLockTTL = time.Minute
const logExportLockRenewScript = `
if redis.call('GET', KEYS[1]) == ARGV[1] then
return redis.call('PEXPIRE', KEYS[1], ARGV[2])
end
return 0`
const logExportLockReleaseScript = `
if redis.call('GET', KEYS[1]) == ARGV[1] then
return redis.call('DEL', KEYS[1])
end
return 0`
// logExportLocks holds the running exports when Redis is disabled.
var logExportLocks sync.Map
// AcquireLogExportLock reserves the one log export a user may run at a time.
// With Redis the reservation spans all nodes and lapses within a minute if
// this node dies. ok is false while another export of the user is running;
// otherwise release must be called when the export ends.
func AcquireLogExportLock(userId int) (release func(), ok bool, err error) {
if !common.RedisEnabled {
if _, running := logExportLocks.LoadOrStore(userId, struct{}{}); running {
return nil, false, nil
}
return func() { logExportLocks.Delete(userId) }, true, nil
}
ctx := context.Background()
key := fmt.Sprintf("log_export:%d", userId)
owner := common.GetRandomString(16)
acquired, err := common.RDB.SetNX(ctx, key, owner, logExportLockTTL).Result()
if err != nil || !acquired {
return nil, false, err
}
stop := make(chan struct{})
go func() {
ticker := time.NewTicker(logExportLockTTL / 3)
defer ticker.Stop()
for {
select {
case <-stop:
return
case <-ticker.C:
common.RDB.Eval(ctx, logExportLockRenewScript, []string{key}, owner, logExportLockTTL.Milliseconds())
}
}
}()
return func() {
close(stop)
common.RDB.Eval(ctx, logExportLockReleaseScript, []string{key}, owner)
}, true, nil
}
+7 -4
View File
@@ -2,8 +2,8 @@ package service
import (
"encoding/base64"
"fmt"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
@@ -42,7 +42,7 @@ func attachQuotaSaturation(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, o
return
}
attachQuotaSaturationToOther(other, clamp)
logger.LogWarn(ctx, fmt.Sprintf("quota saturation on consume log: op=%s kind=%s original=%g clamped=%d user=%d model=%s",
logger.LogWarn(ctx, common.LogText("quota saturation on consume log: op=%s kind=%s original=%g clamped=%d user=%d model=%s",
clamp.Op, clamp.Kind, clamp.Original, clamp.Clamped, relayInfo.UserId, relayInfo.GetBillingModelName()))
}
@@ -74,6 +74,9 @@ func AppendRelayLogAdminInfo(ctx *gin.Context, relayInfo *relaycommon.RelayInfo,
other.SetAdmin("use_channel", ctx.GetStringSlice("use_channel"))
appendRequestDetails(ctx, relayInfo, other)
if relayInfo != nil {
if timing := relayInfo.UpstreamTiming.Snapshot(); timing != nil {
other.SetAdmin("upstream_timing", timing)
}
if billingModel := relayInfo.GetBillingModelName(); billingModel != "" && billingModel != relayInfo.OriginModelName {
other.SetAdmin("billing_model", billingModel)
}
@@ -112,6 +115,7 @@ func GenerateTextOtherInfo(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, m
other.SetPublic("model_price", modelPrice)
other.SetPublic("user_group_ratio", userGroupRatio)
other.SetPublic("frt", float64(relayInfo.FirstResponseTime.UnixMilli()-relayInfo.StartTime.UnixMilli()))
other.SetPublic("latency_ms", time.Since(relayInfo.StartTime).Milliseconds())
if relayInfo.ReasoningEffort != "" {
other.SetPublic("reasoning_effort", relayInfo.ReasoningEffort)
}
@@ -141,8 +145,7 @@ func AppendResponseModelLogInfo(relayInfo *relaycommon.RelayInfo, other *model.L
return
}
observation := relayInfo.ResponseModel
if !observation.Mismatch &&
observation.ReturnedModel == observation.RequestedModel &&
if observation.ReturnedModel == observation.RequestedModel &&
(observation.UpstreamModel == "" || observation.UpstreamModel == observation.RequestedModel) &&
(relayInfo.ChannelMeta == nil || !relayInfo.IsModelMapped) {
return
+4 -5
View File
@@ -100,7 +100,7 @@ func RefundMidjourneyQuota(ctx context.Context, task *model.Midjourney, reason s
}
if err := model.IncreaseUserQuota(task.UserId, quota, false); err != nil {
logger.LogWarn(ctx, fmt.Sprintf("退还 Midjourney 用户额度失败 task %s: %s", task.MjId, err.Error()))
logger.LogWarn(ctx, common.LogText("failed to refund Midjourney user quota, task %s: %s", task.MjId, err.Error()))
return false
}
@@ -108,7 +108,7 @@ func RefundMidjourneyQuota(ctx context.Context, task *model.Midjourney, reason s
tokenKey := resolveTokenKey(ctx, task.TokenId, task.MjId)
if tokenKey != "" {
if err := model.IncreaseTokenQuota(task.TokenId, tokenKey, quota); err != nil {
logger.LogWarn(ctx, fmt.Sprintf("退还 Midjourney 令牌额度失败 task %s: %s", task.MjId, err.Error()))
logger.LogWarn(ctx, common.LogText("failed to refund Midjourney token quota, task %s: %s", task.MjId, err.Error()))
}
}
}
@@ -122,7 +122,6 @@ func RefundMidjourneyQuota(ctx context.Context, task *model.Midjourney, reason s
model.RecordTaskBillingLog(model.RecordTaskBillingLogParams{
UserId: task.UserId,
LogType: model.LogTypeRefund,
Content: "",
ChannelId: billingChannelId,
ModelName: CovertMjpActionToModelName(task.Action),
Quota: quota,
@@ -132,7 +131,7 @@ func RefundMidjourneyQuota(ctx context.Context, task *model.Midjourney, reason s
task.Quota = 0
if err := task.UpdateBillingState(); err != nil {
logger.LogError(ctx, fmt.Sprintf("Midjourney 退款成功但清除 quota 失败 task %s: %s", task.MjId, err.Error()))
logger.LogError(ctx, common.LogText("Midjourney refund succeeded but clearing the quota failed, task %s: %s", task.MjId, err.Error()))
}
return true
}
@@ -348,7 +347,7 @@ func DoMidjourneyHttpRequest(c *gin.Context, timeout time.Duration, fullRequestU
defer cancel()
resp, err := GetHttpClient().Do(req)
if err != nil {
common.SysLog("do request failed: " + err.Error())
common.SysLog(common.LogText("do request failed: %s", err.Error()))
return MidjourneyErrorWithStatusCodeWrapper(constant.MjErrorUnknown, "do_request_failed", http.StatusInternalServerError), nullBytes, err
}
statusCode := resp.StatusCode
+5 -5
View File
@@ -168,7 +168,7 @@ func resolveOrigins(r *http.Request, settings *system_setting.PasskeySettings) (
continue
}
if !settings.AllowInsecureOrigin && strings.HasPrefix(strings.ToLower(trimmed), "http://") {
return nil, fmt.Errorf("Passkey 不允许使用不安全的 Origin: %s", trimmed)
return nil, fmt.Errorf("passkey does not allow the insecure origin %s", trimmed)
}
origins = append(origins, trimmed)
}
@@ -182,7 +182,7 @@ func resolveOrigins(r *http.Request, settings *system_setting.PasskeySettings) (
autoDetect:
scheme := detectScheme(r)
if scheme == "http" && !settings.AllowInsecureOrigin && r.Host != "localhost" && r.Host != "127.0.0.1" && !strings.HasPrefix(r.Host, "127.0.0.1:") && !strings.HasPrefix(r.Host, "localhost:") {
return nil, fmt.Errorf("Passkey 仅支持 HTTPS,当前访问: %s://%s,请在 Passkey 设置中允许不安全 Origin 或配置 HTTPS", scheme, r.Host)
return nil, fmt.Errorf("passkey requires HTTPS, current request: %s://%s; allow insecure origins in the Passkey settings or configure HTTPS", scheme, r.Host)
}
// 优先使用请求的完整Host(包含端口)
host := r.Host
@@ -197,7 +197,7 @@ autoDetect:
}
}
if host == "" {
return nil, fmt.Errorf("无法确定 Passkey 的 Origin,请在系统设置或 Passkey 设置中指定。当前 Host: '%s', ServerAddress: '%s'", r.Host, system_setting.ServerAddress)
return nil, fmt.Errorf("cannot determine the Passkey origin; set it in the system settings or the Passkey settings. Host: '%s', ServerAddress: '%s'", r.Host, system_setting.ServerAddress)
}
if scheme == "" {
scheme = "https"
@@ -212,11 +212,11 @@ func resolveRPID(r *http.Request, settings *system_setting.PasskeySettings, orig
return hostWithoutPort(rpID), nil
}
if len(origins) == 0 {
return "", errors.New("Passkey 未配置 Origin,无法推导 RPID")
return "", errors.New("passkey origin is not configured, cannot derive the RP ID")
}
parsed, err := url.Parse(origins[0])
if err != nil {
return "", fmt.Errorf("无法解析 Passkey Origin: %w", err)
return "", fmt.Errorf("failed to parse the Passkey origin: %w", err)
}
return hostWithoutPort(parsed.Host), nil
}
+1 -1
View File
@@ -29,7 +29,7 @@ type FlowSecurity struct {
func CreateSessionDataFlow(purpose string, security FlowSecurity, data *webauthn.SessionData) (string, int64, error) {
if data == nil {
return "", 0, errors.New("Passkey 会话数据不能为空")
return "", 0, errors.New("passkey session data is empty")
}
if purpose == model.AuthFlowPurposeLoginPasskey {
if security.UserID <= 0 || security.UserAuthVersion <= 0 || security.LoginFlowID <= 0 || security.LoginExpiresAt <= time.Now().Unix() || security.SessionID != "" || security.SessionVersion != 0 {
+8 -1
View File
@@ -2,6 +2,7 @@ package service
import (
"fmt"
"net/http"
"sync"
"time"
@@ -163,7 +164,13 @@ func IsPoolFailoverError(c *gin.Context, err *types.NewAPIError) bool {
if c.Writer != nil && c.Writer.Written() {
return false
}
if err.StatusCode < 500 || err.StatusCode > 599 {
if err.StatusCode == http.StatusForbidden {
// Only upstream refusals participate in pool failover. Local permission
// failures must not gain the pool's retry budget or override strict affinity.
if err.GetErrorType() == types.ErrorTypeNewAPIError {
return false
}
} else if err.StatusCode < 500 || err.StatusCode > 599 {
return false
}
return !operation_setting.IsAlwaysSkipRetryStatusCode(err.StatusCode) && !operation_setting.IsAlwaysSkipRetryCode(err.GetErrorCode()) && operation_setting.ShouldRetryByStatusCode(err.StatusCode)
+1 -1
View File
@@ -162,7 +162,7 @@ func (t *ssrfProtectedRoundTripper) newTransport(proxyURL *url.URL) *http.Transp
DialContext: dialContext,
}
if common.TLSInsecureSkipVerify {
transport.TLSClientConfig = common.InsecureTLSConfig
transport.TLSClientConfig = common.InsecureTLSConfig.Clone()
}
return transport
}
+92 -94
View File
@@ -7,6 +7,7 @@ import (
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/billingexpr"
@@ -135,12 +136,12 @@ func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usag
if err != nil {
return err
}
logger.LogInfo(ctx, "realtime streaming consume quota success, quota: "+fmt.Sprintf("%d", quota))
logger.LogInfo(ctx, common.LogText("realtime streaming consume quota success, quota: %d", quota))
return nil
}
func PostWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, modelName string,
usage *dto.RealtimeUsage, extraContent string) {
usage *dto.RealtimeUsage, extraContent []*common.Message) {
var tieredResult *billingexpr.TieredResult
tieredOk, tieredQuota, tieredRes := TryTieredSettle(relayInfo, billingexpr.TokenParams{
@@ -191,12 +192,20 @@ func PostWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, mod
}
totalTokens := usage.TotalTokens
var logContent string
var logContent []*common.Message
if !usePrice {
logContent = fmt.Sprintf("模型倍率 %.2f,补全倍率 %.2f,音频倍率 %.2f,音频补全倍率 %.2f,分组倍率 %.2f",
modelRatio, completionRatio.InexactFloat64(), audioRatio.InexactFloat64(), audioCompletionRatio.InexactFloat64(), groupRatio)
logContent = append(logContent, common.NewMessage("Model ratio {{model_ratio}}, completion ratio {{completion_ratio}}, audio ratio {{audio_ratio}}, audio completion ratio {{audio_completion_ratio}}, group ratio {{group_ratio}}", map[string]any{
"model_ratio": fmt.Sprintf("%.2f", modelRatio),
"completion_ratio": fmt.Sprintf("%.2f", completionRatio.InexactFloat64()),
"audio_ratio": fmt.Sprintf("%.2f", audioRatio.InexactFloat64()),
"audio_completion_ratio": fmt.Sprintf("%.2f", audioCompletionRatio.InexactFloat64()),
"group_ratio": fmt.Sprintf("%.2f", groupRatio),
}))
} else {
logContent = fmt.Sprintf("模型价格 %.2f,分组倍率 %.2f", modelPrice, groupRatio)
logContent = append(logContent, common.NewMessage("Model price {{model_price}}, group ratio {{group_ratio}}", map[string]any{
"model_price": fmt.Sprintf("%.2f", modelPrice),
"group_ratio": fmt.Sprintf("%.2f", groupRatio),
}))
}
// record all the consume log even if quota is 0
@@ -204,22 +213,20 @@ func PostWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, mod
// in this case, must be some error happened
// we cannot just return, because we may have to return the pre-consumed quota
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, modelName, relayInfo.FinalPreConsumedQuota))
logContent = append(logContent, common.NewMessage("Possibly an upstream timeout"))
logger.LogError(ctx, common.LogText("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, modelName, relayInfo.FinalPreConsumedQuota))
} else {
model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quota)
model.UpdateChannelUsedQuota(relayInfo.ChannelId, quota)
}
if err := SettleBilling(ctx, relayInfo, quota); err != nil {
logger.LogError(ctx, "error settling billing: "+err.Error())
logger.LogError(ctx, common.LogText("error settling billing: %s", err.Error()))
}
logModel := modelName
if extraContent != "" {
logContent += ", " + extraContent
}
logContent = append(logContent, extraContent...)
other := GenerateWssOtherInfo(ctx, relayInfo, usage, modelRatio, groupRatio,
completionRatio.InexactFloat64(), audioRatio.InexactFloat64(), audioCompletionRatio.InexactFloat64(), modelPrice, relayInfo.PriceData.GroupRatioInfo.GroupSpecialRatio)
if tieredResult != nil {
@@ -268,7 +275,7 @@ func CalcOpenRouterCacheCreateTokens(usage dto.Usage, priceData types.PriceData)
return quota
}
func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage, extraContent string) {
func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage, extraContent []*common.Message) {
if usage == nil {
usage = &dto.Usage{PromptTokens: relayInfo.GetEstimatePromptTokens(), TotalTokens: relayInfo.GetEstimatePromptTokens()}
}
@@ -324,12 +331,20 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, u
}
totalTokens := usage.TotalTokens
var logContent string
var logContent []*common.Message
if !usePrice {
logContent = fmt.Sprintf("模型倍率 %.2f,补全倍率 %.2f,音频倍率 %.2f,音频补全倍率 %.2f,分组倍率 %.2f",
modelRatio, completionRatio.InexactFloat64(), audioRatio.InexactFloat64(), audioCompletionRatio.InexactFloat64(), groupRatio)
logContent = append(logContent, common.NewMessage("Model ratio {{model_ratio}}, completion ratio {{completion_ratio}}, audio ratio {{audio_ratio}}, audio completion ratio {{audio_completion_ratio}}, group ratio {{group_ratio}}", map[string]any{
"model_ratio": fmt.Sprintf("%.2f", modelRatio),
"completion_ratio": fmt.Sprintf("%.2f", completionRatio.InexactFloat64()),
"audio_ratio": fmt.Sprintf("%.2f", audioRatio.InexactFloat64()),
"audio_completion_ratio": fmt.Sprintf("%.2f", audioCompletionRatio.InexactFloat64()),
"group_ratio": fmt.Sprintf("%.2f", groupRatio),
}))
} else {
logContent = fmt.Sprintf("模型价格 %.2f,分组倍率 %.2f", modelPrice, groupRatio)
logContent = append(logContent, common.NewMessage("Model price {{model_price}}, group ratio {{group_ratio}}", map[string]any{
"model_price": fmt.Sprintf("%.2f", modelPrice),
"group_ratio": fmt.Sprintf("%.2f", groupRatio),
}))
}
// record all the consume log even if quota is 0
@@ -337,22 +352,20 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, u
// in this case, must be some error happened
// we cannot just return, because we may have to return the pre-consumed quota
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, billingModelName, relayInfo.FinalPreConsumedQuota))
logContent = append(logContent, common.NewMessage("Possibly an upstream timeout"))
logger.LogError(ctx, common.LogText("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, billingModelName, relayInfo.FinalPreConsumedQuota))
} else {
model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quota)
model.UpdateChannelUsedQuota(relayInfo.ChannelId, quota)
}
if err := SettleBilling(ctx, relayInfo, quota); err != nil {
logger.LogError(ctx, "error settling billing: "+err.Error())
logger.LogError(ctx, common.LogText("error settling billing: %s", err.Error()))
}
logModel := billingModelName
if extraContent != "" {
logContent += ", " + extraContent
}
logContent = append(logContent, extraContent...)
other := GenerateAudioOtherInfo(ctx, relayInfo, usage, modelRatio, groupRatio,
completionRatio.InexactFloat64(), audioRatio.InexactFloat64(), audioCompletionRatio.InexactFloat64(), modelPrice, relayInfo.PriceData.GroupRatioInfo.GroupSpecialRatio)
if tieredResult != nil {
@@ -378,7 +391,7 @@ func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, u
func PreConsumeTokenQuota(relayInfo *relaycommon.RelayInfo, quota int) error {
if quota < 0 {
return errors.New("quota 不能为负数!")
return errors.New("quota cannot be negative")
}
if relayInfo.IsPlayground {
return nil
@@ -459,47 +472,18 @@ func postConsumeQuotaWithResult(relayInfo *relaycommon.RelayInfo, quota int, pre
func checkAndSendQuotaNotify(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQuota int) {
gopool.Go(func() {
userSetting := relayInfo.UserSetting
threshold := common.QuotaRemindThreshold
if userSetting.QuotaWarningThreshold != 0 {
threshold = int(userSetting.QuotaWarningThreshold)
// UserQuota is the balance before this request; quota plus
// preConsumedQuota is what the request finally cost.
balanceAfter := relayInfo.UserQuota - (quota + preConsumedQuota)
if !quotaReminderDue(userSetting, int64(relayInfo.UserQuota), int64(balanceAfter)) {
return
}
//noMoreQuota := userCache.Quota-(quota+preConsumedQuota) <= 0
quotaTooLow := false
consumeQuota := quota + preConsumedQuota
if relayInfo.UserQuota-consumeQuota < threshold {
quotaTooLow = true
}
if quotaTooLow {
prompt := "您的额度即将用尽"
topUpLink := PaymentReturnURL("/wallet")
// 根据通知方式生成不同的内容格式
var content string
var values []any
notifyType := userSetting.NotifyType
if notifyType == "" {
notifyType = dto.NotifyTypeEmail
}
if notifyType == dto.NotifyTypeBark {
// Bark推送使用简短文本,不支持HTML
content = "{{value}},剩余额度:{{value}},请及时充值"
values = []any{prompt, logger.FormatQuota(relayInfo.UserQuota)}
} else if notifyType == dto.NotifyTypeGotify {
content = "{{value}},当前剩余额度为 {{value}},请及时充值。"
values = []any{prompt, logger.FormatQuota(relayInfo.UserQuota)}
} else {
// 默认内容格式,适用于Email和Webhook(支持HTML)
content = "{{value}},当前剩余额度为 {{value}},为了不影响您的使用,请及时充值。<br/>充值链接:<a href='{{value}}'>{{value}}</a>"
values = []any{prompt, logger.FormatQuota(relayInfo.UserQuota), topUpLink, topUpLink}
}
err := NotifyUser(relayInfo.UserId, relayInfo.UserEmail, relayInfo.UserSetting, dto.NewNotify(dto.NotifyTypeQuotaExceed, prompt, content, values))
if err != nil {
common.SysError(fmt.Sprintf("failed to send quota notify to user %d: %s", relayInfo.UserId, err.Error()))
}
lang := userSetting.Language
prompt := i18n.Translate(lang, i18n.MsgNotifyQuotaLowTitle)
content, values := quotaNotifyContent(lang, userSetting.NotifyType, prompt, logger.FormatQuota(balanceAfter))
err := NotifyUser(relayInfo.UserId, relayInfo.UserEmail, relayInfo.UserSetting, dto.NewNotify(dto.NotifyTypeQuotaExceed, prompt, content, values))
if err != nil {
common.SysError(common.LogText("failed to send quota notify to user %d: %s", relayInfo.UserId, err.Error()))
}
})
}
@@ -514,40 +498,54 @@ func checkAndSendSubscriptionQuotaNotify(relayInfo *relaycommon.RelayInfo) {
}
userSetting := relayInfo.UserSetting
threshold := common.QuotaRemindThreshold
if userSetting.QuotaWarningThreshold != 0 {
threshold = int(userSetting.QuotaWarningThreshold)
}
usedBefore := relayInfo.SubscriptionAmountUsedAfterPreConsume - relayInfo.SubscriptionPreConsumed
usedAfter := relayInfo.SubscriptionAmountUsedAfterPreConsume + relayInfo.SubscriptionPostDelta
remaining := relayInfo.SubscriptionAmountTotal - usedAfter
if remaining >= int64(threshold) {
if !quotaReminderDue(userSetting, relayInfo.SubscriptionAmountTotal-usedBefore, remaining) {
return
}
prompt := "您的订阅额度即将用尽"
topUpLink := PaymentReturnURL("/wallet")
var content string
var values []any
notifyType := userSetting.NotifyType
if notifyType == "" {
notifyType = dto.NotifyTypeEmail
}
if notifyType == dto.NotifyTypeBark {
content = "{{value}},剩余额度:{{value}},请及时充值"
values = []any{prompt, logger.FormatQuota(int(remaining))}
} else if notifyType == dto.NotifyTypeGotify {
content = "{{value}},当前剩余额度为 {{value}},请及时充值。"
values = []any{prompt, logger.FormatQuota(int(remaining))}
} else {
content = "{{value}},当前剩余额度为 {{value}},为了不影响您的使用,请及时充值。<br/>充值链接:<a href='{{value}}'>{{value}}</a>"
values = []any{prompt, logger.FormatQuota(int(remaining)), topUpLink, topUpLink}
}
lang := userSetting.Language
prompt := i18n.Translate(lang, i18n.MsgNotifySubscriptionQuotaLowTitle)
content, values := quotaNotifyContent(lang, userSetting.NotifyType, prompt, logger.FormatQuota(int(remaining)))
if err := NotifyUser(relayInfo.UserId, relayInfo.UserEmail, relayInfo.UserSetting, dto.NewNotify(dto.NotifyTypeQuotaExceed, prompt, content, values)); err != nil {
common.SysError(fmt.Sprintf("failed to send subscription quota notify to user %d: %s", relayInfo.UserId, err.Error()))
common.SysError(common.LogText("failed to send subscription quota notify to user %d: %s", relayInfo.UserId, err.Error()))
}
})
}
// quotaReminderDue reports whether a charge that moved a balance (wallet or
// subscription) from before to after earns a low-balance reminder. The site
// and the user must both have reminders on, and the charge must cross the
// user's threshold, or the site default when the user set none. A balance
// already below the threshold earns nothing more until a top-up lifts it back
// above, so a user gets one reminder per drop.
func quotaReminderDue(userSetting dto.UserSetting, before int64, after int64) bool {
if !common.QuotaRemindEnabled {
return false
}
if userSetting.QuotaWarningEnabled != nil && !*userSetting.QuotaWarningEnabled {
return false
}
threshold := float64(common.QuotaRemindThreshold)
if userSetting.QuotaWarningThreshold != 0 {
threshold = userSetting.QuotaWarningThreshold
}
return float64(before) >= threshold && float64(after) < threshold
}
// quotaNotifyContent renders a low-quota notice in the user's saved language.
// Bark and Gotify get plain text; email and webhook get HTML with the top-up
// link. values stay in the notify payload for webhook receivers.
func quotaNotifyContent(lang string, notifyType string, prompt string, remaining string) (string, []any) {
switch notifyType {
case dto.NotifyTypeBark:
// Bark推送使用简短文本,不支持HTML
return i18n.Translate(lang, i18n.MsgNotifyQuotaLowBark, map[string]any{"Prompt": prompt, "Quota": remaining}), []any{prompt, remaining}
case dto.NotifyTypeGotify:
return i18n.Translate(lang, i18n.MsgNotifyQuotaLowGotify, map[string]any{"Prompt": prompt, "Quota": remaining}), []any{prompt, remaining}
default:
topUpLink := PaymentReturnURL("/wallet")
return i18n.Translate(lang, i18n.MsgNotifyQuotaLowEmail, map[string]any{"Prompt": prompt, "Quota": remaining, "Link": topUpLink}), []any{prompt, remaining, topUpLink, topUpLink}
}
}
+5 -4
View File
@@ -9,6 +9,7 @@ import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/tokenkit"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
@@ -70,7 +71,7 @@ func resolveAttemptGuardResult(reason string, info *relaycommon.RelayInfo) strin
return reason
}
if info != nil && info.ResponseModel != nil && info.ResponseModel.ReturnedModel != "" {
if info.ResponseModel.Mismatch {
if info.ResponseModel.Mismatch() {
return "mismatch_allowed"
}
return "accepted"
@@ -120,7 +121,7 @@ func recordDurableAttempt(c *gin.Context, info *relaycommon.RelayInfo, usage *dt
if reason != "accepted" && reason != "unchecked" {
params.Other.SetPublic("model_guard_result", reason)
}
log := &model.Log{UserId: info.UserId, Username: common.GetContextKeyString(c, constant.ContextKeyUserName), CreatedAt: time.Now().Unix(), Type: model.LogTypeConsume, Content: params.Content, TokenName: params.TokenName, ModelName: params.ModelName, Quota: params.Quota, PromptTokens: params.PromptTokens, CompletionTokens: params.CompletionTokens, UseTime: params.UseTimeSeconds, IsStream: params.IsStream, ChannelId: params.ChannelId, TokenId: params.TokenId, Group: params.Group, RequestId: info.RequestId, Other: params.Other.JSONString(), CorrelationKey: c.GetString("log_correlation_key")}
log := &model.Log{UserId: info.UserId, Username: common.GetContextKeyString(c, constant.ContextKeyUserName), CreatedAt: time.Now().Unix(), Type: model.LogTypeConsume, Content: params.Other.SetContent(params.Content), TokenName: params.TokenName, ModelName: params.ModelName, Quota: params.Quota, PromptTokens: params.PromptTokens, CompletionTokens: params.CompletionTokens, UseTime: params.UseTimeSeconds, IsStream: params.IsStream, ChannelId: params.ChannelId, TokenId: params.TokenId, Group: params.Group, RequestId: info.RequestId, Other: params.Other.JSONString(), CorrelationKey: c.GetString("log_correlation_key")}
billable := params.Quota > 0 || params.PromptTokens > 0 || params.CompletionTokens > 0
usageBody, _ := common.Marshal(usage)
logBody, _ := common.Marshal(log)
@@ -216,14 +217,14 @@ func (collector *GuardUsageCollector) Observe(payload []byte, stream bool) {
if collector.outputTokens == 0 {
text := ExtractOutputTextFromResponses(event.Response)
if text != "" {
collector.outputTokens = CountTextToken(text, collector.info.OriginModelName)
collector.outputTokens = tokenkit.Count(collector.info.OriginModelName, text)
collector.observed = true
}
}
}
if strings.HasSuffix(event.Type, ".delta") && event.Delta != "" {
collector.observed = true
collector.outputTokens += CountTextToken(event.Delta, collector.info.OriginModelName)
collector.outputTokens += tokenkit.Count(collector.info.OriginModelName, event.Delta)
}
if collector.hasUsage || collector.observed || collector.ctx.GetString("attempt_returned_model") != "" && collector.ctx.GetString("attempt_returned_model") != collector.info.OriginModelName {
if time.Since(collector.lastCheckpoint) > time.Second || collector.hasUsage {
-1
View File
@@ -24,7 +24,6 @@ func TestResolveAttemptGuardResultForDetectionOnly(t *testing.T) {
info := &relaycommon.RelayInfo{ResponseModel: &relaycommon.ResponseModel{
RequestedModel: "gpt-6-astra",
ReturnedModel: "gpt-5.6-luna",
Mismatch: true,
}}
require.Equal(t, "mismatch_allowed", resolveAttemptGuardResult("", info))
})
+1 -1
View File
@@ -71,7 +71,7 @@ func ProcessChannelError(c *gin.Context, channelError types.ChannelError, err *t
if err == nil {
return
}
logger.LogError(c, fmt.Sprintf("channel error (channel #%d, status code: %d): %s", channelError.ChannelId, err.StatusCode, common.LocalLogPreview(err.MaskSensitiveErrorWithStatusCode())))
logger.LogError(c, common.LogText("channel error (channel #%d, status code: %d): %s", channelError.ChannelId, err.StatusCode, common.LocalLogPreview(err.MaskSensitiveErrorWithStatusCode())))
if err.GetErrorCode() != relaycommon.ResponseGuardErrorCode && ShouldDisableChannel(err) && channelError.AutoBan {
reason := err.MaskSensitiveErrorWithStatusCode()
gopool.Go(func() {
+70 -3
View File
@@ -156,6 +156,7 @@ func TestProcessChannelErrorMasksDisableReasonAndNotification(t *testing.T) {
notifyLimitStore.Delete(notifyKey)
t.Cleanup(func() { notifyLimitStore.Delete(notifyKey) })
c, _ := gin.CreateTestContext(httptest.NewRecorder())
common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, "https://private.example.com")
apiErr := types.NewErrorWithStatusCode(errors.New("upstream https://private.example.com/path?token=review-token api_key:review-secret"), types.ErrorCodeChannelNoAvailableKey, http.StatusUnauthorized)
ProcessChannelError(c, types.ChannelError{ChannelId: channel.Id, ChannelName: channel.Name, AutoBan: true}, apiErr, nil)
var notification WebhookPayload
@@ -176,6 +177,72 @@ func TestProcessChannelErrorMasksDisableReasonAndNotification(t *testing.T) {
assert.Equal(t, http.StatusUnauthorized, apiErr.StatusCode)
}
// Only addresses the admin configured are masked: a channel on the base URL
// new-api presets reaches a public official endpoint, and so do adaptors that
// build official endpoints themselves.
func TestUpstreamAddressesMaskOnlyConfiguredAddresses(t *testing.T) {
for _, tc := range []struct {
name string
channelType int
baseURL string
origin string
proxyEnv string
message string
want string
}{
{
name: "preset base URL stays readable",
channelType: constant.ChannelTypeAnthropic,
baseURL: "https://api.anthropic.com",
origin: "https://api.anthropic.com",
message: "request to https://api.anthropic.com/v1/messages failed, see https://platform.claude.com/docs/en/build-with-claude/task-budgets",
want: "request to https://api.anthropic.com/v1/messages failed, see https://platform.claude.com/docs/en/build-with-claude/task-budgets",
},
{
name: "official endpoint built by the adaptor stays readable",
channelType: constant.ChannelTypeVertexAi,
origin: "https://us-central1-aiplatform.googleapis.com",
message: "request to https://us-central1-aiplatform.googleapis.com/v1/projects failed",
want: "request to https://us-central1-aiplatform.googleapis.com/v1/projects failed",
},
{
name: "configured base URL is masked",
channelType: constant.ChannelTypeOpenAI,
baseURL: "https://relay.private.example",
origin: "https://relay.private.example",
message: "request to https://relay.private.example/v1/chat/completions failed",
want: "request to https://***.example/***/***/*** failed",
},
{
name: "absolute Advanced Custom route is masked",
channelType: constant.ChannelTypeAdvancedCustom,
origin: "https://route.private.example",
message: "dial tcp: lookup route.private.example: no such host",
want: "dial tcp: lookup ***.***.example: no such host",
},
{
name: "environment proxy is masked on a preset channel",
channelType: constant.ChannelTypeAnthropic,
baseURL: "https://api.anthropic.com",
proxyEnv: "http://egress.private.invalid:8080",
message: "proxyconnect tcp: dial tcp: lookup egress.private.invalid: no such host",
want: "proxyconnect tcp: dial tcp: lookup ***.***.invalid: no such host",
},
} {
t.Run(tc.name, func(t *testing.T) {
t.Setenv("HTTPS_PROXY", tc.proxyEnv)
t.Setenv("HTTP_PROXY", "")
c, _ := gin.CreateTestContext(httptest.NewRecorder())
common.SetContextKey(c, constant.ContextKeyChannelType, tc.channelType)
common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, tc.baseURL)
common.SetContextKey(c, constant.ContextKeyUpstreamRequestOrigin, tc.origin)
apiErr := types.NewErrorWithStatusCode(errors.New(tc.message), types.ErrorCodeBadResponseStatusCode, http.StatusBadGateway)
assert.Equal(t, tc.want, apiErr.ToOpenAIError(common.UpstreamAddresses(c)...).Message)
})
}
}
func TestDecideRelayRetryReasons(t *testing.T) {
upstream := func(status int) *types.NewAPIError {
return types.NewOpenAIError(errors.New("upstream"), types.ErrorCodeBadResponseStatusCode, status)
@@ -269,9 +336,9 @@ func TestRelayErrorGroup(t *testing.T) {
channel int
actual, auto, selected, base, want string
}{
{"actual group", 102, "example-preferred", "stale", "stale", "default", "example-preferred"},
{"auto group", 102, "", "example-preferred", "", "default", "example-preferred"},
{"pool selection", 102, "", "", "example-preferred", "default", "example-preferred"},
{"actual group", 27, "test-group-b", "stale", "stale", "default", "test-group-b"},
{"auto group", 27, "", "test-group-b", "", "default", "test-group-b"},
{"pool selection", 27, "", "", "test-group-b", "default", "test-group-b"},
{"fixed group", 27, "", "", "", "fixed", "fixed"},
{"no selection", 0, "stale", "stale", "stale", "default", ""},
{"unresolved auto", 27, "", "", "", "auto", ""},
+13 -8
View File
@@ -8,6 +8,7 @@ import (
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/relayconvert"
"github.com/QuantumNous/new-api/tokenkit"
)
// ResponsesUsageAccumulator owns the accounting facts for one Responses stream.
@@ -29,7 +30,9 @@ func NewResponsesUsageAccumulator(info *relaycommon.RelayInfo) *ResponsesUsageAc
return &ResponsesUsageAccumulator{info: info, usage: &dto.Usage{}}
}
func (a *ResponsesUsageAccumulator) Observe(event *dto.ResponsesStreamResponse) {
// Observe feeds one decoded stream event. raw is the same event's wire bytes;
// the vendor tool-usage reader runs on it once, at the terminal event.
func (a *ResponsesUsageAccumulator) Observe(event *dto.ResponsesStreamResponse, raw []byte) {
if a == nil || event == nil || a.finished {
return
}
@@ -50,13 +53,16 @@ func (a *ResponsesUsageAccumulator) Observe(event *dto.ResponsesStreamResponse)
a.outputText.WriteString(relayconvert.ExtractOutputTextFromResponses(event.Response))
}
}
// Vendor counts are cumulative on the terminal event and replace the
// web_search_call items counted from output_item.done.
a.info.ApplyVendorToolUsage(raw)
if a.imageCommitted {
return
}
failed := event.Type != "response.completed" && event.Type != "response.done"
if failed || (event.Response != nil && relaycommon.IsNonBillableResponsesStatus(event.Response.Status)) {
a.imageCounter.Reset()
} else if event.Response != nil {
// Images that completed before any terminal, failed ones included,
// were delivered and stay billable; Observe still skips unfinished
// image items.
if event.Response != nil {
for i := range event.Response.Output {
a.imageCounter.Observe(&event.Response.Output[i], &i)
}
@@ -92,15 +98,14 @@ func (a *ResponsesUsageAccumulator) Finish() *dto.Usage {
a.usage.UsageSource = "estimated"
}
// A final image item can already have reached the client before the stream
// disconnects. Explicit failed/incomplete terminals reset and commit zero in
// Observe; otherwise retain completed tool usage even without a terminal.
// disconnects, so completed tool usage is retained even without a terminal.
if !a.imageCommitted {
a.imageCounter.Commit(a.info)
a.imageCommitted = true
}
if a.usage.CompletionTokens == 0 {
if output := a.outputText.String(); output != "" {
a.usage.CompletionTokens = CountTextToken(output, a.info.GetUpstreamModelName())
a.usage.CompletionTokens = tokenkit.Count(a.info.GetUpstreamModelName(), output)
}
}
// Upstream bills the prompt as soon as it starts generating, so a stream
+17 -16
View File
@@ -7,6 +7,7 @@ import (
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/setting/operation_setting"
"github.com/QuantumNous/new-api/tokenkit"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
@@ -20,10 +21,10 @@ func TestResponsesUsageAccumulatorTerminalAccounting(t *testing.T) {
}{
{eventType: "response.completed", wantImages: 1},
{eventType: "response.done", wantImages: 1},
{eventType: "response.incomplete"},
{eventType: "response.failed"},
{eventType: "response.cancelled"},
{eventType: "response.canceled"},
{eventType: "response.incomplete", wantImages: 1},
{eventType: "response.failed", wantImages: 1},
{eventType: "response.cancelled", wantImages: 1},
{eventType: "response.canceled", wantImages: 1},
} {
t.Run(tc.eventType, func(t *testing.T) {
info := &relaycommon.RelayInfo{OriginModelName: "gpt-5.1", StreamStatus: relaycommon.NewStreamStatus()}
@@ -34,10 +35,10 @@ func TestResponsesUsageAccumulatorTerminalAccounting(t *testing.T) {
{Type: dto.BuildInCallFunctionCall, Name: "responses_priced_fn"},
{Type: dto.BuildInCallFunctionCall, Name: "responses_unpriced_fn"},
} {
accumulator.Observe(&dto.ResponsesStreamResponse{Type: dto.ResponsesOutputTypeItemDone, Item: &item})
accumulator.Observe(&dto.ResponsesStreamResponse{Type: dto.ResponsesOutputTypeItemDone, Item: &item}, nil)
}
image := dto.ResponsesOutput{ID: "image-1", Type: dto.ResponsesOutputTypeImageGenerationCall, Status: "completed", Result: "image-data"}
accumulator.Observe(&dto.ResponsesStreamResponse{Type: dto.ResponsesOutputTypeItemDone, Item: &image})
accumulator.Observe(&dto.ResponsesStreamResponse{Type: dto.ResponsesOutputTypeItemDone, Item: &image}, nil)
upstream := &dto.Usage{InputTokens: 20, OutputTokens: 5, TotalTokens: 25, InputTokensDetails: &dto.InputTokenDetails{CachedTokens: 4}}
upstream.BillingUsage = dto.NewOpenAIResponsesBillingUsage(upstream)
terminal := &dto.ResponsesStreamResponse{
@@ -46,8 +47,8 @@ func TestResponsesUsageAccumulatorTerminalAccounting(t *testing.T) {
Usage: upstream, Output: []dto.ResponsesOutput{image},
},
}
accumulator.Observe(terminal)
accumulator.Observe(terminal)
accumulator.Observe(terminal, nil)
accumulator.Observe(terminal, nil)
usage := accumulator.Finish()
assert.Equal(t, tc.eventType == "response.failed", info.StreamStatus.ResponseFailed())
assert.NotEmpty(t, info.StreamStatus.ResponseOutcome())
@@ -81,11 +82,11 @@ func TestResponsesUsageAccumulatorInterruptedTextFallback(t *testing.T) {
info := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-4o"}}
info.SetEstimatePromptTokens(100)
accumulator := NewResponsesUsageAccumulator(info)
accumulator.Observe(&dto.ResponsesStreamResponse{Type: "response.output_text.delta", Delta: "hello"})
accumulator.Observe(&dto.ResponsesStreamResponse{Type: "response.output_text.delta", Delta: "hello"}, nil)
if withUsage {
upstream := &dto.Usage{InputTokens: 20, InputTokensDetails: &dto.InputTokenDetails{CachedTokens: 4}}
upstream.BillingUsage = dto.NewOpenAIResponsesBillingUsage(upstream)
accumulator.Observe(&dto.ResponsesStreamResponse{Type: "response.failed", Response: &dto.OpenAIResponsesResponse{Usage: upstream}})
accumulator.Observe(&dto.ResponsesStreamResponse{Type: "response.failed", Response: &dto.OpenAIResponsesResponse{Usage: upstream}}, nil)
}
usage := accumulator.Finish()
assert.Equal(t, 1, usage.CompletionTokens)
@@ -103,7 +104,7 @@ func TestResponsesUsageAccumulatorInterruptedTextFallback(t *testing.T) {
assert.Equal(t, 100, usage.PromptTokens)
assert.Equal(t, 101, usage.TotalTokens)
}
accumulator.Observe(&dto.ResponsesStreamResponse{Type: "response.output_text.delta", Delta: " late output"})
accumulator.Observe(&dto.ResponsesStreamResponse{Type: "response.output_text.delta", Delta: " late output"}, nil)
assert.Equal(t, usage, accumulator.Finish())
assert.Equal(t, 1, usage.CompletionTokens)
})
@@ -144,11 +145,11 @@ func TestResponsesUsageAccumulatorDisconnectBillsCompletedImage(t *testing.T) {
accumulator.Observe(&dto.ResponsesStreamResponse{
Type: dto.ResponsesOutputTypeItemDone,
Item: &dto.ResponsesOutput{ID: "complete-image", Type: dto.ResponsesOutputTypeImageGenerationCall, Status: "completed", Result: "final-image-data"},
})
}, nil)
accumulator.Observe(&dto.ResponsesStreamResponse{
Type: dto.ResponsesOutputTypeItemDone,
Item: &dto.ResponsesOutput{ID: "partial-image", Type: dto.ResponsesOutputTypeImageGenerationCall, Status: "partial", Result: "partial-image-data"},
})
}, nil)
usage := accumulator.Finish()
assert.Zero(t, usage.TotalTokens)
require.Contains(t, info.ResponsesUsageInfo.BuiltInTools, dto.BuildInToolImageGeneration)
@@ -272,7 +273,7 @@ func TestResponsesUsageAccumulatorMissingUsageEstimation(t *testing.T) {
{Type: dto.ResponsesOutputTypeItemDone, Item: &dto.ResponsesOutput{Type: dto.BuildInCallFunctionCall, Name: "shell"}},
},
wantPrompt: 100,
wantCompletion: CountTextToken(summary+arguments, model),
wantCompletion: tokenkit.Count(model, summary+arguments),
},
{
name: "created only then disconnect bills the prompt",
@@ -302,7 +303,7 @@ func TestResponsesUsageAccumulatorMissingUsageEstimation(t *testing.T) {
}},
},
wantPrompt: 100,
wantCompletion: CountTextToken("final answer", model),
wantCompletion: tokenkit.Count(model, "final answer"),
},
{
name: "explicit failure without usage bills nothing",
@@ -330,7 +331,7 @@ func TestResponsesUsageAccumulatorMissingUsageEstimation(t *testing.T) {
info.SetEstimatePromptTokens(100)
accumulator := NewResponsesUsageAccumulator(info)
for i := range tc.events {
accumulator.Observe(&tc.events[i])
accumulator.Observe(&tc.events[i], nil)
}
usage := accumulator.Finish()
assert.Equal(t, tc.wantPrompt, usage.PromptTokens)
+166 -7
View File
@@ -4,6 +4,7 @@ import (
"crypto/hmac"
"encoding/json"
"errors"
"slices"
"sort"
"strings"
@@ -29,13 +30,28 @@ const (
VerificationScopeLogin = "auth.login"
VerificationScopeAccessTokenGenerate = "access_token.generate"
VerificationScopeAccessTokenRevoke = "access_token.revoke"
VerificationScopeAccessTokenUpdate = "access_token.update"
VerificationScopeAccountBind = "account.binding.bind"
VerificationScopeAccountUnbind = "account.binding.unbind"
VerificationScopePasswordSet = "account.password.set"
VerificationScopePasswordChange = "account.password.change"
VerificationScopeAccountDelete = "account.delete"
// Administrative step-up scopes bind the managed user into the proof so a
// proof issued for one account cannot be replayed against another.
VerificationScopeAdminUserCreate = "admin.user.create"
VerificationScopeAdminUserUpdate = "admin.user.update"
VerificationScopeAdminUserDelete = "admin.user.delete"
VerificationScopeAdminUserManage = "admin.user.manage"
VerificationScopeAdminUserPasskeyReset = "admin.user.passkey.reset"
VerificationScopeAdminUserTwoFADisable = "admin.user.2fa.disable"
VerificationScopeAdminUserBindingClear = "admin.user.binding.clear"
verificationScopeAdminUserPrefix = "admin.user."
)
// adminUserManageActions are the ManageUser actions that change a user's
// status or role. Quota adjustments are not gated and deletion has its own scope.
var adminUserManageActions = []string{"disable", "enable", "promote", "demote"}
var (
ErrVerificationFailed = errors.New("Verification failed. Please try again.")
ErrVerificationUnavailable = errors.New("This verification method is currently unavailable.")
@@ -66,6 +82,71 @@ type AccountUnbindingContext struct {
ProviderID int `json:"provider_id"`
}
type AdminUserContext struct {
UserID int `json:"user_id"`
}
type AdminUserManageContext struct {
UserID int `json:"user_id"`
Action string `json:"action"`
}
// AdminUserBindingContext names exactly one binding: a built-in binding type
// or a custom OAuth provider ID.
type AdminUserBindingContext struct {
UserID int `json:"user_id"`
BindingType string `json:"binding_type,omitempty"`
ProviderID int `json:"provider_id,omitempty"`
}
type AdminUserCreateContext struct {
Role int `json:"role"`
}
// AccessTokenGenerateContext binds a creation proof to the exact grant. Scopes
// are compared as a sorted, de-duplicated set.
type AccessTokenGenerateContext struct {
Scopes []string `json:"scopes"`
ExpiresAt int64 `json:"expires_at"`
}
// AccessTokenUpdateContext binds a grant change to one token and the exact new
// grant, compared as a sorted, de-duplicated set.
type AccessTokenUpdateContext struct {
TokenID int `json:"token_id"`
Scopes []string `json:"scopes"`
}
// AccessTokenRevokeContext names exactly one token: a scoped token ID or the
// legacy token.
type AccessTokenRevokeContext struct {
TokenID int `json:"token_id,omitempty"`
Legacy bool `json:"legacy,omitempty"`
}
const (
maxAccessTokenContextScopes = 128
maxAccessTokenContextScopeLen = 64
)
// NormalizeAccessTokenScopeList trims, de-duplicates and sorts scope keys. It
// only checks the shape; grantability is checked against the scope catalog.
func NormalizeAccessTokenScopeList(scopes []string) ([]string, bool) {
if len(scopes) == 0 || len(scopes) > maxAccessTokenContextScopes {
return nil, false
}
normalized := make([]string, 0, len(scopes))
for _, scope := range scopes {
scope = strings.TrimSpace(scope)
if scope == "" || len(scope) > maxAccessTokenContextScopeLen {
return nil, false
}
normalized = append(normalized, scope)
}
slices.Sort(normalized)
return slices.Compact(normalized), true
}
// VerificationBinding contains no original operation parameters. It can safely
// travel through a signed proof or a server-owned interactive verification flow.
type VerificationBinding struct {
@@ -117,9 +198,69 @@ func BindVerificationOperation(operation VerificationOperation) (VerificationBin
return VerificationBinding{}, ErrVerificationContextInvalid
}
normalized = context
case VerificationScopeAdminUserUpdate, VerificationScopeAdminUserDelete, VerificationScopeAdminUserPasskeyReset, VerificationScopeAdminUserTwoFADisable:
var context AdminUserContext
if len(fields) != 1 || common.Unmarshal(fields["user_id"], &context.UserID) != nil || context.UserID <= 0 {
return VerificationBinding{}, ErrVerificationContextInvalid
}
normalized = context
case VerificationScopeAdminUserManage:
var context AdminUserManageContext
if len(fields) != 2 || common.Unmarshal(operation.Context, &context) != nil || context.UserID <= 0 || !slices.Contains(adminUserManageActions, context.Action) {
return VerificationBinding{}, ErrVerificationContextInvalid
}
normalized = context
case VerificationScopeAdminUserBindingClear:
var context AdminUserBindingContext
if len(fields) != 2 || common.Unmarshal(operation.Context, &context) != nil || context.UserID <= 0 || context.ProviderID < 0 {
return VerificationBinding{}, ErrVerificationContextInvalid
}
context.BindingType = strings.ToLower(strings.TrimSpace(context.BindingType))
if (context.BindingType == "") == (context.ProviderID == 0) || len(context.BindingType) > 32 {
return VerificationBinding{}, ErrVerificationContextInvalid
}
normalized = context
case VerificationScopeAdminUserCreate:
var context AdminUserCreateContext
if len(fields) != 1 || common.Unmarshal(fields["role"], &context.Role) != nil || context.Role < common.RoleAdminUser || !common.IsValidateRole(context.Role) {
return VerificationBinding{}, ErrVerificationContextInvalid
}
normalized = context
case VerificationScopeAccessTokenGenerate:
var context AccessTokenGenerateContext
if len(fields) != 2 || fields["scopes"] == nil || fields["expires_at"] == nil ||
common.Unmarshal(operation.Context, &context) != nil || context.ExpiresAt < 0 {
return VerificationBinding{}, ErrVerificationContextInvalid
}
scopes, ok := NormalizeAccessTokenScopeList(context.Scopes)
if !ok {
return VerificationBinding{}, ErrVerificationContextInvalid
}
context.Scopes = scopes
normalized = context
case VerificationScopeAccessTokenUpdate:
var context AccessTokenUpdateContext
if len(fields) != 2 || common.Unmarshal(operation.Context, &context) != nil || context.TokenID <= 0 {
return VerificationBinding{}, ErrVerificationContextInvalid
}
scopes, ok := NormalizeAccessTokenScopeList(context.Scopes)
if !ok {
return VerificationBinding{}, ErrVerificationContextInvalid
}
context.Scopes = scopes
normalized = context
case VerificationScopeAccessTokenRevoke:
var context AccessTokenRevokeContext
if len(fields) != 1 || common.Unmarshal(operation.Context, &context) != nil || context.TokenID < 0 {
return VerificationBinding{}, ErrVerificationContextInvalid
}
// Exactly one target: {"token_id": n} or {"legacy": true}.
if (context.TokenID > 0) == context.Legacy {
return VerificationBinding{}, ErrVerificationContextInvalid
}
normalized = context
case VerificationScopePasskeyRegister, VerificationScopePasskeyDelete, VerificationScopeTwoFASetup,
VerificationScopeTwoFADisable, VerificationScopeTwoFABackupCodes,
VerificationScopeAccessTokenGenerate, VerificationScopeAccessTokenRevoke,
VerificationScopePasswordSet, VerificationScopePasswordChange, VerificationScopeAccountDelete:
if len(fields) != 0 {
return VerificationBinding{}, ErrVerificationContextInvalid
@@ -184,9 +325,12 @@ func securityVerificationPolicy(scope string, state model.UserVerificationState)
return nil, model.ErrTwoFANotEnabled
}
case VerificationScopePasskeyRegister, VerificationScopeTwoFASetup,
VerificationScopeAccessTokenGenerate, VerificationScopeAccessTokenRevoke,
VerificationScopeAccessTokenGenerate, VerificationScopeAccessTokenUpdate, VerificationScopeAccessTokenRevoke,
VerificationScopeAccountBind, VerificationScopeAccountUnbind,
VerificationScopePasswordSet, VerificationScopePasswordChange, VerificationScopeAccountDelete:
VerificationScopePasswordSet, VerificationScopePasswordChange, VerificationScopeAccountDelete,
VerificationScopeAdminUserCreate, VerificationScopeAdminUserUpdate, VerificationScopeAdminUserDelete,
VerificationScopeAdminUserManage, VerificationScopeAdminUserPasskeyReset,
VerificationScopeAdminUserTwoFADisable, VerificationScopeAdminUserBindingClear:
if scope == VerificationScopeAccountDelete && state.Role == common.RoleRootUser {
return nil, ErrVerificationForbidden
}
@@ -234,6 +378,12 @@ func GetVerificationRequirements(identity AuthIdentity, scope string) (*Verifica
if scope == VerificationScopeChannelKeyRead && state.Role != common.RoleRootUser {
return nil, ErrVerificationForbidden
}
if strings.HasPrefix(scope, verificationScopeAdminUserPrefix) && state.Role < common.RoleAdminUser {
return nil, ErrVerificationForbidden
}
if err := requireAccessTokenVerificationScope(identity, scope); err != nil {
return nil, err
}
methods, err := securityVerificationPolicy(scope, *state)
if err != nil {
return nil, err
@@ -242,13 +392,22 @@ func GetVerificationRequirements(identity AuthIdentity, scope string) (*Verifica
for i := range methods {
if methods[i].Method == VerificationMethodPassword && !common.PasswordLoginEnabled {
switch scope {
case VerificationScopeAccountBind, VerificationScopeAccountUnbind, VerificationScopePasswordSet, VerificationScopePasswordChange, VerificationScopeAccountDelete:
case VerificationScopeAccountBind, VerificationScopeAccountUnbind, VerificationScopePasswordSet, VerificationScopePasswordChange, VerificationScopeAccountDelete,
VerificationScopeAdminUserCreate, VerificationScopeAdminUserUpdate, VerificationScopeAdminUserDelete,
VerificationScopeAdminUserManage, VerificationScopeAdminUserPasskeyReset,
VerificationScopeAdminUserTwoFADisable, VerificationScopeAdminUserBindingClear:
methods[i].Available, methods[i].Reason = false, "Password authentication is disabled."
}
}
if methods[i].Method != VerificationMethodOAuth {
continue
}
// OAuth verification completes through a browser popup bound to a login
// session, which an access token cannot provide.
if _, ok := model.ParseAccessTokenSessionID(identity.SessionID); ok {
methods[i].Available, methods[i].Reason = false, "OAuth verification requires a browser sign-in. Enable two-factor authentication or a Passkey to continue."
continue
}
user, err := model.GetUserById(identity.UserID, false)
if err != nil {
return nil, err
@@ -330,7 +489,7 @@ func RequireVerificationMethod(identity AuthIdentity, scope, method string) (*Ve
// CompleteSecurityVerification runs after the concrete authentication ceremony.
// Recheck the session and policy after potentially slow external authentication.
func CompleteSecurityVerification(identity AuthIdentity, binding VerificationBinding, method string) (*SecurityProof, error) {
if _, _, err := ValidateLoginSession(identity); err != nil {
if err := ValidateStepUpIdentity(identity); err != nil {
return nil, err
}
if _, err := RequireVerificationMethod(identity, binding.Scope, method); err != nil {
@@ -354,7 +513,7 @@ func ConsumeOperationProof(raw string, identity AuthIdentity, operation Verifica
if err != nil {
return nil, err
}
if _, _, err := ValidateLoginSession(identity); err != nil {
if err := ValidateStepUpIdentity(identity); err != nil {
return nil, err
}
if _, err := RequireVerificationMethod(identity, binding.Scope, claims.Method); err != nil {
@@ -394,7 +553,7 @@ func ValidateFlowAuthorization(identity AuthIdentity, operation VerificationOper
if authorization.Scope != binding.Scope || !hmac.Equal([]byte(authorization.ContextHash), []byte(binding.ContextHash)) {
return model.ErrAuthFlowInvalid
}
if _, _, err := ValidateLoginSession(identity); err != nil {
if err := ValidateStepUpIdentity(identity); err != nil {
return err
}
_, err = RequireVerificationMethod(identity, binding.Scope, authorization.Method)
+3 -4
View File
@@ -2,7 +2,6 @@ package service
import (
"context"
"fmt"
"sync"
"sync/atomic"
"time"
@@ -32,7 +31,7 @@ func StartSubscriptionQuotaResetTask() {
return
}
gopool.Go(func() {
logger.LogInfo(context.Background(), fmt.Sprintf("subscription quota reset task started: tick=%s", subscriptionResetTickInterval))
logger.LogInfo(context.Background(), common.LogText("subscription quota reset task started: tick=%s", subscriptionResetTickInterval))
ticker := time.NewTicker(subscriptionResetTickInterval)
defer ticker.Stop()
@@ -56,7 +55,7 @@ func runSubscriptionQuotaResetOnce() {
for {
n, err := model.ExpireDueSubscriptions(subscriptionResetBatchSize)
if err != nil {
logger.LogWarn(ctx, fmt.Sprintf("subscription expire task failed: %v", err))
logger.LogWarn(ctx, common.LogText("subscription expire task failed: %v", err))
return
}
if n == 0 {
@@ -70,7 +69,7 @@ func runSubscriptionQuotaResetOnce() {
for {
n, err := model.ResetDueSubscriptions(subscriptionResetBatchSize)
if err != nil {
logger.LogWarn(ctx, fmt.Sprintf("subscription quota reset task failed: %v", err))
logger.LogWarn(ctx, common.LogText("subscription quota reset task failed: %v", err))
return
}
if n == 0 {
+1 -1
View File
@@ -125,6 +125,6 @@ func ReportCurrentSystemInstance() error {
func reportSystemInstanceWithLog() {
if err := ReportCurrentSystemInstance(); err != nil {
logger.LogWarn(context.Background(), fmt.Sprintf("system instance report failed: %v", err))
logger.LogWarn(context.Background(), common.LogText("system instance report failed: %v", err))
}
}
+11 -11
View File
@@ -128,7 +128,7 @@ func StartSystemTaskRunner() {
runnerID := fmt.Sprintf("%s-%s", common.NodeName, common.GetRandomString(8))
gopool.Go(func() {
logger.LogInfo(context.Background(), fmt.Sprintf("system task runner started: runner=%s idle_interval=%s", runnerID, systemTaskRunnerIdleInterval))
logger.LogInfo(context.Background(), common.LogText("system task runner started: runner=%s idle_interval=%s", runnerID, systemTaskRunnerIdleInterval))
ticker := time.NewTicker(systemTaskRunnerIdleInterval)
defer ticker.Stop()
@@ -143,7 +143,7 @@ func StartSystemTaskRunner() {
if now.Sub(lastStaleLockCleanup) >= systemTaskStaleLockInterval {
lastStaleLockCleanup = now
if err := model.ExpireStaleSystemTaskLocks(common.GetTimestamp()); err != nil {
logger.LogWarn(context.Background(), fmt.Sprintf("system task stale lock cleanup failed: %v", err))
logger.LogWarn(context.Background(), common.LogText("system task stale lock cleanup failed: %v", err))
}
}
if now.Sub(lastScheduler) >= systemTaskSchedulerInterval {
@@ -230,7 +230,7 @@ func runSystemTaskClaimPass(runnerID string) {
}
pendingTasks, err := model.FindEarliestPendingSystemTasks(taskTypes)
if err != nil {
logger.LogWarn(context.Background(), fmt.Sprintf("system task runner query failed: %v", err))
logger.LogWarn(context.Background(), common.LogText("system task runner query failed: %v", err))
return
}
for _, handler := range handlers {
@@ -240,7 +240,7 @@ func runSystemTaskClaimPass(runnerID string) {
}
claimedTask, claimed, err := model.ClaimSystemTask(task.ID, handler.Type(), runnerID, systemTaskLockUntil())
if err != nil {
logger.LogWarn(context.Background(), fmt.Sprintf("system task claim failed: %v", err))
logger.LogWarn(context.Background(), common.LogText("system task claim failed: %v", err))
continue
}
if !claimed {
@@ -275,7 +275,7 @@ func runSystemTaskScheduler() {
}
latestTasks, err := model.GetLatestSystemTasks(taskTypes)
if err != nil {
logger.LogWarn(context.Background(), fmt.Sprintf("system task scheduler query failed: %v", err))
logger.LogWarn(context.Background(), common.LogText("system task scheduler query failed: %v", err))
return
}
for _, scheduled := range scheduledHandlers {
@@ -310,9 +310,9 @@ func runSystemTaskScheduler() {
continue
}
if activeErr != nil {
logger.LogWarn(context.Background(), fmt.Sprintf("system task scheduler active lookup failed: type=%s err=%v", scheduled.Type(), activeErr))
logger.LogWarn(context.Background(), common.LogText("system task scheduler active lookup failed: type=%s err=%v", scheduled.Type(), activeErr))
}
logger.LogWarn(context.Background(), fmt.Sprintf("system task scheduler create failed: type=%s err=%v", scheduled.Type(), err))
logger.LogWarn(context.Background(), common.LogText("system task scheduler create failed: type=%s err=%v", scheduled.Type(), err))
continue
}
}
@@ -525,16 +525,16 @@ func NewSystemTaskProgressReporter(task *model.SystemTask, runnerID string) func
}
func failSystemTask(task *model.SystemTask, runnerID string, err error) {
logger.LogWarn(context.Background(), fmt.Sprintf("system task %s failed: %v", task.TaskID, err))
logger.LogWarn(context.Background(), common.LogText("system task %s failed: %v", task.TaskID, err))
if finishErr := model.FinishSystemTask(task.TaskID, runnerID, model.SystemTaskStatusFailed, nil, err.Error()); finishErr != nil {
logger.LogWarn(context.Background(), fmt.Sprintf("system task %s failed to save failure state: %v", task.TaskID, finishErr))
logger.LogWarn(context.Background(), common.LogText("system task %s failed to save failure state: %v", task.TaskID, finishErr))
}
}
func logSystemTaskLockError(ctx context.Context, task *model.SystemTask, err error) {
if errors.Is(err, model.ErrSystemTaskLockLost) {
logger.LogWarn(ctx, fmt.Sprintf("system task %s lock lost", task.TaskID))
logger.LogWarn(ctx, common.LogText("system task %s lock lost", task.TaskID))
return
}
logger.LogWarn(ctx, fmt.Sprintf("system task %s update failed: %v", task.TaskID, err))
logger.LogWarn(ctx, common.LogText("system task %s update failed: %v", task.TaskID, err))
}
+70 -19
View File
@@ -13,7 +13,9 @@ import (
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/billingexpr"
"github.com/QuantumNous/new-api/pkg/jsplugin"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/QuantumNous/new-api/types"
"github.com/gin-gonic/gin"
@@ -23,10 +25,10 @@ import (
// 实际扣费已由 BillingSession(PreConsumeBilling + SettleBilling)完成。
func LogTaskConsumption(c *gin.Context, info *relaycommon.RelayInfo, task *model.Task) {
tokenName := c.GetString("token_name")
logContent := fmt.Sprintf("操作 %s", info.Action)
logContent := []*common.Message{common.NewMessage("Action {{action}}", map[string]any{"action": info.Action})}
// 支持任务仅按次计费
if common.StringsContains(constant.TaskPricePatches, info.OriginModelName) {
logContent = fmt.Sprintf("%s,按次计费", logContent)
logContent = append(logContent, common.NewMessage("Billed per request"))
} else {
var contents []string
if otherRatios := info.PriceData.OtherRatios(); len(otherRatios) > 0 {
@@ -42,12 +44,15 @@ func LogTaskConsumption(c *gin.Context, info *relaycommon.RelayInfo, task *model
}
}
if len(contents) > 0 {
logContent = fmt.Sprintf("%s, 计算参数:%s", logContent, strings.Join(contents, ", "))
logContent = append(logContent, common.NewMessage("Billing parameters: {{params}}", map[string]any{"params": strings.Join(contents, ", ")}))
}
}
other := model.NewLogOther()
other.SetPublic("is_task", true)
other.SetPublic("request_path", c.Request.URL.Path)
if taskDeliveredInline(c, task) {
other.SetPublic("task_sync", true)
}
other.SetPublic("model_price", info.PriceData.ModelPrice)
if info.PriceData.ModelRatio > 0 {
other.SetPublic("model_ratio", info.PriceData.ModelRatio)
@@ -67,6 +72,9 @@ func LogTaskConsumption(c *gin.Context, info *relaycommon.RelayInfo, task *model
if len(snap.UsageFacts) > 0 {
other.SetPublic("usage_facts", snap.UsageFacts)
}
setTaskImageCount(other, snap.UsageFacts["image_count"])
} else {
setTaskImageCount(other, info.PriceData.OtherRatios()["image_count"])
}
appendTaskLogInfo(task, other)
attachQuotaSaturation(c, info, other)
@@ -93,7 +101,7 @@ func LogTaskConsumption(c *gin.Context, info *relaycommon.RelayInfo, task *model
func resolveTokenKey(ctx context.Context, tokenId int, taskID string) string {
token, err := model.GetTokenById(tokenId)
if err != nil {
logger.LogWarn(ctx, fmt.Sprintf("获取令牌 key 失败 (tokenId=%d, task=%s): %s", tokenId, taskID, err.Error()))
logger.LogWarn(ctx, common.LogText("failed to get token key (tokenId=%d, task=%s): %s", tokenId, taskID, err.Error()))
return ""
}
return token.Key
@@ -132,13 +140,15 @@ func taskAdjustTokenQuota(ctx context.Context, task *model.Task, delta int) {
err = model.IncreaseTokenQuota(task.PrivateData.TokenId, tokenKey, -delta)
}
if err != nil {
logger.LogWarn(ctx, fmt.Sprintf("调整令牌额度失败 (delta=%d, task=%s): %s", delta, task.TaskID, err.Error()))
logger.LogWarn(ctx, common.LogText("failed to adjust token quota (delta=%d, task=%s): %s", delta, task.TaskID, err.Error()))
}
}
// taskBillingOther 从 task 的 BillingContext 构建日志 Other 字段。
func taskBillingOther(task *model.Task) *model.LogOther {
other := model.NewLogOther()
// 前端列表按 is_task 选择任务档位渲染分支,缺失会把差额/退款日志误判成普通 relay 动态计费
other.SetPublic("is_task", true)
if bc := task.PrivateData.BillingContext; bc != nil {
other.SetPublic("model_price", bc.ModelPrice)
if bc.ModelRatio > 0 {
@@ -148,7 +158,7 @@ func taskBillingOther(task *model.Task) *model.LogOther {
if priceData := taskBillingContextPriceData(bc); priceData != nil {
for k, v := range priceData.OtherRatios() {
if !other.SetPublic(k, v) {
common.SysError("task billing other ratio key rejected: " + k)
common.SysError(common.LogText("task billing other ratio key rejected: %s", k))
}
}
}
@@ -159,6 +169,9 @@ func taskBillingOther(task *model.Task) *model.LogOther {
if len(snap.UsageFacts) > 0 {
other.SetPublic("usage_facts", snap.UsageFacts)
}
setTaskImageCount(other, snap.UsageFacts["image_count"])
} else if priceData := taskBillingContextPriceData(bc); priceData != nil {
setTaskImageCount(other, priceData.OtherRatios()["image_count"])
}
}
props := task.Properties
@@ -170,6 +183,35 @@ func taskBillingOther(task *model.Task) *model.LogOther {
return other
}
// setTaskImageCount publishes the billed image quantity of an image task as
// other.image_count, the same field the HTTP image relay writes, so the log
// detail shows one "billable image count" regardless of the serving path. The
// value is a host-validated count fact (at most dto.MaxImageN), never a quota.
func setTaskImageCount(other *model.LogOther, value any) {
count, ok := value.(float64)
if !ok || count < 0 || count > float64(dto.MaxImageN) {
return
}
other.SetPublic("image_count", common.QuotaRound(count))
}
// taskDeliveredInline reports whether the submitting HTTP request itself
// returns the task deliverable: the upstream completed immediately, or the
// request came through the synchronous OpenAI Images protocol, where the host
// waits for an asynchronous upstream task before answering. The usage log
// presents such requests as synchronous instead of as asynchronous jobs.
func taskDeliveredInline(c *gin.Context, task *model.Task) bool {
if task != nil && (task.Status == model.TaskStatusSuccess || task.Status == model.TaskStatusFailure) {
return true
}
pinnedValue, exists := c.Get(jsplugin.ContextKeyPinnedEndpoint)
if !exists {
return false
}
pinned, ok := pinnedValue.(jsplugin.PinnedEndpoint)
return ok && pinned.Protocol == jsplugin.ProtocolOpenAIImage
}
func appendTaskLogInfo(task *model.Task, other *model.LogOther) {
if task == nil || other == nil {
return
@@ -177,6 +219,11 @@ func appendTaskLogInfo(task *model.Task, other *model.LogOther) {
if task.TaskID != "" {
other.SetPublic("task_id", task.TaskID)
}
if task.PrivateData.ResultDiscarded {
// The result was delivered inline and the upstream snapshot was not
// persisted, so no artifact can be retrieved for this task.
other.SetPublic("result_discarded", true)
}
if task.PrivateData.Execution != nil {
AppendTaskPluginAuditInfo(other, task.PrivateData.Execution.TaskPlugin)
}
@@ -221,7 +268,7 @@ func RefundTaskQuota(ctx context.Context, task *model.Task, reason string) bool
// 1. 退还资金来源(钱包或订阅)
if err := taskAdjustFunding(task, -quota); err != nil {
logger.LogWarn(ctx, fmt.Sprintf("退还资金来源失败 task %s: %s", task.TaskID, err.Error()))
logger.LogWarn(ctx, common.LogText("failed to refund the funding source, task %s: %s", task.TaskID, err.Error()))
return false
}
@@ -239,7 +286,6 @@ func RefundTaskQuota(ctx context.Context, task *model.Task, reason string) bool
model.RecordTaskBillingLog(model.RecordTaskBillingLogParams{
UserId: task.UserId,
LogType: model.LogTypeRefund,
Content: "",
ChannelId: task.ChannelId,
ModelName: taskModelName(task),
Quota: quota,
@@ -252,16 +298,16 @@ func RefundTaskQuota(ctx context.Context, task *model.Task, reason string) bool
// 回写失败必须显式告警,避免漏掉潜在的重复退款风险。
task.Quota = 0
if err := task.UpdateQuota(); err != nil {
logger.LogError(ctx, fmt.Sprintf("退款成功但清除 task quota 失败 task %s: %s", task.TaskID, err.Error()))
logger.LogError(ctx, common.LogText("refund succeeded but clearing the task quota failed, task %s: %s", task.TaskID, err.Error()))
}
return true
}
// RecalculateTaskQuota 通用的异步差额结算。
// actualQuota 是任务完成后的实际应扣额度,与预扣额度 (task.Quota) 做差额结算。
// reason 用于日志记录(例如 "token重算" 或 "adaptor调整")。
// reason 写入日志内容(例如按 token 重算或 adaptor 调整)。
// clamps 可选:若计算 actualQuota 时发生额度饱和,将其记入日志 admin_info(仅管理员可见)。
func RecalculateTaskQuota(ctx context.Context, task *model.Task, actualQuota int, reason string, clamps ...*common.QuotaClamp) {
func RecalculateTaskQuota(ctx context.Context, task *model.Task, actualQuota int, reason *common.Message, clamps ...*common.QuotaClamp) {
if actualQuota < 0 {
return
}
@@ -269,22 +315,22 @@ func RecalculateTaskQuota(ctx context.Context, task *model.Task, actualQuota int
quotaDelta := actualQuota - preConsumedQuota
if quotaDelta == 0 {
logger.LogInfo(ctx, fmt.Sprintf("任务 %s 预扣费准确(%s,%s)",
task.TaskID, logger.LogQuota(actualQuota), reason))
logger.LogInfo(ctx, common.LogText("task %s pre-consumed quota matches the actual quota (%s, %s)",
task.TaskID, logger.LogQuota(actualQuota), reason.Error()))
return
}
logger.LogInfo(ctx, fmt.Sprintf("任务 %s 差额结算:delta=%s(实际:%s,预扣:%s,%s)",
logger.LogInfo(ctx, common.LogText("task %s settles the difference: delta=%s (actual %s, pre-consumed %s, %s)",
task.TaskID,
logger.LogQuota(quotaDelta),
logger.LogQuota(actualQuota),
logger.LogQuota(preConsumedQuota),
reason,
reason.Error(),
))
// 调整资金来源
if err := taskAdjustFunding(task, quotaDelta); err != nil {
logger.LogError(ctx, fmt.Sprintf("差额结算资金调整失败 task %s: %s", task.TaskID, err.Error()))
logger.LogError(ctx, common.LogText("failed to adjust funding for the difference, task %s: %s", task.TaskID, err.Error()))
return
}
@@ -293,7 +339,7 @@ func RecalculateTaskQuota(ctx context.Context, task *model.Task, actualQuota int
task.Quota = actualQuota
if err := task.UpdateQuota(); err != nil {
logger.LogError(ctx, fmt.Sprintf("差额结算回写 quota 失败 task %s: %s", task.TaskID, err.Error()))
logger.LogError(ctx, common.LogText("failed to write back the settled task quota, task %s: %s", task.TaskID, err.Error()))
}
// 提交阶段已经累计过一次请求;结算阶段只调整最终用量。
@@ -319,7 +365,7 @@ func RecalculateTaskQuota(ctx context.Context, task *model.Task, actualQuota int
model.RecordTaskBillingLog(model.RecordTaskBillingLogParams{
UserId: task.UserId,
LogType: logType,
Content: reason,
Content: []*common.Message{reason},
ChannelId: task.ChannelId,
ModelName: taskModelName(task),
Quota: logQuota,
@@ -383,7 +429,12 @@ func RecalculateTaskQuotaByTokens(ctx context.Context, task *model.Task, totalTo
// 计算实际应扣费额度: totalTokens * modelRatio * groupRatio * otherMultiplier(饱和转换,防止溢出成负数)
actualQuota, clamp := common.QuotaFromFloatChecked(float64(totalTokens) * modelRatio * finalGroupRatio * otherMultiplier)
reason := fmt.Sprintf("token重算:tokens=%d, modelRatio=%.2f, groupRatio=%.2f, otherMultiplier=%.4f", totalTokens, modelRatio, finalGroupRatio, otherMultiplier)
reason := common.NewMessage("Recalculated from tokens: tokens {{tokens}}, model ratio {{model_ratio}}, group ratio {{group_ratio}}, other multiplier {{other_multiplier}}", map[string]any{
"tokens": totalTokens,
"model_ratio": fmt.Sprintf("%.2f", modelRatio),
"group_ratio": fmt.Sprintf("%.2f", finalGroupRatio),
"other_multiplier": fmt.Sprintf("%.4f", otherMultiplier),
})
RecalculateTaskQuota(ctx, task, actualQuota, reason, clamp)
return true
}
+77 -11
View File
@@ -13,8 +13,10 @@ import (
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/billingexpr"
"github.com/QuantumNous/new-api/pkg/jsplugin"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/QuantumNous/new-api/types"
@@ -44,6 +46,9 @@ func TestMain(m *testing.M) {
common.RedisEnabled = false
common.BatchUpdateEnabled = false
common.LogConsumeEnabled = true
if err := i18n.Init(); err != nil {
panic("failed to load locales: " + err.Error())
}
if err := db.AutoMigrate(
&model.Task{},
@@ -341,7 +346,7 @@ func TestLogTaskConsumptionIncludesTieredSnapshotUsageFacts(t *testing.T) {
assert.Equal(t, float64(5), facts["seconds"])
assert.NotContains(t, other, "resolution")
assert.NotContains(t, other, "seconds")
assert.Contains(t, log.Content, "计算参数:")
assert.Contains(t, log.Content, "Billing parameters: ")
assert.Contains(t, log.Content, "resolution: 720P")
assert.Contains(t, log.Content, "seconds: 5")
}
@@ -379,8 +384,69 @@ func TestLogTaskConsumptionWithoutSnapshotKeepsRatioMode(t *testing.T) {
assert.NotContains(t, other, "expr_b64")
assert.NotContains(t, other, "matched_tier")
assert.NotContains(t, other, "usage_facts")
assert.Contains(t, log.Content, "计算参数:")
assert.Contains(t, log.Content, "size: 2.00")
assert.Equal(t, "Action GENERATE, Billing parameters: size: 2.00", log.Content)
// The web console renders content_parts in the viewer's language; content
// keeps the English text.
assert.Equal(t, []any{
map[string]any{"key": "Action {{action}}", "params": map[string]any{"action": "GENERATE"}},
map[string]any{"key": "Billing parameters: {{params}}", "params": map[string]any{"params": "size: 2.00"}},
}, other["content_parts"])
}
// Task logs distinguish jobs the client polls from requests whose HTTP call
// returned the deliverable itself, and flag results the gateway did not keep.
func TestLogTaskConsumptionMarksInlineResultsAndDiscardedArtifacts(t *testing.T) {
for _, tc := range []struct {
name string
status model.TaskStatus
discarded bool
pinnedProtocol string
wantSync, wantKept bool
}{
{"asynchronous job", model.TaskStatusNotStart, false, "", false, true},
{"immediate result on a discarding route", model.TaskStatusSuccess, true, "", true, false},
{"openai image request waits for an asynchronous upstream task", model.TaskStatusNotStart, false, jsplugin.ProtocolOpenAIImage, true, true},
{"immediate failure", model.TaskStatusFailure, false, "", true, true},
} {
t.Run(tc.name, func(t *testing.T) {
truncate(t)
const userID, channelID = 43, 43
seedUser(t, userID, 10_000)
seedChannel(t, channelID)
task := makeTask(userID, channelID, 100, 0, BillingSourceWallet, 0)
task.Status = tc.status
task.PrivateData.ResultDiscarded = tc.discarded
info := &relaycommon.RelayInfo{
UserId: userID, OriginModelName: "qwen-image-plus", UsingGroup: "default",
ChannelMeta: &relaycommon.ChannelMeta{ChannelId: channelID},
TaskRelayInfo: &relaycommon.TaskRelayInfo{Action: "text_to_image"},
PriceData: types.PriceData{ModelPrice: 0.03, Quota: 100, GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 1}},
}
gin.SetMode(gin.TestMode)
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
ctx.Set("token_name", "test_token")
if tc.pinnedProtocol != "" {
ctx.Set(jsplugin.ContextKeyPinnedEndpoint, jsplugin.PinnedEndpoint{Protocol: tc.pinnedProtocol})
}
LogTaskConsumption(ctx, info, task)
log := getLastLog(t)
require.NotNil(t, log)
var other map[string]any
require.NoError(t, common.UnmarshalJsonStr(log.Other, &other))
assert.Equal(t, true, other["is_task"])
if tc.wantSync {
assert.Equal(t, true, other["task_sync"])
} else {
assert.NotContains(t, other, "task_sync")
}
if tc.wantKept {
assert.NotContains(t, other, "result_discarded")
} else {
assert.Equal(t, true, other["result_discarded"])
}
})
}
}
func TestTaskBillingOtherSeparatesPluginAndRootDiagnostics(t *testing.T) {
@@ -621,7 +687,7 @@ func TestMidjourneyRefundRestoresEveryAccountingElementOnBillingChannel(t *testi
seedChargedAccounting(t, userID, billingChannelID, tokenID, chargedQuota, 1)
assert.True(t, RefundMidjourneyQuota(ctx, task, "构图失败"))
assert.True(t, RefundMidjourneyQuota(ctx, task, "Composition failed"))
assert.Equal(t, initialUserQuota, getUserQuota(t, userID))
assert.Equal(t, initialTokenQuota, getTokenRemainQuota(t, tokenID))
assert.Zero(t, getTokenUsedQuota(t, tokenID))
@@ -981,7 +1047,7 @@ func TestRecalculate_PositiveDelta(t *testing.T) {
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
RecalculateTaskQuota(ctx, task, actualQuota, "adaptor adjustment")
RecalculateTaskQuota(ctx, task, actualQuota, common.NewMessage("adaptor adjustment"))
// User quota should decrease by the delta (1000 additional charge)
assert.Equal(t, initQuota-(actualQuota-preConsumed), getUserQuota(t, userID))
@@ -1020,7 +1086,7 @@ func TestRecalculate_NegativeDelta(t *testing.T) {
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
RecalculateTaskQuota(ctx, task, actualQuota, "adaptor adjustment")
RecalculateTaskQuota(ctx, task, actualQuota, common.NewMessage("adaptor adjustment"))
// User quota should increase by abs(delta) = 2000 (refund overpayment)
assert.Equal(t, initQuota+(preConsumed-actualQuota), getUserQuota(t, userID))
@@ -1054,7 +1120,7 @@ func TestRecalculate_ZeroDelta(t *testing.T) {
task := makeTask(userID, 0, preConsumed, 0, BillingSourceWallet, 0)
RecalculateTaskQuota(ctx, task, preConsumed, "exact match")
RecalculateTaskQuota(ctx, task, preConsumed, common.NewMessage("exact match"))
// No change to user quota
assert.Equal(t, initQuota, getUserQuota(t, userID))
@@ -1075,7 +1141,7 @@ func TestRecalculate_ActualQuotaZero(t *testing.T) {
task := makeTask(userID, 0, preConsumed, 0, BillingSourceWallet, 0)
require.NoError(t, model.DB.Create(task).Error)
RecalculateTaskQuota(ctx, task, 0, "zero actual")
RecalculateTaskQuota(ctx, task, 0, common.NewMessage("zero actual"))
assert.Equal(t, initQuota+preConsumed, getUserQuota(t, userID))
assert.Zero(t, task.Quota)
@@ -1094,7 +1160,7 @@ func TestRecalculate_RejectsNegativeActualQuota(t *testing.T) {
seedUser(t, userID, initQuota)
task := makeTask(userID, 0, preConsumed, 0, BillingSourceWallet, 0)
RecalculateTaskQuota(ctx, task, -1, "invalid negative actual")
RecalculateTaskQuota(ctx, task, -1, common.NewMessage("invalid negative actual"))
assert.Equal(t, initQuota, getUserQuota(t, userID))
assert.Equal(t, preConsumed, task.Quota)
@@ -1119,7 +1185,7 @@ func TestRecalculate_Subscription_NegativeDelta(t *testing.T) {
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceSubscription, subID)
RecalculateTaskQuota(ctx, task, actualQuota, "subscription over-charge")
RecalculateTaskQuota(ctx, task, actualQuota, common.NewMessage("subscription over-charge"))
// Subscription used should decrease by delta (refund 3000)
assert.Equal(t, subUsed-int64(preConsumed-actualQuota), getSubscriptionUsed(t, subID))
@@ -1186,7 +1252,7 @@ func simulatePollBilling(ctx context.Context, task *model.Task, newStatus model.
}
if shouldSettle && actualQuota > 0 {
RecalculateTaskQuota(ctx, task, actualQuota, "test settle")
RecalculateTaskQuota(ctx, task, actualQuota, common.NewMessage("test settle"))
}
if shouldRefund {
RefundTaskQuota(ctx, task, task.FailReason)
+25
View File
@@ -83,3 +83,28 @@ func TestBuildTaskPluginViewOmitsPrivatePollState(t *testing.T) {
assert.NotContains(t, payload, "poll_failures")
assert.NotContains(t, payload, "private_data")
}
func TestBuildTaskPluginViewEncodesKeysInMapOrder(t *testing.T) {
task := &model.Task{
TaskID: "task_public_order",
Platform: "fixture",
Status: model.TaskStatusSuccess,
Progress: "100%",
FailReason: "none",
CreatedAt: 1700000000,
UpdatedAt: 1700000100,
FinishTime: 1700000200,
Data: []byte(`{"b":1,"a":{"10":true,"2":false}}`),
}
view, err := BuildTaskPluginView(task)
require.NoError(t, err)
encoded, err := common.Marshal(view)
require.NoError(t, err)
var payload map[string]any
require.NoError(t, common.Unmarshal(encoded, &payload))
require.Len(t, payload, 9)
reencoded, err := common.Marshal(payload)
require.NoError(t, err)
assert.Equal(t, string(reencoded), string(encoded))
}
+80 -78
View File
@@ -15,6 +15,7 @@ import (
taskdto "github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model"
perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
"github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
relaycommon "github.com/QuantumNous/new-api/relay/common"
@@ -76,8 +77,8 @@ func sweepTimedOutTasks(ctx context.Context) {
return
}
reason := fmt.Sprintf("任务超时(%d分钟)", constant.TaskTimeoutMinutes)
legacyReason := "任务超时(旧系统遗留任务,不进行退款,请联系管理员)"
reason := "Task timed out"
legacyReason := "Task timed out (a task left by the old system is not refunded; contact the administrator)"
now := time.Now().Unix()
timedOutCount := 0
@@ -99,11 +100,11 @@ func sweepTimedOutTasks(ctx context.Context) {
won, err := task.UpdateWithStatus(oldStatus)
if err != nil {
logger.LogError(ctx, fmt.Sprintf("sweepTimedOutTasks CAS update error for task %s: %v", task.TaskID, err))
logger.LogError(ctx, common.LogText("sweepTimedOutTasks CAS update error for task %s: %v", task.TaskID, err))
continue
}
if !won {
logger.LogInfo(ctx, fmt.Sprintf("sweepTimedOutTasks: task %s already transitioned, skip", task.TaskID))
logger.LogInfo(ctx, common.LogText("sweepTimedOutTasks: task %s already transitioned, skip", task.TaskID))
continue
}
timedOutCount++
@@ -113,7 +114,7 @@ func sweepTimedOutTasks(ctx context.Context) {
}
if timedOutCount > 0 {
logger.LogInfo(ctx, fmt.Sprintf("sweepTimedOutTasks: timed out %d tasks", timedOutCount))
logger.LogInfo(ctx, common.LogText("sweepTimedOutTasks: timed out %d tasks", timedOutCount))
}
}
@@ -139,7 +140,7 @@ func RunTaskPollingOnce(ctx context.Context, report func(processed, total int))
ctx = context.Background()
}
common.SysLog("任务进度轮询开始")
common.SysLog(common.LogText("task progress polling started"))
sweepTimedOutTasks(ctx)
allTasks := model.GetAllUnFinishSyncTasks(constant.TaskQueryLimit)
summary.UnfinishedTasks = len(allTasks)
@@ -182,9 +183,9 @@ func RunTaskPollingOnce(ctx context.Context, report func(processed, total int))
"progress": "100%",
})
if err != nil {
logger.LogError(ctx, fmt.Sprintf("Fix null task_id task error: %v", err))
logger.LogError(ctx, common.LogText("Fix null task_id task error: %v", err))
} else {
logger.LogInfo(ctx, fmt.Sprintf("Fix null task_id task success: %v", nullTaskIds))
logger.LogInfo(ctx, common.LogText("Fix null task_id task success: %v", nullTaskIds))
}
}
if len(taskChannelM) == 0 {
@@ -196,7 +197,7 @@ func RunTaskPollingOnce(ctx context.Context, report func(processed, total int))
if report != nil && ctx.Err() == nil {
report(totalPlatforms, totalPlatforms)
}
common.SysLog("任务进度轮询完成")
common.SysLog(common.LogText("task progress polling finished"))
return summary
}
@@ -212,12 +213,12 @@ func DispatchPlatformUpdate(ctx context.Context, platform constant.TaskPlatform,
adaptor := GetTaskAdaptorFunc(platform)
if batchAdaptor, ok := adaptor.(BatchTaskPollingAdaptor); ok && batchAdaptor.FetchMode() == "batch" {
if err := UpdateBatchTasks(ctx, batchAdaptor, taskChannelM, taskM); err != nil {
common.SysLog(fmt.Sprintf("UpdateBatchTasks fail: %s", err))
common.SysLog(common.LogText("UpdateBatchTasks fail: %s", err))
}
return
}
if err := UpdateVideoTasks(ctx, platform, taskChannelM, taskM); err != nil {
common.SysLog(fmt.Sprintf("UpdateVideoTasks fail: %s", err))
common.SysLog(common.LogText("UpdateVideoTasks fail: %s", err))
}
}
@@ -228,14 +229,14 @@ func UpdateBatchTasks(ctx context.Context, adaptor BatchTaskPollingAdaptor, task
}
err := updateBatchTasks(ctx, adaptor, channelId, taskIds, taskM)
if err != nil {
logger.LogError(ctx, fmt.Sprintf("渠道 #%d 更新异步任务失败: %s", channelId, err.Error()))
logger.LogError(ctx, common.LogText("channel #%d failed to update async tasks: %s", channelId, err.Error()))
}
}
return nil
}
func updateBatchTasks(ctx context.Context, adaptor BatchTaskPollingAdaptor, channelId int, taskIds []string, taskM map[string]*model.Task) error {
logger.LogInfo(ctx, fmt.Sprintf("渠道 #%d 未完成的任务有: %d", channelId, len(taskIds)))
logger.LogInfo(ctx, common.LogText("channel #%d has %d unfinished tasks", channelId, len(taskIds)))
if ctx.Err() != nil {
return ctx.Err()
}
@@ -244,7 +245,7 @@ func updateBatchTasks(ctx context.Context, adaptor BatchTaskPollingAdaptor, chan
}
ch, err := model.CacheGetChannel(channelId)
if err != nil {
common.SysLog(fmt.Sprintf("CacheGetChannel: %v", err))
common.SysLog(common.LogText("CacheGetChannel: %v", err))
// Collect DB primary key IDs for bulk update (taskIds are upstream IDs, not task_id column values)
var failedIDs []int64
for _, upstreamID := range taskIds {
@@ -253,12 +254,12 @@ func updateBatchTasks(ctx context.Context, adaptor BatchTaskPollingAdaptor, chan
}
}
err = model.TaskBulkUpdateByID(failedIDs, map[string]any{
"fail_reason": fmt.Sprintf("获取渠道信息失败,请联系管理员,渠道ID:%d", channelId),
"fail_reason": "Failed to get the channel information. Please contact the administrator",
"status": "FAILURE",
"progress": "100%",
})
if err != nil {
common.SysLog(fmt.Sprintf("UpdateSunoTask error: %v", err))
common.SysLog(common.LogText("UpdateSunoTask error: %v", err))
}
return err
}
@@ -273,26 +274,33 @@ func updateBatchTasks(ctx context.Context, adaptor BatchTaskPollingAdaptor, chan
tasks = append(tasks, task)
}
}
// The channel type tells plugin adaptors whether the upstream is another
// New API gateway, the same signal submission derives from the request.
info := &relaycommon.RelayInfo{}
info.ChannelMeta = &relaycommon.ChannelMeta{ChannelBaseUrl: baseURL}
info.ChannelMeta = &relaycommon.ChannelMeta{ChannelType: ch.Type, ChannelId: ch.Id, ChannelBaseUrl: baseURL, ChannelSetting: ch.GetSetting()}
info.ApiKey = ch.Key
adaptor.Init(info)
resp, err := adaptor.FetchBatchTasks(baseURL, ch.Key, tasks, proxy)
if err != nil {
common.SysLog(fmt.Sprintf("Get Task Do req error: %v", err))
common.SysLog(common.LogText("Get Task Do req error: %v", err))
return recordPollFailureForTasks(ctx, adaptor, tasks, pollClassTransport, 0, err.Error())
}
defer resp.Body.Close()
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
common.SysLog(fmt.Sprintf("Get Suno Task parse body error: %v", err))
common.SysLog(common.LogText("Get Suno Task parse body error: %v", err))
return recordPollFailureForTasks(ctx, adaptor, tasks, pollClassTransport, resp.StatusCode, err.Error())
}
switch classifyPollHTTP(resp.StatusCode) {
case pollClassNotFound:
return failTasksFromPoll(ctx, adaptor, tasks, fmt.Sprintf("upstream task not found (HTTP %d)", resp.StatusCode))
for _, task := range tasks {
if task != nil {
logger.LogWarn(ctx, common.LogText("task %s poll %s (failures=%d, http=%d): %s", task.TaskID, pollClassNotFound, task.PrivateData.PollFailures, resp.StatusCode, unrecognizedPollDetail("", responseBody)))
}
}
return failTasksFromPoll(ctx, adaptor, tasks, "Upstream task not found")
case pollClassAuth:
logger.LogWarn(ctx, fmt.Sprintf("task poll auth failure channel_id=%d http=%d", channelId, resp.StatusCode))
logger.LogWarn(ctx, common.LogText("task poll auth failure channel_id=%d http=%d", channelId, resp.StatusCode))
return recordPollFailureForTasks(ctx, adaptor, tasks, pollClassAuth, resp.StatusCode, "")
case pollClassTransient:
return recordPollFailureForTasks(ctx, adaptor, tasks, pollClassTransient, resp.StatusCode, "")
@@ -307,7 +315,7 @@ func updateBatchTasks(ctx context.Context, adaptor BatchTaskPollingAdaptor, chan
}
task := taskM[upstreamID]
if task == nil {
logger.LogWarn(ctx, fmt.Sprintf("Batch task response ignored: unknown task_id=%s", upstreamID))
logger.LogWarn(ctx, common.LogText("Batch task response ignored: unknown task_id=%s", upstreamID))
continue
}
snap := task.Snapshot()
@@ -315,13 +323,13 @@ func updateBatchTasks(ctx context.Context, adaptor BatchTaskPollingAdaptor, chan
parsedStatus := model.TaskStatus(responseItem.TaskInfo.Status)
if parsedStatus == model.TaskStatusUnknown || parsedStatus == "" || !knownPollStatus(parsedStatus) {
if err := recordPollFailure(ctx, adaptor, task, snap.Status, pollClassUnrecognized, resp.StatusCode, responseItem.TaskInfo.Reason); err != nil {
common.SysLog("UpdateSunoTask task error: " + err.Error())
common.SysLog(common.LogText("UpdateSunoTask task error: %s", err.Error()))
}
continue
}
if httpClass == pollClassOtherClient && isNonTerminalPollStatus(parsedStatus) {
if err := recordPollFailure(ctx, adaptor, task, snap.Status, pollClassUnrecognized, resp.StatusCode, responseItem.TaskInfo.Reason); err != nil {
common.SysLog("UpdateSunoTask task error: " + err.Error())
common.SysLog(common.LogText("UpdateSunoTask task error: %s", err.Error()))
}
continue
}
@@ -340,7 +348,7 @@ func updateBatchTasks(ctx context.Context, adaptor BatchTaskPollingAdaptor, chan
task.Progress = responseItem.TaskInfo.Progress
}
if responseItem.TaskInfo.Reason != "" || task.Status == model.TaskStatusFailure {
logger.LogInfo(ctx, task.TaskID+" 构建失败,"+task.FailReason)
logger.LogInfo(ctx, common.LogText("%s failed: %s", task.TaskID, task.FailReason))
task.Status = model.TaskStatusFailure
task.Progress = "100%"
}
@@ -350,7 +358,7 @@ func updateBatchTasks(ctx context.Context, adaptor BatchTaskPollingAdaptor, chan
if responseItem.Data != nil {
task.SetData(responseItem.Data)
} else if task.Status == model.TaskStatusSuccess || task.Status == model.TaskStatusFailure {
logger.LogWarn(ctx, fmt.Sprintf(
logger.LogWarn(ctx, common.LogText(
"Batch task %s reached terminal status without data; preserving existing task data",
task.TaskID,
))
@@ -363,18 +371,15 @@ func updateBatchTasks(ctx context.Context, adaptor BatchTaskPollingAdaptor, chan
terminalTransition := isDone && snap.Status != task.Status
won, updateErr := task.UpdateWithStatus(snap.Status)
if updateErr != nil {
common.SysLog("UpdateSunoTask task error: " + updateErr.Error())
common.SysLog(common.LogText("UpdateSunoTask task error: %s", updateErr.Error()))
continue
}
if !won {
logger.LogWarn(ctx, fmt.Sprintf("Batch task %s already transitioned by another process, skip billing", task.TaskID))
logger.LogWarn(ctx, common.LogText("Batch task %s already transitioned by another process, skip billing", task.TaskID))
continue
}
if terminalTransition {
billingSettled := settleTaskBillingOnComplete(ctx, adaptor, task, &responseItem.TaskInfo)
if task.Status == model.TaskStatusFailure && !billingSettled && task.Quota != 0 {
RefundTaskQuota(ctx, task, task.FailReason)
}
finalizeTerminalTask(ctx, adaptor, task, &responseItem.TaskInfo)
}
}
return nil
@@ -400,7 +405,7 @@ func UpdateVideoTasks(ctx context.Context, platform constant.TaskPlatform, taskC
gopool.Go(func() {
defer wg.Done()
if err := updateVideoTasks(ctx, platform, channelId, taskIds, taskM); err != nil {
logger.LogError(ctx, fmt.Sprintf("Channel #%d failed to update video async tasks: %s", channelId, err.Error()))
logger.LogError(ctx, common.LogText("Channel #%d failed to update video async tasks: %s", channelId, err.Error()))
}
})
}
@@ -412,7 +417,7 @@ func UpdateVideoTasks(ctx context.Context, platform constant.TaskPlatform, taskC
}
func updateVideoTasks(ctx context.Context, platform constant.TaskPlatform, channelId int, taskIds []string, taskM map[string]*model.Task) error {
logger.LogInfo(ctx, fmt.Sprintf("Channel #%d pending video tasks: %d", channelId, len(taskIds)))
logger.LogInfo(ctx, common.LogText("Channel #%d pending video tasks: %d", channelId, len(taskIds)))
if ctx.Err() != nil {
return ctx.Err()
}
@@ -429,12 +434,12 @@ func updateVideoTasks(ctx context.Context, platform constant.TaskPlatform, chann
}
}
errUpdate := model.TaskBulkUpdateByID(failedIDs, map[string]any{
"fail_reason": fmt.Sprintf("Failed to get channel info, channel ID: %d", channelId),
"fail_reason": "Failed to get the channel information. Please contact the administrator",
"status": "FAILURE",
"progress": "100%",
})
if errUpdate != nil {
common.SysLog(fmt.Sprintf("UpdateVideoTask error: %v", errUpdate))
common.SysLog(common.LogText("UpdateVideoTask error: %v", errUpdate))
}
return fmt.Errorf("CacheGetChannel failed: %w", err)
}
@@ -444,7 +449,10 @@ func updateVideoTasks(ctx context.Context, platform constant.TaskPlatform, chann
}
info := &relaycommon.RelayInfo{}
info.ChannelMeta = &relaycommon.ChannelMeta{
ChannelType: cacheGetChannel.Type,
ChannelId: cacheGetChannel.Id,
ChannelBaseUrl: cacheGetChannel.GetBaseURL(),
ChannelSetting: cacheGetChannel.GetSetting(),
}
info.ApiKey = cacheGetChannel.Key
adaptor.Init(info)
@@ -454,7 +462,7 @@ func updateVideoTasks(ctx context.Context, platform constant.TaskPlatform, chann
return ctx.Err()
}
if err := updateVideoSingleTask(ctx, adaptor, cacheGetChannel, taskId, taskM); err != nil {
logger.LogError(ctx, fmt.Sprintf("Failed to update video task %s: %s", taskId, err.Error()))
logger.LogError(ctx, common.LogText("Failed to update video task %s: %s", taskId, err.Error()))
}
if disablePollingSleep || i == len(taskIds)-1 {
continue
@@ -482,7 +490,7 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch *
task := taskM[taskId]
if task == nil {
logger.LogError(ctx, fmt.Sprintf("Task %s not found in taskM", taskId))
logger.LogError(ctx, common.LogText("Task %s not found in taskM", taskId))
return fmt.Errorf("task %s not found", taskId)
}
key := ch.Key
@@ -506,9 +514,10 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch *
switch classifyPollHTTP(resp.StatusCode) {
case pollClassNotFound:
return failTaskFromPoll(ctx, adaptor, task, snap.Status, fmt.Sprintf("upstream task not found (HTTP %d)", resp.StatusCode))
logger.LogWarn(ctx, common.LogText("task %s poll %s (failures=%d, http=%d): %s", task.TaskID, pollClassNotFound, task.PrivateData.PollFailures, resp.StatusCode, unrecognizedPollDetail("", responseBody)))
return failTaskFromPoll(ctx, adaptor, task, snap.Status, "Upstream task not found")
case pollClassAuth:
logger.LogWarn(ctx, fmt.Sprintf("task poll auth failure channel_id=%d task=%s http=%d", ch.Id, task.TaskID, resp.StatusCode))
logger.LogWarn(ctx, common.LogText("task poll auth failure channel_id=%d task=%s http=%d", ch.Id, task.TaskID, resp.StatusCode))
return recordPollFailure(ctx, adaptor, task, snap.Status, pollClassAuth, resp.StatusCode, "")
case pollClassTransient:
return recordPollFailure(ctx, adaptor, task, snap.Status, pollClassTransient, resp.StatusCode, "")
@@ -579,14 +588,14 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch *
}
shouldFinalizeBilling = true
case model.TaskStatusFailure:
logger.LogJson(ctx, fmt.Sprintf("Task %s failed", taskId), task)
logger.LogJson(ctx, common.LogText("Task %s failed", taskId), task)
task.Status = model.TaskStatusFailure
task.Progress = taskcommon.ProgressComplete
if task.FinishTime == 0 {
task.FinishTime = now
}
task.FailReason = taskResult.Reason
logger.LogInfo(ctx, fmt.Sprintf("Task %s failed: %s", task.TaskID, task.FailReason))
logger.LogInfo(ctx, common.LogText("Task %s failed: %s", task.TaskID, task.FailReason))
taskResult.Progress = taskcommon.ProgressComplete
shouldFinalizeBilling = true
}
@@ -598,15 +607,15 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch *
if isDone && snap.Status != task.Status {
won, err := task.UpdateWithStatus(snap.Status)
if err != nil {
logger.LogError(ctx, fmt.Sprintf("UpdateWithStatus failed for task %s: %s", task.TaskID, err.Error()))
logger.LogError(ctx, common.LogText("UpdateWithStatus failed for task %s: %s", task.TaskID, err.Error()))
shouldFinalizeBilling = false
} else if !won {
logger.LogWarn(ctx, fmt.Sprintf("Task %s CAS lost or no-op update, skip billing", task.TaskID))
logger.LogWarn(ctx, common.LogText("Task %s CAS lost or no-op update, skip billing", task.TaskID))
shouldFinalizeBilling = false
}
} else if !snap.Equal(task.Snapshot()) {
if _, err := task.UpdateWithStatus(snap.Status); err != nil {
logger.LogError(ctx, fmt.Sprintf("Failed to update task %s: %s", task.TaskID, err.Error()))
logger.LogError(ctx, common.LogText("Failed to update task %s: %s", task.TaskID, err.Error()))
}
} else {
// No changes, skip update
@@ -614,15 +623,21 @@ func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch *
}
if shouldFinalizeBilling {
billingSettled := settleTaskBillingOnComplete(ctx, adaptor, task, taskResult)
if task.Status == model.TaskStatusFailure && !billingSettled && task.Quota != 0 {
RefundTaskQuota(ctx, task, task.FailReason)
}
finalizeTerminalTask(ctx, adaptor, task, taskResult)
}
return nil
}
// finalizeTerminalTask 终态统一收尾(状态 CAS 赢家调用,恰好一次):采样 + 结算 + 失败兜底退款。
func finalizeTerminalTask(ctx context.Context, adaptor TaskPollingAdaptor, task *model.Task, taskResult *relaycommon.TaskInfo) {
perfmetrics.RecordTaskResult(task, taskResult)
billingSettled := settleTaskBillingOnComplete(ctx, adaptor, task, taskResult)
if task.Status == model.TaskStatusFailure && !billingSettled && task.Quota != 0 {
RefundTaskQuota(ctx, task, task.FailReason)
}
}
func redactVideoResponseBody(body []byte) []byte {
var m map[string]any
if err := common.Unmarshal(body, &m); err != nil {
@@ -670,25 +685,25 @@ func settleTaskBillingOnComplete(ctx context.Context, adaptor TaskPollingAdaptor
}
result, usageFacts, err := EvaluateTaskCompletionUsage(bc.TieredSnapshot, taskResult.UsageFacts)
if err != nil {
logger.LogWarn(ctx, fmt.Sprintf("任务 %s 表达式结算失败,保留预扣额度: %v", task.TaskID, err))
logger.LogWarn(ctx, common.LogText("task %s expression settlement failed, keeping the pre-consumed quota: %v", task.TaskID, err))
return true
}
if result.Clamp != nil {
logger.LogWarn(ctx, fmt.Sprintf("任务 %s 表达式结算额度发生饱和: %+v", task.TaskID, result.Clamp))
logger.LogWarn(ctx, common.LogText("task %s expression settlement quota was clamped: %+v", task.TaskID, result.Clamp))
}
bc.TieredSnapshot.UsageFacts = usageFacts
bc.TieredSnapshot.EstimatedTier = result.MatchedTier
RecalculateTaskQuota(ctx, task, result.ActualQuotaAfterGroup, "任务用量表达式结算", result.Clamp)
RecalculateTaskQuota(ctx, task, result.ActualQuotaAfterGroup, common.NewMessage("Settled by the task usage expression"), result.Clamp)
return true
}
// 按次计费的成功任务保持预扣;失败任务由调用方全额退款。
if bc := task.PrivateData.BillingContext; bc != nil && bc.PerCallBilling {
logger.LogInfo(ctx, fmt.Sprintf("任务 %s 按次计费,跳过差额结算", task.TaskID))
logger.LogInfo(ctx, common.LogText("task %s is billed per request, skipping difference settlement", task.TaskID))
return false
}
// 优先让 adaptor 决定最终额度。
if actualQuota := adaptor.AdjustBillingOnComplete(task, taskResult); actualQuota > 0 {
RecalculateTaskQuota(ctx, task, actualQuota, "adaptor计费调整")
RecalculateTaskQuota(ctx, task, actualQuota, common.NewMessage("Adjusted by the channel adaptor"))
return true
}
// 回退到 token 重算。
@@ -737,17 +752,6 @@ func isNonTerminalPollStatus(status model.TaskStatus) bool {
}
}
func pollFailureReason(class string, statusCode int, detail string) string {
reason := fmt.Sprintf("poll failed: %s", class)
if statusCode > 0 {
reason = fmt.Sprintf("poll failed: %s (HTTP %d)", class, statusCode)
}
if detail != "" {
reason = reason + ": " + detail
}
return reason
}
// unrecognizedPollDetail pairs the plugin's reason with a bounded copy of the
// upstream body so the WARN line is enough to diagnose a parser gap.
func unrecognizedPollDetail(reason string, body []byte) string {
@@ -764,16 +768,18 @@ func unrecognizedPollDetail(reason string, body []byte) string {
func recordPollFailure(ctx context.Context, adaptor TaskPollingAdaptor, task *model.Task, fromStatus model.TaskStatus, class string, statusCode int, detail string) error {
task.PrivateData.PollFailures++
if class == pollClassUnrecognized || class == pollClassHookError {
// The redacted body is intentionally not persisted to Task.Data on these
// paths, so the WARN line is the only operator-visible copy of what the
// plugin could not interpret.
logger.LogWarn(ctx, fmt.Sprintf("task %s poll %s (failures=%d, http=%d): %s", task.TaskID, class, task.PrivateData.PollFailures, statusCode, detail))
}
// TASK_POLL_MAX_FAILURES <= 0 disables the consecutive-failure cutoff, matching
// TASK_TIMEOUT_MINUTES semantics; the 24h sweep remains the only backstop.
if constant.TaskPollMaxFailures > 0 && task.PrivateData.PollFailures >= constant.TaskPollMaxFailures {
return failTaskFromPoll(ctx, adaptor, task, fromStatus, pollFailureReason(class, statusCode, detail))
giveUp := constant.TaskPollMaxFailures > 0 && task.PrivateData.PollFailures >= constant.TaskPollMaxFailures
if giveUp || class == pollClassUnrecognized || class == pollClassHookError {
// The redacted body is intentionally not persisted to Task.Data on these
// paths, and the fail reason of a task that is given up is a fixed
// sentence, so the WARN line is the only operator-visible copy of the
// class, the status and what the plugin could not interpret.
logger.LogWarn(ctx, common.LogText("task %s poll %s (failures=%d, http=%d): %s", task.TaskID, class, task.PrivateData.PollFailures, statusCode, detail))
}
if giveUp {
return failTaskFromPoll(ctx, adaptor, task, fromStatus, "Failed to query the upstream task status")
}
if _, err := task.UpdateWithStatus(fromStatus); err != nil {
return err
@@ -809,11 +815,7 @@ func failTaskFromPoll(ctx context.Context, adaptor TaskPollingAdaptor, task *mod
if !won {
return nil
}
taskResult := relaycommon.FailTaskInfo(reason)
billingSettled := settleTaskBillingOnComplete(ctx, adaptor, task, taskResult)
if !billingSettled && task.Quota != 0 {
RefundTaskQuota(ctx, task, reason)
}
finalizeTerminalTask(ctx, adaptor, task, relaycommon.FailTaskInfo(reason))
return nil
}
+52 -8
View File
@@ -24,6 +24,7 @@ import (
type taskPollingFetchAdaptor struct {
mu sync.Mutex
initInfo *relaycommon.RelayInfo
taskIDs []string
fetched chan string
blockTaskID string
@@ -59,7 +60,19 @@ func (a *batchPollingAdaptor) ParseBatchResult(_ []*model.Task, _ *http.Response
return results, nil
}
func (a *taskPollingFetchAdaptor) Init(_ *relaycommon.RelayInfo) {}
func (a *taskPollingFetchAdaptor) Init(info *relaycommon.RelayInfo) {
a.mu.Lock()
defer a.mu.Unlock()
a.initInfo = info
}
func (a *taskPollingFetchAdaptor) initChannelMeta() *relaycommon.ChannelMeta {
a.mu.Lock()
defer a.mu.Unlock()
if a.initInfo == nil {
return nil
}
return a.initInfo.ChannelMeta
}
func (a *taskPollingFetchAdaptor) FetchTask(_ string, _ string, task *model.Task, _ string) (*http.Response, error) {
taskID := ""
@@ -225,6 +238,37 @@ func TestUpdateVideoTasksDefaultSleepWaitsBetweenTasks(t *testing.T) {
assert.Equal(t, 1, adaptor.fetchCount())
}
func TestPollingPassesExecutingChannelTypeToAdaptor(t *testing.T) {
truncate(t)
const channelID = 112
baseURL := "https://gateway.example"
gateway := &model.Channel{Id: channelID, Type: constant.ChannelTypeNewAPI, Name: "gateway", Key: "sk-gateway", Status: common.ChannelStatusEnabled, BaseURL: &baseURL}
gateway.SetOtherSettings(dto.ChannelOtherSettings{DisableTaskPollingSleep: true})
require.NoError(t, model.DB.Create(gateway).Error)
task := seedPollingTask(t, channelID, "task_gateway", "upstream_gateway")
taskChannels := map[int][]string{channelID: {task.GetUpstreamTaskID()}}
tasks := map[string]*model.Task{task.GetUpstreamTaskID(): task}
perTask := &taskPollingFetchAdaptor{}
previousFactory := GetTaskAdaptorFunc
GetTaskAdaptorFunc = func(constant.TaskPlatform) TaskPollingAdaptor { return perTask }
t.Cleanup(func() { GetTaskAdaptorFunc = previousFactory })
require.NoError(t, UpdateVideoTasks(context.Background(), constant.TaskPlatform("kling"), taskChannels, tasks))
meta := perTask.initChannelMeta()
require.NotNil(t, meta, "per-task polling initializes the adaptor with channel metadata")
assert.Equal(t, constant.ChannelTypeNewAPI, meta.ChannelType, "the adaptor derives the upstream kind from the channel type")
assert.Equal(t, channelID, meta.ChannelId)
assert.Equal(t, baseURL, meta.ChannelBaseUrl)
batch := &batchPollingAdaptor{}
require.NoError(t, UpdateBatchTasks(context.Background(), batch, taskChannels, tasks))
meta = batch.initChannelMeta()
require.NotNil(t, meta, "batch polling initializes the adaptor with channel metadata")
assert.Equal(t, constant.ChannelTypeNewAPI, meta.ChannelType)
assert.Equal(t, channelID, meta.ChannelId)
assert.Equal(t, baseURL, meta.ChannelBaseUrl)
}
func TestDispatchPlatformUpdateUsesFetchMode(t *testing.T) {
truncate(t)
const channelID = 109
@@ -722,8 +766,8 @@ func TestSweepTimedOutTasksHonorsRefundRolloutBoundary(t *testing.T) {
assert.EqualValues(t, model.TaskStatusFailure, reloadedModern.Status)
assert.Zero(t, reloadedLegacy.Quota)
assert.Zero(t, reloadedModern.Quota)
assert.Contains(t, reloadedLegacy.FailReason, "旧系统遗留任务")
assert.Contains(t, reloadedModern.FailReason, "任务超时")
assert.Contains(t, reloadedLegacy.FailReason, "old system")
assert.Equal(t, "Task timed out", reloadedModern.FailReason)
assert.Equal(t, initialQuota+modernTaskQuota, getUserQuota(t, userID))
assert.Equal(t, int64(1), countLogs(t))
}
@@ -802,7 +846,7 @@ func TestUpdateVideoSingleTaskPollClassification(t *testing.T) {
statusCode: http.StatusNotFound,
wantStatus: model.TaskStatusFailure,
wantRefund: true,
wantReason: "upstream task not found (HTTP 404)",
wantReason: "Upstream task not found",
wantUnchanged: false,
},
{
@@ -820,7 +864,7 @@ func TestUpdateVideoSingleTaskPollClassification(t *testing.T) {
wantStatus: model.TaskStatusFailure,
wantFailures: 3,
wantRefund: true,
wantReason: "poll failed: transient (HTTP 429)",
wantReason: "Failed to query the upstream task status",
},
{
name: "UNKNOWN increments",
@@ -904,7 +948,7 @@ func TestUpdateVideoSingleTaskPollClassification(t *testing.T) {
assert.Empty(t, persisted.FailReason)
}
if testCase.wantReason != "" {
assert.Contains(t, persisted.FailReason, testCase.wantReason)
assert.Equal(t, testCase.wantReason, persisted.FailReason)
}
if testCase.wantState != "" {
assert.JSONEq(t, testCase.wantState, string(persisted.PrivateData.PluginState))
@@ -939,7 +983,7 @@ func TestUpdateBatchTasksPollClassification(t *testing.T) {
statusCode: http.StatusNotFound,
wantStatus: model.TaskStatusFailure,
wantRefund: true,
wantReason: "upstream task not found (HTTP 404)",
wantReason: "Upstream task not found",
},
{
name: "401 increments every task",
@@ -986,7 +1030,7 @@ func TestUpdateBatchTasksPollClassification(t *testing.T) {
assert.EqualValues(t, testCase.wantStatus, persisted.Status)
assert.Equal(t, testCase.wantFailures, persisted.PrivateData.PollFailures)
if testCase.wantReason != "" {
assert.Contains(t, persisted.FailReason, testCase.wantReason)
assert.Equal(t, testCase.wantReason, persisted.FailReason)
}
if testCase.wantRefund {
assert.Equal(t, initialQuota+preConsumed, getUserQuota(t, userID))
+37 -32
View File
@@ -1,7 +1,6 @@
package service
import (
"fmt"
"math"
"sort"
"strings"
@@ -68,11 +67,15 @@ type textQuotaSummary struct {
}
// hasBillableUsage reports whether this request should incur any charge.
// A request can carry zero tokens yet still be billable via a tool-call
// Claude reports cache tokens separately from ordinary input and output.
// A request can carry zero ordinary tokens yet still be billable via cache or a tool-call
// surcharge (e.g. /v1/alpha/search returns no usage but bills one web_search
// call), so token count alone is not sufficient to decide.
func (s *textQuotaSummary) hasBillableUsage() bool {
return s.FixedPriceBilling || s.TotalTokens > 0 || !s.ToolCallSurchargeQuota.IsZero()
return s.FixedPriceBilling || s.TotalTokens > 0 ||
(s.IsClaudeUsageSemantic && (s.CacheTokens > 0 || s.CacheCreationTokens > 0 ||
s.CacheCreationTokens5m > 0 || s.CacheCreationTokens1h > 0)) ||
!s.ToolCallSurchargeQuota.IsZero()
}
func cacheWriteTokensTotal(summary textQuotaSummary) int {
@@ -132,7 +135,7 @@ func mergeToolSurchargeItems(items []ToolSurchargeItem) []ToolSurchargeItem {
merged[lastIndex].Name == item.Name &&
merged[lastIndex].Price == item.Price {
if item.Count > math.MaxInt-merged[lastIndex].Count {
common.SysError("tool surcharge call count overflow for " + item.Name)
common.SysError(common.LogText("tool surcharge call count overflow for %s", item.Name))
merged[lastIndex].Count = math.MaxInt
} else {
merged[lastIndex].Count += item.Count
@@ -155,7 +158,12 @@ func calculateTextToolCallSurcharge(ctx *gin.Context, relayInfo *relaycommon.Rel
if tool == nil {
continue
}
items = collectToolSurchargeItem(items, name, tool.CallCount, summary.ModelName)
count := tool.CallCount
if count > relaycommon.MaxBillableToolCallCount {
logger.LogWarn(ctx, common.LogText("tool surcharge call count clamped: tool=%s count=%d", name, count))
count = relaycommon.MaxBillableToolCallCount
}
items = collectToolSurchargeItem(items, name, count, summary.ModelName)
}
}
if relayInfo.RelayMode != relayconstant.RelayModeResponses &&
@@ -163,17 +171,6 @@ func calculateTextToolCallSurcharge(ctx *gin.Context, relayInfo *relaycommon.Rel
items = collectToolSurchargeItem(items, dto.BuildInToolWebSearchPreview, 1, summary.ModelName)
}
items = collectToolSurchargeItem(
items,
dto.BuildInToolWebSearch,
ctx.GetInt("claude_web_search_requests"),
summary.ModelName,
)
if ctx.GetBool("gemini_google_search_call") {
items = collectToolSurchargeItem(items, dto.BuildInToolGoogleSearch, 1, summary.ModelName)
}
summary.ToolSurchargeItems = mergeToolSurchargeItems(items)
var surcharge decimal.Decimal
for _, item := range summary.ToolSurchargeItems {
@@ -390,11 +387,18 @@ func usageSemanticFromUsage(relayInfo *relaycommon.RelayInfo, usage *dto.Usage)
return "openai"
}
func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage, extraContent []string) {
func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage, extraContent []*common.Message) {
originUsage := usage
billingUsage := effectiveBillingUsage(usage)
var cacheReadAdjustment *claudeCacheReadAdjustment
if !common.GetContextKeyBool(ctx, constant.ContextKeyLocalCountTokens) {
billingUsage, cacheReadAdjustment = normalizeClaudeCacheReadUsage(relayInfo, billingUsage)
if cacheReadAdjustment != nil && cacheReadAdjustment.Status == "invalid_usage" {
logger.LogWarn(ctx, "invalid Claude cache-read usage; channel compatibility adjustment skipped")
}
}
if usage == nil {
extraContent = append(extraContent, "上游无计费信息")
extraContent = append(extraContent, common.NewMessage("Upstream returned no usage"))
}
if originUsage != nil && !ctx.GetBool("attempt_checkpoint") {
ObserveChannelAffinityUsageCacheByRelayFormat(ctx, billingUsage, relayInfo.GetFinalRequestRelayFormat())
@@ -437,21 +441,20 @@ func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us
Div(decimal.NewFromInt(1000)).
Mul(decimal.NewFromFloat(summary.GroupRatio)).
Mul(decimal.NewFromFloat(common.QuotaPerUnit))
extraContent = append(extraContent, fmt.Sprintf(
"%s 调用 %d 次,调用花费 %s",
item.Name,
item.Count,
logger.LogQuota(common.QuotaFromDecimal(q)),
))
extraContent = append(extraContent, common.NewMessage("{{name}} called {{count}} times, cost {{quota}}", map[string]any{
"name": item.Name,
"count": item.Count,
"quota": logger.FormatQuota(common.QuotaFromDecimal(q)),
}))
}
if summary.AudioInputPrice > 0 && summary.AudioTokens > 0 {
q := decimal.NewFromFloat(summary.AudioInputPrice).Div(decimal.NewFromInt(1000000)).Mul(decimal.NewFromInt(int64(summary.AudioTokens))).Mul(decimal.NewFromFloat(summary.GroupRatio)).Mul(decimal.NewFromFloat(common.QuotaPerUnit))
extraContent = append(extraContent, fmt.Sprintf("Audio Input 花费 %s", logger.LogQuota(common.QuotaFromDecimal(q))))
extraContent = append(extraContent, common.NewMessage("Audio input cost {{quota}}", map[string]any{"quota": logger.FormatQuota(common.QuotaFromDecimal(q))}))
}
if !summary.hasBillableUsage() {
extraContent = append(extraContent, "上游没有返回计费信息,无法扣费(可能是上游超时)")
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, summary.ModelName, relayInfo.FinalPreConsumedQuota))
extraContent = append(extraContent, common.NewMessage("Upstream returned no usage, so nothing was charged (possibly an upstream timeout)"))
logger.LogError(ctx, common.LogText("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, summary.ModelName, relayInfo.FinalPreConsumedQuota))
} else if relayInfo.DurableAttemptID == 0 {
model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, summary.Quota)
model.UpdateChannelUsedQuota(relayInfo.ChannelId, summary.Quota)
@@ -459,21 +462,20 @@ func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us
if relayInfo.DurableAttemptID == 0 {
if err := SettleBilling(ctx, relayInfo, summary.Quota); err != nil {
logger.LogError(ctx, "error settling billing: "+err.Error())
logger.LogError(ctx, common.LogText("error settling billing: %s", err.Error()))
}
}
logModel := summary.ModelName
if strings.HasPrefix(logModel, "gpt-4-gizmo") {
logModel = "gpt-4-gizmo-*"
extraContent = append(extraContent, fmt.Sprintf("模型 %s", summary.ModelName))
extraContent = append(extraContent, common.NewMessage("Model {{model}}", map[string]any{"model": summary.ModelName}))
}
if strings.HasPrefix(logModel, "gpt-4o-gizmo") {
logModel = "gpt-4o-gizmo-*"
extraContent = append(extraContent, fmt.Sprintf("模型 %s", summary.ModelName))
extraContent = append(extraContent, common.NewMessage("Model {{model}}", map[string]any{"model": summary.ModelName}))
}
logContent := strings.Join(extraContent, ", ")
var other *model.LogOther
if summary.IsClaudeUsageSemantic {
other = GenerateClaudeOtherInfo(ctx, relayInfo,
@@ -488,6 +490,9 @@ func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us
other = GenerateTextOtherInfo(ctx, relayInfo, summary.ModelRatio, summary.GroupRatio, summary.CompletionRatio, summary.CacheTokens, summary.CacheRatio, summary.ModelPrice, relayInfo.PriceData.GroupRatioInfo.GroupSpecialRatio)
}
appendUsageBillingPathForLog(other, common.GetContextKeyBool(ctx, constant.ContextKeyLocalCountTokens), originUsage)
if cacheReadAdjustment != nil {
other.SetAdmin("cache_read_adjustment", cacheReadAdjustment)
}
if adminRejectReason != "" {
other.SetAdmin("reject_reason", adminRejectReason)
}
@@ -542,7 +547,7 @@ func PostTextConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, us
ModelName: logModel,
TokenName: summary.TokenName,
Quota: summary.Quota,
Content: logContent,
Content: extraContent,
TokenId: relayInfo.TokenId,
UseTimeSeconds: int(summary.UseTimeSeconds),
IsStream: relayInfo.IsStream,
+307 -27
View File
@@ -15,6 +15,7 @@ import (
relaycommon "github.com/QuantumNous/new-api/relay/common"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/relayconvert"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/setting/operation_setting"
hosttypes "github.com/QuantumNous/new-api/types"
@@ -90,10 +91,207 @@ func TestFixedPriceBillingDatabaseMatrix(t *testing.T) {
require.NoError(t, db.Raw(versionQuery).Scan(&version).Error)
t.Logf("database: %s", version)
runFixedPriceAccountingCases(t, db, logDB)
t.Run("claude cache read compatibility", func(t *testing.T) {
runClaudeCacheReadAccountingCases(t, db, logDB)
})
})
}
}
func claudeCacheReadTestInfo() *relaycommon.RelayInfo {
return &relaycommon.RelayInfo{
OriginModelName: "deepseek-v4.1-flash", RelayFormat: types.RelayFormatClaude,
ChannelMeta: &relaycommon.ChannelMeta{
ChannelId: 22, UpstreamModelName: "deepseek-v4.1-flash",
ChannelOtherSettings: dto.ChannelOtherSettings{ClaudeInputIncludesCacheReadModels: []string{"deepseek-v4.1-flash"}},
},
PriceData: hosttypes.PriceData{ModelRatio: 1, CompletionRatio: 4, CacheRatio: 0.02, CacheCreationRatio: 1.25,
GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 0.162}},
StartTime: time.Now(),
}
}
func TestClaudeCacheReadCompatibility(t *testing.T) {
for _, tc := range []struct {
name string
prompt, cached, output int
wantQuota int
}{
{"screenshot 1", 275773, 275200, 287, 1170},
{"screenshot 2", 275328, 274944, 2117, 2325},
{"screenshot 3", 275049, 270336, 3926, 4183},
{"screenshot 4", 270431, 270208, 58, 949},
} {
t.Run(tc.name, func(t *testing.T) {
info := claudeCacheReadTestInfo()
original := &dto.Usage{BillingUsage: dto.NewClaudeMessagesBillingUsage(&dto.ClaudeUsage{
InputTokens: tc.prompt, CacheReadInputTokens: tc.cached, OutputTokens: tc.output,
})}
before, err := common.Marshal(original)
require.NoError(t, err)
canonical := effectiveBillingUsage(original)
adjusted, audit := normalizeClaudeCacheReadUsage(info, canonical)
require.NotNil(t, audit)
assert.Equal(t, "applied", audit.Status)
assert.Equal(t, tc.prompt, audit.OriginalInputTokens)
assert.Equal(t, tc.prompt-tc.cached, adjusted.PromptTokens)
assert.Equal(t, tc.prompt, adjusted.InputTokens)
assert.Equal(t, tc.prompt-tc.cached, adjusted.BillingUsage.ClaudeUsage.InputTokens)
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
assert.Equal(t, tc.wantQuota, calculateTextQuotaSummary(ctx, info, adjusted).Quota)
params := BuildTieredTokenParams(adjusted, true, billingexpr.UsedVars("p * 2 + cr * 0.04 + c * 8"))
assert.Equal(t, float64(tc.prompt), params.Len)
assert.Equal(t, float64(tc.prompt-tc.cached), params.P)
cost, _, err := billingexpr.RunExprWithRequest("p * 2 + cr * 0.04 + c * 8", params, billingexpr.RequestInput{})
require.NoError(t, err)
assert.Equal(t, tc.wantQuota, common.QuotaRound(cost/1_000_000*common.QuotaPerUnit*0.162))
withoutSeparateCache := BuildTieredTokenParams(adjusted, true, billingexpr.UsedVars("p * 2 + c * 8"))
assert.Equal(t, float64(tc.prompt), withoutSeparateCache.P)
again, _ := normalizeClaudeCacheReadUsage(info, effectiveBillingUsage(original))
assert.Equal(t, adjusted, again, "re-reading the same upstream snapshot must not accumulate corrections")
after, err := common.Marshal(original)
require.NoError(t, err)
assert.Equal(t, before, after)
assert.Equal(t, tc.prompt, canonical.PromptTokens)
})
}
}
func TestClaudeCacheReadCompatibilityBoundaries(t *testing.T) {
for _, tc := range []struct {
name string
prompt, cached, creation int
estimated, disabled, openAI bool
model, returnedModel, status string
wantPrompt, wantInput int
}{
{name: "disabled standard Anthropic", prompt: 100, cached: 80, disabled: true, wantPrompt: 100, wantInput: 180},
{name: "another model", prompt: 100, cached: 80, model: "deepseek-v4-pro", wantPrompt: 100, wantInput: 180},
{name: "exact names only", prompt: 100, cached: 80, model: "deepseek-v4.1-flash-extra", wantPrompt: 100, wantInput: 180},
{name: "returned model cannot opt in", prompt: 100, cached: 80, model: "other", returnedModel: "deepseek-v4.1-flash", wantPrompt: 100, wantInput: 180},
{name: "returned model cannot opt out", prompt: 100, cached: 80, returnedModel: "other", status: "applied", wantPrompt: 20, wantInput: 100},
{name: "OpenAI snapshot remains authoritative", prompt: 100, cached: 80, openAI: true, wantPrompt: 100, wantInput: 100},
{name: "estimated usage", prompt: 100, cached: 80, estimated: true, wantPrompt: 100, wantInput: 180},
{name: "zero cache", prompt: 100, wantPrompt: 100, wantInput: 100},
{name: "full cache", prompt: 100, cached: 100, status: "applied", wantInput: 100},
{name: "cache writes remain separate", prompt: 100, cached: 80, creation: 10, status: "applied", wantPrompt: 20, wantInput: 110},
{name: "negative prompt", prompt: -1, cached: 80, status: "invalid_usage", wantPrompt: -1, wantInput: 79},
{name: "negative cache", prompt: 100, cached: -1, status: "invalid_usage", wantPrompt: 100, wantInput: 99},
{name: "cache exceeds input", prompt: 100, cached: 101, status: "invalid_usage", wantPrompt: 100, wantInput: 201},
} {
t.Run(tc.name, func(t *testing.T) {
info := claudeCacheReadTestInfo()
if tc.disabled {
info.ChannelOtherSettings.ClaudeInputIncludesCacheReadModels = nil
}
if tc.model != "" {
info.UpstreamModelName = tc.model
}
if tc.returnedModel != "" {
info.ObserveResponseModel(tc.returnedModel)
info.UpstreamModelName = tc.returnedModel
}
usage := &dto.Usage{BillingUsage: dto.NewClaudeMessagesBillingUsage(&dto.ClaudeUsage{
InputTokens: tc.prompt, CacheReadInputTokens: tc.cached, CacheCreationInputTokens: tc.creation, OutputTokens: 1,
})}
if tc.openAI {
usage.BillingUsage = dto.NewOpenAIChatBillingUsage(&dto.Usage{PromptTokens: tc.prompt, CompletionTokens: 1,
PromptTokensDetails: dto.InputTokenDetails{CachedTokens: tc.cached}})
}
usage.BillingUsage.Estimated = tc.estimated
canonical := effectiveBillingUsage(usage)
adjusted, audit := normalizeClaudeCacheReadUsage(info, canonical)
assert.Equal(t, tc.wantPrompt, adjusted.PromptTokens)
assert.Equal(t, tc.wantInput, adjusted.InputTokens)
if tc.status == "" {
assert.Nil(t, audit)
assert.Same(t, canonical, adjusted)
} else {
require.NotNil(t, audit)
assert.Equal(t, tc.status, audit.Status)
}
assert.Equal(t, tc.creation, adjusted.PromptTokensDetails.CacheCreationTokensTotal())
})
}
}
func runClaudeCacheReadAccountingCases(t *testing.T, db, logDB *gorm.DB) {
t.Helper()
for _, stream := range []bool{false, true} {
for _, expression := range []bool{false, true} {
t.Run(fmt.Sprintf("stream=%t expression=%t", stream, expression), func(t *testing.T) {
const initialQuota = 1_000_000
info := claudeCacheReadTestInfo()
channel := model.Channel{Name: "cache-read-test", Key: "unused", Status: common.ChannelStatusEnabled}
channel.SetOtherSettings(info.ChannelOtherSettings)
require.NoError(t, db.Create(&channel).Error)
user := model.User{Username: fmt.Sprintf("cache_read_%t_%t", stream, expression), Quota: initialQuota, Status: common.UserStatusEnabled}
require.NoError(t, db.Create(&user).Error)
token := model.Token{UserId: user.Id, Key: fmt.Sprintf("cache-read-%t-%t", stream, expression), RemainQuota: initialQuota, Status: common.TokenStatusEnabled}
require.NoError(t, db.Create(&token).Error)
t.Cleanup(func() {
require.NoError(t, logDB.Where("user_id = ?", user.Id).Delete(&model.Log{}).Error)
require.NoError(t, db.Unscoped().Delete(&token).Error)
require.NoError(t, db.Unscoped().Delete(&user).Error)
require.NoError(t, db.Delete(&channel).Error)
})
require.NoError(t, db.First(&channel, channel.Id).Error)
settings := channel.GetOtherSettings()
settings.AllowSpeed = true
channel.SetOtherSettings(settings)
require.NoError(t, db.Model(&channel).Update("settings", channel.OtherSettings).Error)
require.NoError(t, db.First(&channel, channel.Id).Error)
assert.True(t, channel.GetOtherSettings().AllowSpeed)
assert.Equal(t, []string{"deepseek-v4.1-flash"}, channel.GetOtherSettings().ClaudeInputIncludesCacheReadModels)
info.ChannelOtherSettings = channel.GetOtherSettings()
info.ChannelId, info.UserId, info.TokenId, info.TokenKey = channel.Id, user.Id, token.Id, token.Key
info.UserGroup, info.UsingGroup = "default", "default"
info.UserSetting.BillingPreference = "wallet_only"
info.ForcePreConsume, info.IsStream = true, stream
if expression {
const expr = `tier("base", p * 2 + cr * 0.04 + c * 8)`
info.TieredBillingSnapshot = &billingexpr.BillingSnapshot{BillingMode: "tiered_expr", ExprString: expr,
ExprHash: billingexpr.ExprHashString(expr), QuotaPerUnit: common.QuotaPerUnit, GroupRatio: 0.162}
}
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
ctx.Request = httptest.NewRequest("POST", "/v1/messages", nil)
require.Nil(t, PreConsumeBilling(ctx, 50_000, info))
wireUsage := &dto.ClaudeUsage{InputTokens: 275328, CacheReadInputTokens: 274944, OutputTokens: 2117}
usage := &dto.Usage{BillingUsage: dto.NewClaudeMessagesBillingUsage(wireUsage)}
if stream {
state := &relayconvert.ClaudeResponseInfo{Usage: &dto.Usage{}}
relayconvert.FormatClaudeResponseInfo(&dto.ClaudeResponse{Type: "message_start", Message: &dto.ClaudeMediaMessage{Usage: &dto.ClaudeUsage{InputTokens: 275328}}}, nil, state)
relayconvert.FormatClaudeResponseInfo(&dto.ClaudeResponse{Type: "message_delta", Usage: wireUsage}, nil, state)
relayconvert.FinalizeClaudeStreamBillingUsage(state)
usage = state.Usage
}
PostTextConsumeQuota(ctx, info, usage, nil)
require.NoError(t, db.First(&user, user.Id).Error)
require.NoError(t, db.First(&token, token.Id).Error)
require.NoError(t, db.First(&channel, channel.Id).Error)
assert.Equal(t, initialQuota-2325, user.Quota)
assert.Equal(t, initialQuota-2325, token.RemainQuota)
assert.Equal(t, 2325, user.UsedQuota)
assert.Equal(t, int64(2325), channel.UsedQuota)
var log model.Log
require.NoError(t, logDB.Where("user_id = ?", user.Id).Take(&log).Error)
assert.Equal(t, 2325, log.Quota)
assert.Equal(t, 384, log.PromptTokens)
var other struct {
Admin struct {
Adjustment claudeCacheReadAdjustment `json:"cache_read_adjustment"`
} `json:"admin_info"`
}
require.NoError(t, common.UnmarshalJsonStr(log.Other, &other))
assert.Equal(t, "applied", other.Admin.Adjustment.Status)
assert.Equal(t, 275328, other.Admin.Adjustment.OriginalInputTokens)
assert.Equal(t, 384, other.Admin.Adjustment.InputTokens)
assert.Equal(t, 275328, wireUsage.InputTokens)
})
}
}
}
func runFixedPriceAccountingCases(t *testing.T, db, logDB *gorm.DB) {
t.Helper()
const mixed = `len <= 32000 ? tier("short", fixed(0.01)) : tier("long", p * 2)`
@@ -193,7 +391,7 @@ func runFixedPriceAccountingCases(t *testing.T, db, logDB *gorm.DB) {
require.NoError(t, err)
assert.Equal(t, quota-reservation, held)
if tc.outboundImages > 0 {
reserveErr := PrepareImageBillingForRequest(ctx, info, tc.outboundImages, false)
reserveErr := PrepareImageBillingForRequest(ctx, info, tc.outboundImages)
if tc.reserveInsufficient {
require.NotNil(t, reserveErr)
assert.Equal(t, types.ErrorCodeInsufficientUserQuota, reserveErr.GetErrorCode())
@@ -230,9 +428,9 @@ func runFixedPriceAccountingCases(t *testing.T, db, logDB *gorm.DB) {
if tc.realtime {
PostWssConsumeQuota(ctx, info, info.OriginModelName, &dto.RealtimeUsage{
InputTokens: tc.usage.PromptTokens, OutputTokens: tc.usage.CompletionTokens, TotalTokens: tc.usage.TotalTokens,
}, "")
}, nil)
} else if tc.audio {
PostAudioConsumeQuota(ctx, info, tc.usage, "")
PostAudioConsumeQuota(ctx, info, tc.usage, nil)
} else {
PostTextConsumeQuota(ctx, info, tc.usage, nil)
}
@@ -241,7 +439,7 @@ func runFixedPriceAccountingCases(t *testing.T, db, logDB *gorm.DB) {
require.NoError(t, logDB.Where("user_id = ?", user.Id).Take(&log).Error)
assert.Equal(t, tc.want, log.Quota)
assert.Equal(t, tc.stream, log.IsStream)
assert.NotContains(t, log.Content, "无法扣费")
assert.NotContains(t, log.Content, "nothing was charged")
var other map[string]any
require.NoError(t, common.UnmarshalJsonStr(log.Other, &other))
assert.Equal(t, string(tc.unit), other["billing_unit"])
@@ -1021,7 +1219,6 @@ func TestComposeTieredTextQuotaFallbackKeepsToolCallSurcharges(t *testing.T) {
gin.SetMode(gin.TestMode)
w := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(w)
ctx.Set("claude_web_search_requests", 2)
relayInfo := &relaycommon.RelayInfo{
OriginModelName: "claude-3-7-sonnet",
@@ -1030,6 +1227,11 @@ func TestComposeTieredTextQuotaFallbackKeepsToolCallSurcharges(t *testing.T) {
CompletionRatio: 1,
GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1.25},
},
ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{
BuiltInTools: map[string]*relaycommon.BuildInToolInfo{
dto.BuildInToolWebSearch: {ToolName: dto.BuildInToolWebSearch, CallCount: 2},
},
},
TieredBillingSnapshot: &billingexpr.BillingSnapshot{
BillingMode: "tiered_expr",
GroupRatio: 1.25,
@@ -1055,7 +1257,6 @@ func TestComposeTieredTextQuotaErrorFallbackUsesPreConsumedQuota(t *testing.T) {
gin.SetMode(gin.TestMode)
w := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(w)
ctx.Set("claude_web_search_requests", 2)
relayInfo := &relaycommon.RelayInfo{
OriginModelName: "claude-3-7-sonnet",
@@ -1064,6 +1265,11 @@ func TestComposeTieredTextQuotaErrorFallbackUsesPreConsumedQuota(t *testing.T) {
CompletionRatio: 1,
GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1.25},
},
ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{
BuiltInTools: map[string]*relaycommon.BuildInToolInfo{
dto.BuildInToolWebSearch: {ToolName: dto.BuildInToolWebSearch, CallCount: 2},
},
},
TieredBillingSnapshot: &billingexpr.BillingSnapshot{
BillingMode: "tiered_expr",
GroupRatio: 1.25,
@@ -1257,29 +1463,30 @@ func TestCalculateTextToolCallSurchargeDoesNotInferSearchForResponses(t *testing
assert.Empty(t, summary.ToolSurchargeItems)
}
func TestCalculateTextToolCallSurchargeMergesSameNameAndPrice(t *testing.T) {
// Every BuiltInTools count is bounded by MaxBillableToolCallCount at
// settlement, whatever recorded it (upstream-reported counts and per-item
// counts alike).
func TestCalculateTextToolCallSurchargeClampsCallCount(t *testing.T) {
gin.SetMode(gin.TestMode)
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
ctx.Set("claude_web_search_requests", 3)
relayInfo := &relaycommon.RelayInfo{
OriginModelName: "claude-3-7-sonnet",
ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{
BuiltInTools: map[string]*relaycommon.BuildInToolInfo{
dto.BuildInToolWebSearch: {CallCount: 2},
for recorded, billed := range map[int]int{5: 5, 10_000: 10_000, 10_005: 10_000} {
relayInfo := &relaycommon.RelayInfo{
OriginModelName: "claude-3-7-sonnet",
ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{
BuiltInTools: map[string]*relaycommon.BuildInToolInfo{
dto.BuildInToolWebSearch: {CallCount: recorded},
},
},
},
}
summary := &textQuotaSummary{ModelName: relayInfo.OriginModelName, GroupRatio: 1}
surcharge := calculateTextToolCallSurcharge(ctx, relayInfo, summary)
assert.Equal(t, []ToolSurchargeItem{{Name: dto.BuildInToolWebSearch, Count: billed, Price: 10}}, summary.ToolSurchargeItems, recorded)
expected := decimal.NewFromFloat(10.0 * float64(billed) / 1000).Mul(decimal.NewFromFloat(common.QuotaPerUnit))
assert.True(t, expected.Equal(surcharge), "recorded %d: got %s want %s", recorded, surcharge, expected)
}
summary := &textQuotaSummary{ModelName: relayInfo.OriginModelName, GroupRatio: 1}
surcharge := calculateTextToolCallSurcharge(ctx, relayInfo, summary)
require.Len(t, summary.ToolSurchargeItems, 1)
assert.Equal(t, dto.BuildInToolWebSearch, summary.ToolSurchargeItems[0].Name)
assert.Equal(t, 5, summary.ToolSurchargeItems[0].Count)
assert.Equal(t, 10.0, summary.ToolSurchargeItems[0].Price)
expected := decimal.NewFromFloat(10.0 * 5 / 1000).Mul(decimal.NewFromFloat(common.QuotaPerUnit))
assert.True(t, expected.Equal(surcharge), "got %s want %s", surcharge, expected)
}
func TestMergeToolSurchargeItemsSaturatesCountOverflow(t *testing.T) {
@@ -1321,6 +1528,40 @@ func TestCalculateTextQuotaSummaryZeroTokensStillBillsToolSurcharge(t *testing.T
assert.Equal(t, expected, summary.Quota)
}
func TestClaudeCacheOnlyUsageRemainsBillable(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, tc := range []struct {
name string
usage dto.ClaudeUsage
wantQuota int
}{
{name: "one-hour cache write", usage: dto.ClaudeUsage{CacheCreationInputTokens: 20000, CacheCreation: &dto.ClaudeCacheCreationUsage{Ephemeral1hInputTokens: 20000}}, wantQuota: 100000},
{name: "five-minute cache write", usage: dto.ClaudeUsage{CacheCreationInputTokens: 100, CacheCreation: &dto.ClaudeCacheCreationUsage{Ephemeral5mInputTokens: 100}}, wantQuota: 313},
{name: "cache read", usage: dto.ClaudeUsage{CacheReadInputTokens: 100}, wantQuota: 25},
{name: "no usage", usage: dto.ClaudeUsage{}, wantQuota: 0},
} {
t.Run(tc.name, func(t *testing.T) {
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
info := &relaycommon.RelayInfo{
RelayFormat: types.RelayFormatClaude,
OriginModelName: "claude-opus-4-8",
PriceData: hosttypes.PriceData{
ModelRatio: 2.5, CompletionRatio: 5, CacheRatio: 0.1,
CacheCreationRatio: 1.25, CacheCreation5mRatio: 1.25, CacheCreation1hRatio: 2,
GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1},
},
}
usage := &dto.Usage{UsageSemantic: dto.BillingUsageSemanticAnthropic, BillingUsage: dto.NewClaudeMessagesBillingUsage(&tc.usage)}
summary := calculateTextQuotaSummary(ctx, info, effectiveBillingUsage(usage))
assert.Zero(t, summary.PromptTokens)
assert.Zero(t, summary.CompletionTokens)
assert.Equal(t, tc.wantQuota, summary.Quota)
assert.Equal(t, tc.wantQuota > 0, summary.hasBillableUsage())
t.Logf("input=%d output=%d quota=%d", summary.PromptTokens, summary.CompletionTokens, summary.Quota)
})
}
}
func TestCalculateTextQuotaSummaryDoesNotApplyRequestMultipliersToToolSurcharge(t *testing.T) {
gin.SetMode(gin.TestMode)
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
@@ -1350,10 +1591,16 @@ func TestCalculateTextQuotaSummaryDoesNotApplyRequestMultipliersToToolSurcharge(
func TestCalculateTextToolCallSurchargeGeminiGoogleSearch(t *testing.T) {
gin.SetMode(gin.TestMode)
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
ctx.Set("gemini_google_search_call", true)
relayInfo := &relaycommon.RelayInfo{OriginModelName: "gemini-2.5-flash"}
summary := &textQuotaSummary{ModelName: "gemini-2.5-flash", GroupRatio: 1}
relayInfo := &relaycommon.RelayInfo{
OriginModelName: "gemini-3.7-flash",
ResponsesUsageInfo: &relaycommon.ResponsesUsageInfo{
BuiltInTools: map[string]*relaycommon.BuildInToolInfo{
dto.BuildInToolGoogleSearch: {ToolName: dto.BuildInToolGoogleSearch, CallCount: 1},
},
},
}
summary := &textQuotaSummary{ModelName: "gemini-3.7-flash", GroupRatio: 1}
surcharge := calculateTextToolCallSurcharge(ctx, relayInfo, summary)
expected := decimal.NewFromFloat(14.0 / 1000).Mul(decimal.NewFromFloat(common.QuotaPerUnit))
@@ -1500,3 +1747,36 @@ func TestAppendToolSurchargeLogInfoWritesOnlyStructuredFields(t *testing.T) {
assert.NotContains(t, fields, "image_generation_call")
assert.NotContains(t, fields, "image_generation_call_price")
}
// A request whose upstream returns no usage is logged without a charge. The
// reason is stored as content_parts, which the web console renders in the
// viewer's language; content keeps the English text.
func TestPostTextConsumeQuotaLogsMissingUsageAsStructuredContent(t *testing.T) {
truncate(t)
const userID, channelID = 62, 62
seedUser(t, userID, 10_000)
seedChannel(t, channelID)
gin.SetMode(gin.TestMode)
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
ctx.Request = httptest.NewRequest("POST", "/v1/chat/completions", nil)
info := &relaycommon.RelayInfo{
UserId: userID,
OriginModelName: "missing-usage-model",
UsingGroup: "default",
StartTime: time.Now(),
ChannelMeta: &relaycommon.ChannelMeta{ChannelId: channelID},
PriceData: hosttypes.PriceData{ModelRatio: 1, CompletionRatio: 1, GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: 1}},
}
PostTextConsumeQuota(ctx, info, &dto.Usage{}, nil)
log := getLastLog(t)
require.NotNil(t, log)
assert.Zero(t, log.Quota)
assert.Equal(t, "Upstream returned no usage, so nothing was charged (possibly an upstream timeout)", log.Content)
var other map[string]any
require.NoError(t, common.UnmarshalJsonStr(log.Other, &other))
assert.Equal(t, []any{
map[string]any{"key": "Upstream returned no usage, so nothing was charged (possibly an upstream timeout)"},
}, other["content_parts"])
}
+1 -1
View File
@@ -63,7 +63,7 @@ func BuildTieredTokenParams(usage *dto.Usage, isClaudeUsageSemantic bool, usedVa
cr -= imgCR
img -= imgCR
} else {
common.SysError("invalid image cache token breakdown; using aggregate cache billing")
common.SysError(common.LogText("invalid image cache token breakdown; using aggregate cache billing"))
}
}
}
+4 -1
View File
@@ -3,6 +3,8 @@ package service
import (
"math"
"math/rand"
"net/http"
"net/http/httptest"
"testing"
"github.com/QuantumNous/new-api/common"
@@ -394,7 +396,8 @@ func TestPrepareTieredBillingForSelectedGroupStartsBillingAfterFreeGroup(t *test
GroupRatioInfo: types.GroupRatioInfo{GroupRatio: 0.20},
},
}
ctx, _ := gin.CreateTestContext(nil)
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
require.Nil(t, PrepareTieredBillingForSelectedGroup(ctx, relayInfo))
require.NotNil(t, relayInfo.Billing)
+23 -169
View File
@@ -15,6 +15,7 @@ import (
constant2 "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/tokenkit"
"github.com/gin-gonic/gin"
)
@@ -24,153 +25,24 @@ func getImageToken(c *gin.Context, fileMeta *types.FileMeta, model string, strea
return 0, fmt.Errorf("image_url_is_nil")
}
// Defaults for 4o/4.1/4.5 family unless overridden below
baseTokens := 85
tileTokens := 170
// Model classification
lowerModel := strings.ToLower(model)
// Special cases from existing behavior
if strings.HasPrefix(lowerModel, "glm-4") {
return 1047, nil
img := tokenkit.Image{Detail: fileMeta.Detail}
// Reading the size may download a URL image; operators can turn that off,
// and the estimate then assumes a typical image size.
measure := constant.GetMediaToken && (stream || constant.GetMediaTokenNotStream)
if !measure || !tokenkit.ImageNeedsSize(model, img.Detail) {
return tokenkit.ImageTokens(model, img), nil
}
// Patch-based models (32x32 patches, capped at 1536, with multiplier)
isPatchBased := false
multiplier := 1.0
switch {
case strings.Contains(lowerModel, "gpt-4.1-mini"):
isPatchBased = true
multiplier = 1.62
case strings.Contains(lowerModel, "gpt-4.1-nano"):
isPatchBased = true
multiplier = 2.46
case strings.HasPrefix(lowerModel, "o4-mini"):
isPatchBased = true
multiplier = 1.72
case strings.HasPrefix(lowerModel, "gpt-5-mini"):
isPatchBased = true
multiplier = 1.62
case strings.HasPrefix(lowerModel, "gpt-5-nano"):
isPatchBased = true
multiplier = 2.46
}
// Tile-based model tokens and bases per doc
if !isPatchBased {
if strings.HasPrefix(lowerModel, "gpt-4o-mini") {
baseTokens = 2833
tileTokens = 5667
} else if strings.HasPrefix(lowerModel, "gpt-5-chat-latest") || (strings.HasPrefix(lowerModel, "gpt-5") && !strings.Contains(lowerModel, "mini") && !strings.Contains(lowerModel, "nano")) {
baseTokens = 70
tileTokens = 140
} else if strings.HasPrefix(lowerModel, "o1") || strings.HasPrefix(lowerModel, "o3") || strings.HasPrefix(lowerModel, "o1-pro") {
baseTokens = 75
tileTokens = 150
} else if strings.Contains(lowerModel, "computer-use-preview") {
baseTokens = 65
tileTokens = 129
} else if strings.Contains(lowerModel, "4.1") || strings.Contains(lowerModel, "4o") || strings.Contains(lowerModel, "4.5") {
baseTokens = 85
tileTokens = 170
}
}
// Respect existing feature flags/short-circuits
if fileMeta.Detail == "low" && !isPatchBased {
return baseTokens, nil
}
// Whether to count image tokens at all
if !constant.GetMediaToken {
return 3 * baseTokens, nil
}
if !constant.GetMediaTokenNotStream && !stream {
return 3 * baseTokens, nil
}
// Normalize detail
if fileMeta.Detail == "auto" || fileMeta.Detail == "" {
fileMeta.Detail = "high"
}
// 使用统一的文件服务获取图片配置
config, format, err := GetImageConfig(c, fileMeta.Source)
if err != nil {
return 0, err
}
if config.Width == 0 || config.Height == 0 {
// not an image, but might be a valid file
if format != "" {
// file type
return 3 * baseTokens, nil
}
return 0, errors.New(fmt.Sprintf("fail to decode image config: %s", fileMeta.GetIdentifier()))
if (config.Width == 0 || config.Height == 0) && format == "" {
return 0, fmt.Errorf("fail to decode image config: %s", fileMeta.GetIdentifier())
}
width := config.Width
height := config.Height
logger.LogDebug(c, "image token input: format=%s, width=%d, height=%d", format, width, height)
if isPatchBased {
// 32x32 patch-based calculation with 1536 cap and model multiplier
ceilDiv := func(a, b int) int { return (a + b - 1) / b }
rawPatchesW := ceilDiv(width, 32)
rawPatchesH := ceilDiv(height, 32)
rawPatches := rawPatchesW * rawPatchesH
if rawPatches > 1536 {
// scale down
area := float64(width * height)
r := math.Sqrt(float64(32*32*1536) / area)
wScaled := float64(width) * r
hScaled := float64(height) * r
// adjust to fit whole number of patches after scaling
adjW := math.Floor(wScaled/32.0) / (wScaled / 32.0)
adjH := math.Floor(hScaled/32.0) / (hScaled / 32.0)
adj := math.Min(adjW, adjH)
if !math.IsNaN(adj) && adj > 0 {
r = r * adj
}
wScaled = float64(width) * r
hScaled = float64(height) * r
patchesW := math.Ceil(wScaled / 32.0)
patchesH := math.Ceil(hScaled / 32.0)
imageTokens := min(int(patchesW*patchesH), 1536)
return common.QuotaRound(float64(imageTokens) * multiplier), nil
}
// below cap
imageTokens := rawPatches
return common.QuotaRound(float64(imageTokens) * multiplier), nil
}
// Tile-based calculation for 4o/4.1/4.5/o1/o3/etc.
// Step 1: fit within 2048x2048 square
maxSide := math.Max(float64(width), float64(height))
fitScale := 1.0
if maxSide > 2048 {
fitScale = maxSide / 2048.0
}
fitW := int(math.Round(float64(width) / fitScale))
fitH := int(math.Round(float64(height) / fitScale))
// Step 2: scale so that shortest side is exactly 768
minSide := math.Min(float64(fitW), float64(fitH))
if minSide == 0 {
return baseTokens, nil
}
shortScale := 768.0 / minSide
finalW := int(math.Round(float64(fitW) * shortScale))
finalH := int(math.Round(float64(fitH) * shortScale))
// Count 512px tiles
tilesW := (finalW + 512 - 1) / 512
tilesH := (finalH + 512 - 1) / 512
tiles := tilesW * tilesH
logger.LogDebug(c, "image token scaled size: width=%d, height=%d, tiles=%d", finalW, finalH, tiles)
return tiles*tileTokens + baseTokens, nil
img.Width, img.Height = config.Width, config.Height
logger.LogDebug(c, "image token input: format=%s, width=%d, height=%d", format, img.Width, img.Height)
return tokenkit.ImageTokens(model, img), nil
}
func EstimateRequestToken(c *gin.Context, meta *types.TokenCountMeta, info *relaycommon.RelayInfo) (int, error) {
@@ -228,7 +100,7 @@ func CountRequestToken(c *gin.Context, meta *types.TokenCountMeta, info *relayco
if meta.TokenType == types.TokenTypeTextNumber {
tkm += utf8.RuneCountInString(meta.CombineText)
} else {
tkm += CountTextToken(meta.CombineText, model)
tkm += tokenkit.Count(model, meta.CombineText)
}
if info.RelayFormat == types.RelayFormatOpenAI {
@@ -279,15 +151,11 @@ func CountRequestToken(c *gin.Context, meta *types.TokenCountMeta, info *relayco
for i, file := range meta.Files {
switch file.FileType {
case types.FileTypeImage:
if common.IsOpenAITextModel(model) {
token, err := getImageToken(c, file, model, info.IsStream)
if err != nil {
return 0, fmt.Errorf("error counting image token, media index[%d], identifier[%s], err: %v", i, file.GetIdentifier(), err)
}
tkm += token
} else {
tkm += 520
token, err := getImageToken(c, file, model, info.IsStream)
if err != nil {
return 0, fmt.Errorf("error counting image token, media index[%d], identifier[%s], err: %v", i, file.GetIdentifier(), err)
}
tkm += token
case types.FileTypeAudio:
tkm += 256
case types.FileTypeVideo:
@@ -309,7 +177,7 @@ func CountTokenRealtime(info *relaycommon.RelayInfo, request dto.RealtimeEvent,
switch request.Type {
case dto.RealtimeEventTypeSessionUpdate:
if request.Session != nil {
msgTokens := CountTextToken(request.Session.Instructions, model)
msgTokens := tokenkit.Count(model, request.Session.Instructions)
textToken += msgTokens
}
case dto.RealtimeEventResponseAudioDelta:
@@ -321,7 +189,7 @@ func CountTokenRealtime(info *relaycommon.RelayInfo, request dto.RealtimeEvent,
audioToken += atk
case dto.RealtimeEventResponseAudioTranscriptionDelta, dto.RealtimeEventResponseFunctionCallArgumentsDelta:
// count text token
tkm := CountTextToken(request.Delta, model)
tkm := tokenkit.Count(model, request.Delta)
textToken += tkm
case dto.RealtimeEventInputAudioBufferAppend:
// count audio token
@@ -336,7 +204,7 @@ func CountTokenRealtime(info *relaycommon.RelayInfo, request dto.RealtimeEvent,
case "message":
for _, content := range request.Item.Content {
if content.Type == "input_text" {
tokens := CountTextToken(content.Text, model)
tokens := tokenkit.Count(model, content.Text)
textToken += tokens
}
}
@@ -360,19 +228,19 @@ func CountTokenRealtime(info *relaycommon.RelayInfo, request dto.RealtimeEvent,
func CountTokenInput(input any, model string) int {
switch v := input.(type) {
case string:
return CountTextToken(v, model)
return tokenkit.Count(model, v)
case []string:
var text strings.Builder
for _, s := range v {
text.WriteString(s)
}
return CountTextToken(text.String(), model)
return tokenkit.Count(model, text.String())
case []any:
var text strings.Builder
for _, item := range v {
text.WriteString(fmt.Sprintf("%v", item))
}
return CountTextToken(text.String(), model)
return tokenkit.Count(model, text.String())
}
return CountTokenInput(fmt.Sprintf("%v", input), model)
}
@@ -400,17 +268,3 @@ func CountAudioTokenOutput(audioBase64 string, audioFormat string) (int, error)
// duration 来自上游返回的音频元数据,饱和转换防止 int 回绕
return common.QuotaFromFloat(duration / 60 * 200 / 0.24), nil
}
// CountTextToken 统计文本的token数量,仅OpenAI模型使用tokenizer,其余模型使用估算
func CountTextToken(text string, model string) int {
if text == "" {
return 0
}
if common.IsOpenAITextModel(model) {
tokenEncoder := getTokenEncoder(model)
return getTokenNum(tokenEncoder, text)
} else {
// 非openai模型,使用tiktoken-go计算没有意义,使用估算节省资源
return EstimateTokenByModel(model, text)
}
}
-230
View File
@@ -1,230 +0,0 @@
package service
import (
"math"
"strings"
"sync"
"unicode"
)
// Provider 定义模型厂商大类
type Provider string
const (
OpenAI Provider = "openai" // 代表 GPT-3.5, GPT-4, GPT-4o
Gemini Provider = "gemini" // 代表 Gemini 1.0, 1.5 Pro/Flash
Claude Provider = "claude" // 代表 Claude 3, 3.5 Sonnet
Unknown Provider = "unknown" // 兜底默认
)
// multipliers 定义不同厂商的计费权重
type multipliers struct {
Word float64 // 英文单词 (每词)
Number float64 // 数字 (每连续数字串)
CJK float64 // 中日韩字符 (每字)
Symbol float64 // 普通标点符号 (每个)
MathSymbol float64 // 数学符号 (∑,∫,∂,√等,每个)
URLDelim float64 // URL分隔符 (/,:,?,&,=,#,%) - tokenizer优化好
AtSign float64 // @符号 - 导致单词切分,消耗较高
Emoji float64 // Emoji表情 (每个)
Newline float64 // 换行符/制表符 (每个)
Space float64 // 空格 (每个)
BasePad int // 基础起步消耗 (Start/End tokens)
}
var (
multipliersMap = map[Provider]multipliers{
Gemini: {
Word: 1.15, Number: 2.8, CJK: 0.68, Symbol: 0.38, MathSymbol: 1.05, URLDelim: 1.2, AtSign: 2.5, Emoji: 1.08, Newline: 1.15, Space: 0.2, BasePad: 0,
},
Claude: {
Word: 1.13, Number: 1.63, CJK: 1.21, Symbol: 0.4, MathSymbol: 4.52, URLDelim: 1.26, AtSign: 2.82, Emoji: 2.6, Newline: 0.89, Space: 0.39, BasePad: 0,
},
OpenAI: {
Word: 1.02, Number: 1.55, CJK: 0.85, Symbol: 0.4, MathSymbol: 2.68, URLDelim: 1.0, AtSign: 2.0, Emoji: 2.12, Newline: 0.5, Space: 0.42, BasePad: 0,
},
}
multipliersLock sync.RWMutex
)
// getMultipliers 根据厂商获取权重配置
func getMultipliers(p Provider) multipliers {
multipliersLock.RLock()
defer multipliersLock.RUnlock()
switch p {
case Gemini:
return multipliersMap[Gemini]
case Claude:
return multipliersMap[Claude]
case OpenAI:
return multipliersMap[OpenAI]
default:
// 默认兜底 (按 OpenAI 的算)
return multipliersMap[OpenAI]
}
}
// EstimateToken 计算 Token 数量
func EstimateToken(provider Provider, text string) int {
m := getMultipliers(provider)
var count float64
// 状态机变量
type WordType int
const (
None WordType = iota
Latin
Number
)
currentWordType := None
for _, r := range text {
// 1. 处理空格和换行符
if unicode.IsSpace(r) {
currentWordType = None
// 换行符和制表符使用Newline权重
if r == '\n' || r == '\t' {
count += m.Newline
} else {
// 普通空格使用Space权重
count += m.Space
}
continue
}
// 2. 处理 CJK (中日韩) - 按字符计费
if isCJK(r) {
currentWordType = None
count += m.CJK
continue
}
// 3. 处理Emoji - 使用专门的Emoji权重
if isEmoji(r) {
currentWordType = None
count += m.Emoji
continue
}
// 4. 处理拉丁字母/数字 (英文单词)
if isLatinOrNumber(r) {
isNum := unicode.IsNumber(r)
newType := Latin
if isNum {
newType = Number
}
// 如果之前不在单词中,或者类型发生变化(字母<->数字),则视为新token
// 注意:对于OpenAI,通常"version 3.5"会切分,"abc123xyz"有时也会切分
// 这里简单起见,字母和数字切换时增加权重
if currentWordType == None || currentWordType != newType {
if newType == Number {
count += m.Number
} else {
count += m.Word
}
currentWordType = newType
}
// 单词中间的字符不额外计费
continue
}
// 5. 处理标点符号/特殊字符 - 按类型使用不同权重
currentWordType = None
if isMathSymbol(r) {
count += m.MathSymbol
} else if r == '@' {
count += m.AtSign
} else if isURLDelim(r) {
count += m.URLDelim
} else {
count += m.Symbol
}
}
// 向上取整并加上基础 padding
return int(math.Ceil(count)) + m.BasePad
}
// 辅助:判断是否为 CJK 字符
func isCJK(r rune) bool {
return unicode.Is(unicode.Han, r) ||
(r >= 0x3040 && r <= 0x30FF) || // 日文
(r >= 0xAC00 && r <= 0xD7A3) // 韩文
}
// 辅助:判断是否为单词主体 (字母或数字)
func isLatinOrNumber(r rune) bool {
return unicode.IsLetter(r) || unicode.IsNumber(r)
}
// 辅助:判断是否为Emoji字符
func isEmoji(r rune) bool {
// Emoji的Unicode范围
// 基本范围:0x1F300-0x1F9FF (Emoticons, Symbols, Pictographs)
// 补充范围:0x2600-0x26FF (Misc Symbols), 0x2700-0x27BF (Dingbats)
// 表情符号:0x1F600-0x1F64F (Emoticons)
// 其他:0x1F900-0x1F9FF (Supplemental Symbols and Pictographs)
return (r >= 0x1F300 && r <= 0x1F9FF) ||
(r >= 0x2600 && r <= 0x26FF) ||
(r >= 0x2700 && r <= 0x27BF) ||
(r >= 0x1F600 && r <= 0x1F64F) ||
(r >= 0x1F900 && r <= 0x1F9FF) ||
(r >= 0x1FA00 && r <= 0x1FAFF) // Symbols and Pictographs Extended-A
}
// 辅助:判断是否为数学符号
func isMathSymbol(r rune) bool {
// 数学运算符和符号
// 基本数学符号:∑ ∫ ∂ √ ∞ ≤ ≥ ≠ ≈ ± × ÷
// 上下标数字:² ³ ¹ ⁴ ⁵ ⁶ ⁷ ⁸ ⁹ ⁰
// 希腊字母等也常用于数学
mathSymbols := "∑∫∂√∞≤≥≠≈±×÷∈∉∋∌⊂⊃⊆⊇∪∩∧∨¬∀∃∄∅∆∇∝∟∠∡∢°′″‴⁺⁻⁼⁽⁾ⁿ₀₁₂₃₄₅₆₇₈₉₊₋₌₍₎²³¹⁴⁵⁶⁷⁸⁹⁰"
for _, m := range mathSymbols {
if r == m {
return true
}
}
// Mathematical Operators (U+2200–U+22FF)
if r >= 0x2200 && r <= 0x22FF {
return true
}
// Supplemental Mathematical Operators (U+2A00–U+2AFF)
if r >= 0x2A00 && r <= 0x2AFF {
return true
}
// Mathematical Alphanumeric Symbols (U+1D400–U+1D7FF)
if r >= 0x1D400 && r <= 0x1D7FF {
return true
}
return false
}
// 辅助:判断是否为URL分隔符(tokenizer对这些优化较好)
func isURLDelim(r rune) bool {
// URL中常见的分隔符,tokenizer通常优化处理
urlDelims := "/:?&=;#%"
for _, d := range urlDelims {
if r == d {
return true
}
}
return false
}
func EstimateTokenByModel(model, text string) int {
// strings.Contains(model, "gpt-4o")
if text == "" {
return 0
}
model = strings.ToLower(model)
if strings.Contains(model, "gemini") {
return EstimateToken(Gemini, text)
} else if strings.Contains(model, "claude") {
return EstimateToken(Claude, text)
} else {
return EstimateToken(OpenAI, text)
}
}
+20 -3
View File
@@ -49,10 +49,17 @@ func PreferredPoolChannel(c *gin.Context, modelName string) (*model.Channel, str
// Snapshot the routing policy before selection. An older in-flight request
// may finish after a manual switch, but its binding must not override the
// new policy on the next request.
policy := RequestRoutingPolicy(c)
var fingerprintPool []model.TokenGroupEntry
if policy.Mode != "smart" {
fingerprintPool = ResolveRequestPool(c, modelName)
}
// Automatic rankings change with live metrics and availability. They must
// not invalidate a healthy session binding; explicit policy edits still do.
policyJSON, err := common.Marshal(struct {
Policy *model.TokenRoutingPolicy
Pool []model.TokenGroupEntry
}{RequestRoutingPolicy(c), ResolveRequestPool(c, modelName)})
}{policy, fingerprintPool})
if err != nil {
return nil, ""
}
@@ -111,7 +118,7 @@ func PreferredPoolChannel(c *gin.Context, modelName string) (*model.Channel, str
recordDirectSelection(c, nil, binding.Channel, entry.Group, modelName, "pool_group_degraded")
return nil, ""
}
if poolGroupSkipped(c, modelName, entry.Group) || poolHigherPriorityRecovered(c, modelName, entry.Group) {
if poolGroupSkipped(c, modelName, entry.Group) || (policy.Mode != "smart" && poolHigherPriorityRecovered(c, modelName, entry.Group)) {
ClearPoolAffinity(c)
return nil, ""
}
@@ -174,7 +181,17 @@ func ResolveRequestPool(c *gin.Context, modelName string) []model.TokenGroupEntr
}
userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup)
if policy.Mode == "smart" {
return WithImageGenerationGroup(c, modelName, resolveAutomaticPool(c, modelName))
entries := WithImageGenerationGroup(c, modelName, resolveAutomaticPool(c, modelName))
if len(policy.ExcludedGroups) == 0 {
return entries
}
filtered := make([]model.TokenGroupEntry, 0, len(entries))
for _, entry := range entries {
if !policy.ExcludesGroup(entry.Group) {
filtered = append(filtered, entry)
}
}
return filtered
}
inherited := GetUserAutoGroup(userGroup)
if policy.Mode == "user" {
+116 -2
View File
@@ -8,14 +8,17 @@ import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
"github.com/QuantumNous/new-api/setting"
"github.com/QuantumNous/new-api/setting/operation_setting"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestPoolAffinityManualRoutingSwitch(t *testing.T) {
for _, change := range []string{"priority", "model_rule"} {
for _, change := range []string{"priority", "model_rule", "blacklist"} {
t.Run(change, func(t *testing.T) {
db := setupChannelSelectAutoGroupsTest(t)
const modelName = "manual-switch-model"
@@ -45,10 +48,17 @@ func TestPoolAffinityManualRoutingSwitch(t *testing.T) {
c.Set("token_id", 100)
c.Set("original_model", modelName)
policy := &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}}
if change == "blacklist" {
policy.Mode = "smart"
policy.Groups = nil
automaticMetricCache.Set(modelName, perfmetrics.QueryResult{})
}
if switched {
if change == "priority" {
policy.Groups[0].Priority = 1
policy.Groups[1].Priority = 0
} else if change == "blacklist" {
policy.ExcludedGroups = []string{"default"}
} else {
policy.ModelRules = []model.ModelRoutingRule{{Model: modelName, Groups: []model.TokenGroupEntry{{Group: "vip", Priority: 0}}}}
}
@@ -76,7 +86,7 @@ func TestPoolAffinityManualRoutingSwitch(t *testing.T) {
log := model.Log{UserId: 100, TokenId: 100, ModelName: modelName, CorrelationKey: old.GetString("log_correlation_key"), Type: model.LogTypeConsume, CreatedAt: time.Now().Unix(), ChannelId: 6201, Group: "default", Other: string(body)}
require.NoError(t, db.Create(&log).Error)
// The old group is still eligible: only its priority/preference changed.
// The edited policy must take precedence over the existing binding.
switched := ctx(true)
selected, group, selectErr := SelectChannelForRequest(switched, modelName, &RetryParam{Ctx: switched, TokenGroup: "default", ModelName: modelName, Retry: common.GetPointer(0)})
require.Nil(t, selectErr)
@@ -114,3 +124,107 @@ func TestPoolAffinityManualRoutingSwitch(t *testing.T) {
})
}
}
func TestAutomaticPoolKeepsHealthySession(t *testing.T) {
for _, scenario := range []string{"cheaper_recovered", "ranking_changed", "candidate_added", "degraded", "cooldown", "permission_removed", "channel_disabled", "policy_changed", "new_session", "session_off", "custom_recovery"} {
t.Run(scenario, func(t *testing.T) {
db := setupChannelSelectAutoGroupsTest(t)
const modelName = "smart-sticky-model"
createChannelSelectAutoGroupsChannel(t, db, 6301, "default", modelName)
createChannelSelectAutoGroupsChannel(t, db, 6302, "vip", modelName)
model.InitChannelCache()
oldRedis, oldLogDB := common.RedisEnabled, model.LOG_DB
common.RedisEnabled, model.LOG_DB = false, nil
config := operation_setting.GetChannelAffinitySetting()
oldConfig := *config
*config = operation_setting.ChannelAffinitySetting{Enabled: true, DefaultTTLSeconds: 3600, SessionMode: "prefer", Rules: []operation_setting.ChannelAffinityRule{{Name: "smart-sticky", ModelRegex: []string{".*"}, KeySources: []operation_setting.ChannelAffinityKeySource{{Type: "request_header", Key: "Session-Id"}}, SessionMode: "prefer"}}}
require.NoError(t, poolBindings().Purge())
require.NoError(t, poolCooldowns().Purge())
automaticMetricCache.Set(modelName, perfmetrics.QueryResult{})
t.Cleanup(func() {
*config = oldConfig
common.RedisEnabled, model.LOG_DB = oldRedis, oldLogDB
_ = poolBindings().Purge()
_ = poolCooldowns().Purge()
automaticMetricCache.Purge()
})
ctx := func() *gin.Context {
c := setupGroupPoolContext(t)
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
c.Request.Header.Set("Session-Id", "same-session")
c.Set("id", 101)
c.Set("token_id", 102)
c.Set("original_model", modelName)
policy := &model.TokenRoutingPolicy{Version: 1, Mode: "smart", AutoPreference: "cost", SkipDegradedNonAstra: true}
if scenario == "custom_recovery" {
policy.Mode = "custom"
policy.Groups = []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}
}
c.Set(model.TokenRoutingContextKey, policy)
c.Set("pool_group_fingerprints", map[string]AstraSamplingHistory{})
c.Set("pool_group_visions", map[string]DetectionHistory{})
return c
}
// A temporary failure sends the session to the more expensive group.
first := ctx()
cheapKey := poolCooldownKey(first, modelName, "default")
require.NoError(t, poolCooldowns().SetWithTTL(cheapKey, poolCooldown{Until: time.Now().Add(time.Minute).UnixMilli()}, 0))
if scenario == "candidate_added" {
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"vip":"VIP"}`))
}
ch, group, selectErr := SelectChannelForRequest(first, modelName, &RetryParam{Ctx: first, ModelName: modelName, Retry: common.GetPointer(0)})
require.Nil(t, selectErr)
require.NotNil(t, ch)
require.Equal(t, "vip", group)
first.Set("channel_id", ch.Id)
RequestPolicy(first).Successful = true
RecordPoolAffinity(first)
require.NoError(t, poolCooldowns().SetWithTTL(cheapKey, poolCooldown{Until: time.Now().Add(-time.Second).UnixMilli()}, 0))
next := ctx()
want := "vip"
switch scenario {
case "ranking_changed":
// The initial pool already placed default first. Reverse it to prove ranking changes are irrelevant.
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":4,"vip":3}`))
case "candidate_added":
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`))
case "degraded":
next.Set("pool_group_visions", map[string]DetectionHistory{"vip": {Quality: "degraded"}})
want = "default"
case "cooldown":
require.NoError(t, poolCooldowns().SetWithTTL(poolCooldownKey(next, modelName, "vip"), poolCooldown{Until: time.Now().Add(time.Minute).UnixMilli()}, 0))
want = "default"
case "permission_removed":
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default"}`))
want = "default"
case "channel_disabled":
require.NoError(t, db.Model(&model.Channel{}).Where("id = ?", 6302).Update("status", common.ChannelStatusAutoDisabled).Error)
model.InitChannelCache()
want = "default"
case "policy_changed":
RequestRoutingPolicy(next).AutoPreference = "stable"
want = "default"
case "new_session":
next.Request.Header.Set("Session-Id", "new-session")
want = "default"
case "session_off":
config.Rules[0].SessionMode = "off"
config.SessionMode = "off"
want = "default"
case "custom_recovery":
want = "default"
}
selected, selectedGroup, selectErr := SelectChannelForRequest(next, modelName, &RetryParam{Ctx: next, ModelName: modelName, Retry: common.GetPointer(0)})
require.Nil(t, selectErr)
require.NotNil(t, selected)
assert.Equal(t, want, selectedGroup)
if want == "vip" {
assert.Equal(t, 6302, selected.Id)
state, found, err := poolCooldowns().Get(cheapKey)
require.NoError(t, err)
require.True(t, found)
assert.Empty(t, state.ProbeID, "keeping a healthy route must not claim a cheaper group's recovery probe")
}
})
}
}
+173 -5
View File
@@ -6,13 +6,16 @@ import (
"fmt"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/model"
perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
apiTypes "github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/setting/operation_setting"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"math"
"net/http"
"net/http/httptest"
"slices"
"testing"
@@ -80,8 +83,8 @@ func TestTokenRoutingSingleChannelRetryCooldownAndRecovery(t *testing.T) {
}
func TestAutomaticGroupPreferenceRanking(t *testing.T) {
original := []automaticCandidate{{group: "cheap", cost: .1, ttft: 900, success: 90, ratio: .2}, {group: "fast", cost: .2, ttft: 100, success: 95, ratio: .3}, {group: "stable", cost: .3, ttft: 400, success: 99, ratio: .4}, {group: "unknown", cost: math.Inf(1), ttft: math.Inf(1), success: -1, ratio: .01}}
for preference, want := range map[string]string{"cost": "unknown", "speed": "fast", "stable": "stable"} {
original := []automaticCandidate{{group: "cheap", cost: .1, ttft: 900, throughput: 20, success: 90, ratio: .2}, {group: "fast", cost: .2, ttft: 100, throughput: 30, success: 95, ratio: .3}, {group: "stable", cost: .3, ttft: 400, throughput: 80, success: 99, ratio: .4}, {group: "unknown", cost: math.Inf(1), ttft: math.Inf(1), success: -1, ratio: .01}}
for preference, want := range map[string]string{"cost": "unknown", "speed": "fast", "throughput": "stable", "stable": "stable"} {
t.Run(preference, func(t *testing.T) {
items := slices.Clone(original)
sortAutomaticCandidates(items, preference)
@@ -93,6 +96,45 @@ func TestAutomaticGroupPreferenceRanking(t *testing.T) {
}
}
func TestAutomaticThroughputUsesRequestedModelCallsAndTestFallback(t *testing.T) {
db := setupChannelSelectAutoGroupsTest(t)
require.NoError(t, db.AutoMigrate(&model.Log{}))
previousLogs := model.LOG_DB
model.LOG_DB = db
t.Cleanup(func() { model.LOG_DB = previousLogs; groupModelPerformanceCache.Purge(); automaticMetricCache.Purge() })
for _, group := range []string{"default", "vip"} {
id := 6951
if group == "vip" {
id = 6952
}
createChannelSelectAutoGroupsChannel(t, db, id, group, "throughput-model")
}
model.InitChannelCache()
logs := []model.Log{
{Type: model.LogTypeConsume, TokenId: 1, Group: "default", ModelName: "throughput-model", IsStream: true, UseTime: 2, CompletionTokens: 20, Other: `{"frt":100}`},
{Type: model.LogTypeConsume, ChannelId: 6952, ModelName: "throughput-model", Content: "Model test", IsStream: true, UseTime: 2, CompletionTokens: 100, Other: `{"frt":500}`},
{Type: model.LogTypeConsume, TokenId: 1, Group: "default", ModelName: "another-model", IsStream: true, UseTime: 1, CompletionTokens: 999, Other: `{"frt":1}`},
}
require.NoError(t, db.Create(&logs).Error)
automaticMetricCache.Set("throughput-model", perfmetrics.QueryResult{})
makeCtx := func() *gin.Context {
c := setupGroupPoolContext(t)
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "smart", AutoPreference: "throughput"})
return c
}
c := makeCtx()
assert.Equal(t, []model.TokenGroupEntry{{Group: "vip", Priority: 0}, {Group: "default", Priority: 1}}, ResolveRequestPool(c, "throughput-model"))
// New traffic is reflected after snapshot refresh, but never reorders retries
// of an already-ranked request.
require.NoError(t, db.Create(&model.Log{Type: model.LogTypeConsume, TokenId: 1, Group: "vip", ModelName: "throughput-model", IsStream: true, UseTime: 10, CompletionTokens: 5, Other: `{"frt":100}`}).Error)
groupModelPerformanceCache.Purge()
assert.Equal(t, "vip", ResolveRequestPool(c, "throughput-model")[0].Group)
assert.Equal(t, "default", ResolveRequestPool(makeCtx(), "throughput-model")[0].Group)
policy := &model.TokenRoutingPolicy{Version: 1, Mode: "smart", AutoPreference: "throughput"}
require.NoError(t, policy.Validate(5, func(string) bool { return true }))
}
func TestAutomaticGroupUsesAllAuthorizedModelGroups(t *testing.T) {
db := setupChannelSelectAutoGroupsTest(t)
for _, fixture := range []struct {
@@ -132,13 +174,13 @@ func TestAutomaticGroupUsesAllAuthorizedModelGroups(t *testing.T) {
}
func TestPoolFailoverRespectsRetrySafetyBoundaries(t *testing.T) {
for _, code := range []int{400, 401, 429, 500, 502, 503, 504, 524} {
for _, code := range []int{400, 401, 403, 429, 500, 502, 503, 504, 524} {
t.Run(fmt.Sprint(code), func(t *testing.T) {
c := setupGroupPoolContext(t)
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "smart"})
err := apiTypes.NewErrorWithStatusCode(errors.New("test upstream"), apiTypes.ErrorCodeBadResponseStatusCode, code)
require.Equal(t, code == 500 || code == 502 || code == 503, IsPoolFailoverError(c, err))
err := apiTypes.NewOpenAIError(errors.New("test upstream"), apiTypes.ErrorCodeBadResponseStatusCode, code)
require.Equal(t, code == 403 || code == 500 || code == 502 || code == 503, IsPoolFailoverError(c, err))
ctx, cancel := context.WithCancel(c.Request.Context())
cancel()
c.Request = c.Request.WithContext(ctx)
@@ -211,3 +253,129 @@ func TestTokenRoutingPoolPreferenceExclusionAndAffinityIsolation(t *testing.T) {
require.NoError(t, err2)
require.Nil(t, last)
}
func TestPoolUpstreamForbiddenFailover(t *testing.T) {
for _, mode := range []string{"smart", "custom"} {
t.Run(mode, func(t *testing.T) {
db := setupChannelSelectAutoGroupsTest(t)
const modelName = "forbidden-pool-model"
createChannelSelectAutoGroupsChannel(t, db, 6931, "default", modelName)
createChannelSelectAutoGroupsChannel(t, db, 6932, "vip", modelName)
model.InitChannelCache()
config := operation_setting.GetChannelAffinitySetting()
oldConfig := *config
oldRedis, oldLogDB := common.RedisEnabled, model.LOG_DB
common.RedisEnabled, model.LOG_DB = false, nil
*config = operation_setting.ChannelAffinitySetting{Enabled: true, DefaultTTLSeconds: 3600, SessionMode: "strict", Rules: []operation_setting.ChannelAffinityRule{{Name: "forbidden-pool", ModelRegex: []string{".*"}, KeySources: []operation_setting.ChannelAffinityKeySource{{Type: "request_header", Key: "Session-Id"}}, SessionMode: "strict"}}}
automaticMetricCache.Set(modelName, perfmetrics.QueryResult{})
require.NoError(t, poolBindings().Purge())
require.NoError(t, poolCooldowns().Purge())
t.Cleanup(func() {
_ = poolBindings().Purge()
_ = poolCooldowns().Purge()
automaticMetricCache.Purge()
*config = oldConfig
common.RedisEnabled, model.LOG_DB = oldRedis, oldLogDB
})
makeCtx := func() *gin.Context {
c := setupGroupPoolContext(t)
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
c.Request.Header.Set("Session-Id", "forbidden-session")
common.SetContextKey(c, constant.ContextKeyTokenGroup, "auto")
c.Set("id", 103)
c.Set("token_id", 104)
c.Set("original_model", modelName)
policy := &model.TokenRoutingPolicy{Version: 1, Mode: mode, AutoPreference: "cost", AllowDegraded: true}
if mode == "custom" {
policy.Groups = []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}
}
c.Set(model.TokenRoutingContextKey, policy)
return c
}
// Establish the same successful session binding used by real requests.
first := makeCtx()
channel, group, selectionErr := SelectChannelForRequest(first, modelName, &RetryParam{Ctx: first, ModelName: modelName, Retry: common.GetPointer(0)})
require.Nil(t, selectionErr)
require.NotNil(t, channel)
require.Equal(t, "default", group)
first.Set("channel_id", channel.Id)
RequestPolicy(first).Successful = true
RecordPoolAffinity(first)
current := makeCtx()
param := &RetryParam{Ctx: current, ModelName: modelName, Retry: common.GetPointer(0)}
channel, group, selectionErr = SelectChannelForRequest(current, modelName, param)
require.Nil(t, selectionErr)
require.Equal(t, "default", group)
require.True(t, ShouldSkipRetryAfterChannelAffinityFailure(current))
budget := PoolRetryBudget(current, modelName)
for attempt := range 2 {
require.Equal(t, 6931, channel.Id)
AppendUsedChannel(current, channel.Id)
response := httptest.NewRecorder()
response.WriteHeader(http.StatusForbidden)
_, err := response.WriteString(`{"error":{"message":"upstream access denied","type":"permission_error","code":"forbidden"}}`)
require.NoError(t, err)
upstreamErr := RelayErrorHandler(current, response.Result(), false)
require.True(t, IsPoolFailoverError(current, upstreamErr))
RecordPoolRetryableFailure(current, modelName, group)
require.Equal(t, "retry", DecideRelayRetry(current, upstreamErr, budget-attempt).Action)
param.IncreaseRetry()
channel, group, selectionErr = SelectChannelForRequest(current, modelName, param)
require.Nil(t, selectionErr)
require.NotNil(t, channel)
}
assert.Equal(t, "vip", group)
assert.Equal(t, 6932, channel.Id)
current.Set("channel_id", channel.Id)
RequestPolicy(current).Successful = true
RecordPoolAffinity(current)
next := makeCtx()
bound, boundGroup := PreferredPoolChannel(next, modelName)
require.NotNil(t, bound)
assert.Equal(t, "vip", boundGroup)
assert.Equal(t, 6932, bound.Id)
assert.True(t, poolGroupSkipped(next, modelName, "default"))
})
}
}
func TestPoolForbiddenRetryBoundaries(t *testing.T) {
for _, scenario := range []string{"local", "skip_retry", "disabled_status", "pinned", "response_started", "cancelled", "no_policy", "budget_exhausted"} {
t.Run(scenario, func(t *testing.T) {
c := setupGroupPoolContext(t)
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
common.SetContextKey(c, constant.ContextKeyTokenGroup, "auto")
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "smart"})
c.Set(ginKeyChannelAffinitySkipRetry, true)
upstreamErr := apiTypes.NewOpenAIError(errors.New("upstream forbidden"), apiTypes.ErrorCodeBadResponseStatusCode, http.StatusForbidden)
switch scenario {
case "local":
upstreamErr = apiTypes.NewErrorWithStatusCode(errors.New("local access denied"), apiTypes.ErrorCodeAccessDenied, http.StatusForbidden)
case "skip_retry":
upstreamErr = apiTypes.NewOpenAIError(errors.New("upstream forbidden"), apiTypes.ErrorCodeBadResponseStatusCode, http.StatusForbidden, apiTypes.ErrOptionWithSkipRetry())
case "disabled_status":
previous := operation_setting.AutomaticRetryStatusCodesToString()
require.NoError(t, operation_setting.AutomaticRetryStatusCodesFromString("500-503"))
t.Cleanup(func() { require.NoError(t, operation_setting.AutomaticRetryStatusCodesFromString(previous)) })
case "pinned":
GetChannelConstraints(c).AddPin(dto.ChannelPin{ChannelId: 6931, Source: dto.PinSourceToken, RetryMode: dto.PinRetrySingleAttempt})
case "response_started":
_, err := c.Writer.WriteString("already streaming")
require.NoError(t, err)
case "cancelled":
ctx, cancel := context.WithCancel(c.Request.Context())
cancel()
c.Request = c.Request.WithContext(ctx)
case "no_policy":
c.Set(model.TokenRoutingContextKey, (*model.TokenRoutingPolicy)(nil))
}
assert.Equal(t, scenario == "budget_exhausted", IsPoolFailoverError(c, upstreamErr))
remaining := 3
if scenario == "budget_exhausted" {
remaining = 0
}
assert.Equal(t, "stop", DecideRelayRetry(c, upstreamErr, remaining).Action)
})
}
}
-63
View File
@@ -1,63 +0,0 @@
package service
import (
"sync"
"github.com/QuantumNous/new-api/common"
"github.com/tiktoken-go/tokenizer"
"github.com/tiktoken-go/tokenizer/codec"
)
// tokenEncoderMap won't grow after initialization
var defaultTokenEncoder tokenizer.Codec
// tokenEncoderMap is used to store token encoders for different models
var tokenEncoderMap = make(map[string]tokenizer.Codec)
// tokenEncoderMutex protects tokenEncoderMap for concurrent access
var tokenEncoderMutex sync.RWMutex
func InitTokenEncoders() {
common.SysLog("initializing token encoders")
defaultTokenEncoder = codec.NewCl100kBase()
common.SysLog("token encoders initialized")
}
func getTokenEncoder(model string) tokenizer.Codec {
// First, try to get the encoder from cache with read lock
tokenEncoderMutex.RLock()
if encoder, exists := tokenEncoderMap[model]; exists {
tokenEncoderMutex.RUnlock()
return encoder
}
tokenEncoderMutex.RUnlock()
// If not in cache, create new encoder with write lock
tokenEncoderMutex.Lock()
defer tokenEncoderMutex.Unlock()
// Double-check if another goroutine already created the encoder
if encoder, exists := tokenEncoderMap[model]; exists {
return encoder
}
// Create new encoder
modelCodec, err := tokenizer.ForModel(tokenizer.Model(model))
if err != nil {
// Cache the default encoder for this model to avoid repeated failures
tokenEncoderMap[model] = defaultTokenEncoder
return defaultTokenEncoder
}
// Cache the new encoder
tokenEncoderMap[model] = modelCodec
return modelCodec
}
func getTokenNum(tokenEncoder tokenizer.Codec, text string) int {
if text == "" {
return 0
}
tkm, _ := tokenEncoder.Count(text)
return tkm
}
+2 -1
View File
@@ -4,6 +4,7 @@ import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/tokenkit"
"github.com/gin-gonic/gin"
)
@@ -23,7 +24,7 @@ func ResponseText2Usage(c *gin.Context, responseText string, modeName string, pr
common.SetContextKey(c, constant.ContextKeyLocalCountTokens, true)
usage := &dto.Usage{}
usage.PromptTokens = promptTokens
usage.CompletionTokens = EstimateTokenByModel(modeName, responseText)
usage.CompletionTokens = tokenkit.Estimate(modeName, responseText)
usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
return usage
}
+37 -19
View File
@@ -13,38 +13,56 @@ import (
"github.com/QuantumNous/new-api/setting/system_setting"
)
func NotifyRootUser(t string, subject string, content string) {
user := model.GetRootUser().ToBaseUser()
err := NotifyUser(user.Id, user.Email, user.GetSetting(), dto.NewNotify(t, subject, content, nil))
if err != nil {
common.SysLog(fmt.Sprintf("failed to notify root user: %s", err.Error()))
}
// NotifyChannelWatchers sends a channel notice (a channel disabled or enabled
// automatically, or a test of all channels finished) to every administrator
// who receives channel notices. Until they choose, the root user receives them
// and other administrators do not. render builds the subject and content in
// the recipient's saved language; an empty language means DEFAULT_LANGUAGE.
func NotifyChannelWatchers(notifyType string, render func(lang string) (subject string, content string)) {
notifyAdmins(notifyType, func(role int, setting dto.UserSetting) bool {
if setting.ChannelNotifyEnabled != nil {
return *setting.ChannelNotifyEnabled
}
return role == common.RoleRootUser
}, render)
}
func NotifyUpstreamModelUpdateWatchers(subject string, content string) {
// NotifyUpstreamModelUpdateWatchers sends the upstream model update notice to
// every administrator who enabled it. render builds the subject and content in
// the recipient's saved language; an empty language means DEFAULT_LANGUAGE.
func NotifyUpstreamModelUpdateWatchers(render func(lang string) (subject string, content string)) {
sentCount := notifyAdmins(dto.NotifyTypeChannelUpdate, func(_ int, setting dto.UserSetting) bool {
return setting.UpstreamModelUpdateNotifyEnabled
}, render)
common.SysLog(common.LogText("upstream model update notifications sent: %d", sentCount))
}
// notifyAdmins sends a notice to each enabled administrator for whom receives
// returns true, and returns how many notices it sent.
func notifyAdmins(notifyType string, receives func(role int, setting dto.UserSetting) bool, render func(lang string) (subject string, content string)) int {
var users []model.User
if err := model.DB.
Select("id", "email", "role", "status", "setting").
Where("status = ? AND role >= ?", common.UserStatusEnabled, common.RoleAdminUser).
Find(&users).Error; err != nil {
common.SysLog(fmt.Sprintf("failed to query upstream update notification users: %s", err.Error()))
return
common.SysLog(common.LogText("failed to query administrators to notify: %s", err.Error()))
return 0
}
notification := dto.NewNotify(dto.NotifyTypeChannelUpdate, subject, content, nil)
sentCount := 0
for _, user := range users {
userSetting := user.GetSetting()
if !userSetting.UpstreamModelUpdateNotifyEnabled {
if !receives(user.Role, userSetting) {
continue
}
if err := NotifyUser(user.Id, user.Email, userSetting, notification); err != nil {
common.SysLog(fmt.Sprintf("failed to notify user %d for upstream model update: %s", user.Id, err.Error()))
subject, content := render(userSetting.Language)
if err := NotifyUser(user.Id, user.Email, userSetting, dto.NewNotify(notifyType, subject, content, nil)); err != nil {
common.SysLog(common.LogText("failed to send notice %s to user %d: %s", notifyType, user.Id, err.Error()))
continue
}
sentCount++
}
common.SysLog(fmt.Sprintf("upstream model update notifications sent: %d", sentCount))
return sentCount
}
func NotifyUser(userId int, userEmail string, userSetting dto.UserSetting, data dto.Notify) error {
@@ -56,7 +74,7 @@ func NotifyUser(userId int, userEmail string, userSetting dto.UserSetting, data
// Check notification limit
canSend, err := CheckNotificationLimit(userId, data.Type)
if err != nil {
common.SysLog(fmt.Sprintf("failed to check notification limit: %s", err.Error()))
common.SysLog(common.LogText("failed to check notification limit: %s", err.Error()))
return err
}
if !canSend {
@@ -71,14 +89,14 @@ func NotifyUser(userId int, userEmail string, userSetting dto.UserSetting, data
emailToUse = userEmail
}
if emailToUse == "" {
common.SysLog(fmt.Sprintf("user %d has no email, skip sending email", userId))
common.SysLog(common.LogText("user %d has no email, skip sending email", userId))
return nil
}
return sendEmailNotify(emailToUse, data)
case dto.NotifyTypeWebhook:
webhookURLStr := userSetting.WebhookUrl
if webhookURLStr == "" {
common.SysLog(fmt.Sprintf("user %d has no webhook url, skip sending webhook", userId))
common.SysLog(common.LogText("user %d has no webhook url, skip sending webhook", userId))
return nil
}
@@ -88,7 +106,7 @@ func NotifyUser(userId int, userEmail string, userSetting dto.UserSetting, data
case dto.NotifyTypeBark:
barkURL := userSetting.BarkUrl
if barkURL == "" {
common.SysLog(fmt.Sprintf("user %d has no bark url, skip sending bark", userId))
common.SysLog(common.LogText("user %d has no bark url, skip sending bark", userId))
return nil
}
return sendBarkNotify(barkURL, data)
@@ -96,7 +114,7 @@ func NotifyUser(userId int, userEmail string, userSetting dto.UserSetting, data
gotifyUrl := userSetting.GotifyUrl
gotifyToken := userSetting.GotifyToken
if gotifyUrl == "" || gotifyToken == "" {
common.SysLog(fmt.Sprintf("user %d has no gotify url or token, skip sending gotify", userId))
common.SysLog(common.LogText("user %d has no gotify url or token, skip sending gotify", userId))
return nil
}
return sendGotifyNotify(gotifyUrl, gotifyToken, userSetting.GotifyPriority, data)
+191
View File
@@ -0,0 +1,191 @@
package service
import (
"fmt"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/setting/system_setting"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)
// A low-balance reminder goes out once, on the charge that takes a wallet or
// subscription balance below the threshold. A balance already below it earns
// nothing more until a top-up lifts it back above the line. The site switch
// and the user's switch both turn reminders off; a user who never chose gets
// them.
func TestQuotaReminderFiresOncePerDrop(t *testing.T) {
previousEnabled, previousThreshold := common.QuotaRemindEnabled, common.QuotaRemindThreshold
t.Cleanup(func() {
common.QuotaRemindEnabled, common.QuotaRemindThreshold = previousEnabled, previousThreshold
})
common.QuotaRemindThreshold = 1000
on, off := true, false
// changes move the balance from start: negative ones are charges, which
// may remind; positive ones are top-ups. want holds one entry per charge.
for _, tc := range []struct {
name string
siteEnabled bool
setting dto.UserSetting
start int64
changes []int64
want []bool
}{
{
name: "never chosen: one reminder per drop, again after a top-up",
siteEnabled: true,
start: 1500,
changes: []int64{-400, -200, -300, +1500, -1200},
want: []bool{false, true, false, true},
},
{
name: "switched on with its own threshold",
siteEnabled: true,
setting: dto.UserSetting{QuotaWarningEnabled: &on, QuotaWarningThreshold: 5000},
start: 6000,
changes: []int64{-500, -600, -2000},
want: []bool{false, true, false},
},
{
name: "a charge that lands exactly on the threshold does not remind",
siteEnabled: true,
start: 1200,
changes: []int64{-200, -1},
want: []bool{false, true},
},
{
name: "user switched off",
siteEnabled: true,
setting: dto.UserSetting{QuotaWarningEnabled: &off},
start: 1500,
changes: []int64{-600},
want: []bool{false},
},
{
name: "site switched off",
setting: dto.UserSetting{QuotaWarningEnabled: &on},
start: 1500,
changes: []int64{-600},
want: []bool{false},
},
} {
t.Run(tc.name, func(t *testing.T) {
common.QuotaRemindEnabled = tc.siteEnabled
got := make([]bool, 0, len(tc.want))
balance := tc.start
for _, change := range tc.changes {
if change < 0 {
got = append(got, quotaReminderDue(tc.setting, balance, balance+change))
}
balance += change
}
assert.Equal(t, tc.want, got)
})
}
}
// Channel notices go to every enabled administrator whose channel switch is
// on, in that administrator's saved language. Until an administrator chooses,
// root receives them and other administrators do not; other users never do.
func TestChannelNoticesFollowEachAdministratorsSwitch(t *testing.T) {
previousDB := model.DB
previousCache, previousRedis := common.MemoryCacheEnabled, common.RedisEnabled
previousNotifyLimit := constant.NotifyLimitCount
previousClient, previousWorker := httpClient, system_setting.WorkerUrl
fetch := system_setting.GetFetchSetting()
previousFetch := *fetch
t.Cleanup(func() {
model.DB = previousDB
common.MemoryCacheEnabled, common.RedisEnabled = previousCache, previousRedis
constant.NotifyLimitCount = previousNotifyLimit
httpClient, system_setting.WorkerUrl = previousClient, previousWorker
*fetch = previousFetch
})
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
sqlDB, err := database.DB()
require.NoError(t, err)
sqlDB.SetMaxOpenConns(1)
t.Cleanup(func() { require.NoError(t, sqlDB.Close()) })
require.NoError(t, database.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.User{}))
model.DB = database
common.MemoryCacheEnabled, common.RedisEnabled = false, false
constant.NotifyLimitCount = 100
var mutex sync.Mutex
received := map[string][]string{}
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var payload WebhookPayload
if err := common.DecodeJson(r.Body, &payload); err == nil {
mutex.Lock()
recipient := strings.TrimPrefix(r.URL.Path, "/")
received[recipient] = append(received[recipient], payload.Type+" "+payload.Title)
mutex.Unlock()
}
w.WriteHeader(http.StatusNoContent)
}))
t.Cleanup(server.Close)
httpClient, system_setting.WorkerUrl = server.Client(), ""
fetch.EnableSSRFProtection = false
on, off := true, false
for _, user := range []struct {
name string
role int
status int
receives *bool
language string
}{
{"root-never-chose", common.RoleRootUser, common.UserStatusEnabled, nil, ""},
{"root-turned-off", common.RoleRootUser, common.UserStatusEnabled, &off, ""},
{"admin-never-chose", common.RoleAdminUser, common.UserStatusEnabled, nil, ""},
{"admin-turned-on", common.RoleAdminUser, common.UserStatusEnabled, &on, "zhCN"},
{"disabled-admin-turned-on", common.RoleAdminUser, common.UserStatusDisabled, &on, ""},
{"user-with-switch-stored-on", common.RoleCommonUser, common.UserStatusEnabled, &on, ""},
} {
setting, err := common.Marshal(dto.UserSetting{
NotifyType: dto.NotifyTypeWebhook,
WebhookUrl: server.URL + "/" + user.name,
ChannelNotifyEnabled: user.receives,
Language: user.language,
})
require.NoError(t, err)
require.NoError(t, database.Create(&model.User{Username: user.name, Role: user.role, Status: user.status, Setting: string(setting), AffCode: user.name}).Error)
}
channel := &model.Channel{Name: "notice-review", Key: "fixture-key", Type: 1, Status: common.ChannelStatusEnabled, Group: "default", Models: "test-model"}
require.NoError(t, channel.Insert())
DisableChannel(types.ChannelError{ChannelId: channel.Id, ChannelName: channel.Name, AutoBan: true}, "upstream failed")
EnableChannel(channel.Id, "", channel.Name)
NotifyChannelWatchers(dto.NotifyTypeChannelTest, func(lang string) (string, string) {
return i18n.Translate(lang, i18n.MsgChannelTestCompletedSubject), i18n.Translate(lang, i18n.MsgChannelTestCompletedContent)
})
disabledType := formatNotifyType(channel.Id, common.ChannelStatusAutoDisabled)
enabledType := formatNotifyType(channel.Id, common.ChannelStatusEnabled)
assert.Equal(t, map[string][]string{
"root-never-chose": {
fmt.Sprintf(`%s Channel "notice-review" (#%d) has been disabled`, disabledType, channel.Id),
fmt.Sprintf(`%s Channel "notice-review" (#%d) has been enabled`, enabledType, channel.Id),
"channel_test Channel test completed",
},
"admin-turned-on": {
fmt.Sprintf("%s 通道「notice-review」(#%d)已被禁用", disabledType, channel.Id),
fmt.Sprintf("%s 通道「notice-review」(#%d)已被启用", enabledType, channel.Id),
"channel_test 通道测试完成",
},
}, received)
}
+3 -3
View File
@@ -123,7 +123,7 @@ func ChargeViolationFeeIfNeeded(ctx *gin.Context, relayInfo *relaycommon.RelayIn
}
if err := PostConsumeQuota(relayInfo, feeQuota, 0, true); err != nil {
logger.LogError(ctx, fmt.Sprintf("failed to charge violation fee: %s", err.Error()))
logger.LogError(ctx, common.LogText("failed to charge violation fee: %s", err.Error()))
return false
}
@@ -132,7 +132,7 @@ func ChargeViolationFeeIfNeeded(ctx *gin.Context, relayInfo *relaycommon.RelayIn
useTimeSeconds := time.Now().Unix() - relayInfo.StartTime.Unix()
tokenName := ctx.GetString("token_name")
oai := apiErr.ToOpenAIError()
oai := apiErr.ToOpenAIError(common.UpstreamAddresses(ctx)...)
other := model.NewLogOther()
other.MergePublic(map[string]any{
@@ -152,7 +152,7 @@ func ChargeViolationFeeIfNeeded(ctx *gin.Context, relayInfo *relaycommon.RelayIn
ModelName: relayInfo.OriginModelName,
TokenName: tokenName,
Quota: feeQuota,
Content: "Violation fee charged",
Content: []*common.Message{common.NewMessage("Violation fee charged")},
TokenId: relayInfo.TokenId,
UseTimeSeconds: int(useTimeSeconds),
IsStream: relayInfo.IsStream,
+71
View File
@@ -0,0 +1,71 @@
package service
import (
_ "embed"
"encoding/base64"
"regexp"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
)
const VisionOutputTokens = 4096
const CialloPrompt = "ciallo"
func CialloQuality(answer string) string {
if strings.Contains(answer, "Ciallo") {
return "normal"
}
return "degraded"
}
const VisionPrompt = "不搜索,告诉我这是什么"
//go:embed vision.png
var visionImage []byte
var VisionImageURL = "data:image/png;base64," + base64.StdEncoding.EncodeToString(visionImage)
var visionUnavailableAnswer = regexp.MustCompile(`(?i)` +
`(?:我(?:这边|這邊|这里|這裡)?(?:目前|现在|現在|暂时|暫時)?(?:还|還|也|并|並)?|^(?:抱歉[,, ]*)?(?:目前|当前|當前|现在|現在|暂时|暫時)?)` +
`(?:没(?:有)?收到|沒有收到|未收到|(?:无法|無法|不能)(?:直接)?(?:查看|看到|读取|讀取|处理|處理|识别|識別)|看不到|收不到|没有看到|沒有看到)` +
`[^。!?!?\n,,;;]{0,24}(?:图片|圖片|图像|圖像|附件)(?:内容|內容)?(?:$|[。!?!?\n,,;;或和及])` +
`|(?:^|你(?:还|還)?|(?:消息|訊息|请求|請求)(?:中|里|裡)?)(?:没有|沒有|未)(?:附带|附帶|附上|上传|上傳|提供)[^。!?!?\n,,;;]{0,12}(?:图片|圖片|图像|圖像|附件)` +
`|\b(?:i|we)\s+(?:(?:currently|still)\s+)?(?:(?:cannot|can't|can’t|am unable to|are unable to|am not able to)\s+(?:see|view|access|read|process|receive)|(?:did not|didn't|didn’t|have not|haven't|haven’t)\s+(?:receive|received)|(?:do not|don't|don’t)\s+(?:see|have))[^.!?\n]{0,60}\b(?:images?|pictures?|photos?|attachments?)\b` +
`|\b(?:i['’]m|we['’]re)\s+unable to\s+(?:see|view|access|read|process)[^.!?\n]{0,60}\b(?:images?|pictures?|photos?|attachments?)\b` +
`|\bno\s+(?:images?|pictures?|photos?|attachments?)\s+(?:(?:was|is|has been|were|are)\s+)?(?:attached|provided|received|uploaded)\b`)
func VisionQuality(answer string) string {
answer = strings.Join(strings.Fields(answer), " ")
if visionUnavailableAnswer.MatchString(answer) {
return "incomplete"
}
if strings.Contains(answer, "音响") || strings.Contains(answer, "音箱") {
return "normal"
}
return "degraded"
}
// Before custom.79 keyword misses and transport failures both used failed.
// Interpret that history without rewriting persisted detection records.
func NormalizeVisionRun(run model.DetectionRun) model.DetectionRun {
if run.Kind != "vision" {
return run
}
var result fingerprint.Result
hasAnswer := common.UnmarshalJsonStr(run.Result, &result) == nil && strings.TrimSpace(result.Closest) != ""
if run.Status == "failed" {
run.Quality = "unknown"
if run.Error == "vision_keyword_missing" && hasAnswer {
run.Status = "succeeded"
run.Error = ""
}
}
if run.Status == "succeeded" && hasAnswer {
run.Quality = VisionQuality(result.Closest)
} else if run.Status == "succeeded" {
run.Status, run.Quality = "failed", "unknown"
}
return run
}
Binary file not shown.

After

Width:  |  Height:  |  Size: 2.0 MiB

+65
View File
@@ -0,0 +1,65 @@
package service
import (
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestVisionCapabilityClassification(t *testing.T) {
for _, tc := range []struct{ answer, quality string }{
{"这是音箱", "normal"},
{"这是石头,请重新上传图片", "degraded"},
{"我这边没收到可查看的图片或文件,暂时无法判断是什么。请重新上传图片", "incomplete"},
{"我这边没有收到可查看的图片或附件,所以暂时无法判断是什么", "incomplete"},
{"我无法查看图片,不能判断是不是音箱", "incomplete"},
{"我看不到你上传的图片", "incomplete"},
{"我這邊沒有收到可查看的圖片或附件", "incomplete"},
{"I didn't receive an image. Please upload it again.", "incomplete"},
{"I can't view the attached image.", "incomplete"},
{"I cannot process images.", "incomplete"},
{"I'm unable to see the image.", "incomplete"},
{"No image was attached.", "incomplete"},
{"目前无法直接查看图片,请重新上传。", "incomplete"},
{"你还没有上传任何图片。", "incomplete"},
{"本次请求中未附带图片。", "incomplete"},
{"The image shows a rock.", "degraded"},
{"我看不到图片上的品牌,但这是音箱", "normal"},
{"图片里是音箱,需要更清楚的细节可以重新上传图片", "normal"},
} {
t.Run(tc.answer, func(t *testing.T) {
assert.Equal(t, tc.quality, VisionQuality(tc.answer))
result, err := common.Marshal(map[string]string{"closest": tc.answer})
require.NoError(t, err)
for _, status := range []string{"succeeded", "failed"} {
run := NormalizeVisionRun(model.DetectionRun{Kind: "vision", Status: status, Quality: "degraded", Error: "vision_keyword_missing", Result: string(result)})
assert.Equal(t, "succeeded", run.Status)
assert.Equal(t, tc.quality, run.Quality)
}
})
}
}
func TestVisionLegacyFailureClassification(t *testing.T) {
for _, tc := range []struct{ reason, status, quality string }{
{"vision_keyword_missing", "succeeded", "degraded"},
{"vision_request_failed", "failed", "unknown"},
{"vision_response_invalid", "failed", "unknown"},
{"deadline_exceeded", "failed", "unknown"},
} {
run := NormalizeVisionRun(model.DetectionRun{Kind: "vision", Status: "failed", Quality: "degraded", Error: tc.reason, Result: `{"closest":"石头"}`})
assert.Equal(t, tc.status, run.Status)
assert.Equal(t, tc.quality, run.Quality)
}
empty := NormalizeVisionRun(model.DetectionRun{Kind: "vision", Status: "failed", Quality: "degraded", Error: "vision_keyword_missing", Result: `{"closest":""}`})
assert.Equal(t, "failed", empty.Status)
assert.Equal(t, "unknown", empty.Quality)
for _, result := range []string{"", `{}`, `{"closest":""}`} {
empty := NormalizeVisionRun(model.DetectionRun{Kind: "vision", Status: "succeeded", Quality: "normal", Result: result})
assert.Equal(t, "failed", empty.Status)
assert.Equal(t, "unknown", empty.Quality)
}
}
+6
View File
@@ -106,6 +106,9 @@ func CreateWaffoPancakeCheckoutSession(ctx context.Context, params *WaffoPancake
if strings.TrimSpace(params.OrderMerchantExternalID) == "" {
return nil, fmt.Errorf("missing order merchant external id")
}
if strings.TrimSpace(setting.WaffoPancakeStoreID) == "" {
return nil, fmt.Errorf("missing Waffo Pancake store id")
}
client, err := newWaffoPancakeClient()
if err != nil {
return nil, fmt.Errorf("build Waffo Pancake client: %w", err)
@@ -166,6 +169,9 @@ func VerifyConfiguredWaffoPancakeWebhook(payload string, signatureHeader string)
if err != nil {
return nil, err
}
if storeID := strings.TrimSpace(setting.WaffoPancakeStoreID); storeID == "" || evt.StoreID != storeID {
return nil, fmt.Errorf("Waffo Pancake webhook store mismatch")
}
identity := ""
if evt.Data.MerchantProvidedBuyerIdentity != nil {
identity = *evt.Data.MerchantProvidedBuyerIdentity
+2 -1
View File
@@ -7,6 +7,7 @@ import (
"encoding/hex"
"fmt"
"net/http"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
@@ -35,7 +36,7 @@ func SendWebhookNotify(webhookURL string, secret string, data dto.Notify) error
// 处理占位符
content := data.Content
for _, value := range data.Values {
content = fmt.Sprintf(content, value)
content = strings.Replace(content, dto.ContentValueParam, fmt.Sprintf("%v", value), 1)
}
// 构建 webhook 负载