feat: publish sanitized custom.97 source
This commit is contained in:
1022 files changed
+69903
-19207
No files matched your search
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
@@ -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
@@ -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,
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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()))
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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
@@ -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
@@ -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,
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
@@ -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
@@ -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))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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,
|
||||
))
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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{
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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"
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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
@@ -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()
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
@@ -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 != "" {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
@@ -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}
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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))
|
||||
})
|
||||
|
||||
@@ -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() {
|
||||
|
||||
@@ -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", ""},
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
@@ -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
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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
@@ -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
@@ -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"])
|
||||
}
|
||||
@@ -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"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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" {
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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 |
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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 负载
|
||||
|
||||
Reference in new issue
Block a user