feat: sync deployed v1.0.0-rc.38-custom.42
Publish Docker image (Multi-arch) / Build & push (amd64) (push) Failing after 1m40s
Release (Linux, macOS, Windows) / Linux Release (push) Failing after 42m16s
Publish Docker image (Multi-arch) / Build & push (arm64) (push) Canceled after 0s
Release (Linux, macOS, Windows) / macOS Release (push) Canceled after 0s
Release (Linux, macOS, Windows) / Windows Release (push) Canceled after 0s
Publish Docker image (Multi-arch) / Create multi-arch manifests (push) Canceled after 0s

Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
This commit is contained in:
cneicy committed 2026-09-26 04:49:11 +08:00
1 parent c2b44332d6
commit a5f7b8e586
828 files changed
+122253 -13702

No files matched your search

+256
View File
@@ -0,0 +1,256 @@
package controller
import (
"bytes"
"context"
"crypto/rand"
"errors"
"fmt"
"math/big"
"net/http"
"net/http/httptest"
"sync"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/middleware"
"github.com/QuantumNous/new-api/model"
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
"github.com/QuantumNous/new-api/relay"
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/service"
"github.com/gin-gonic/gin"
)
type astraSamplingHandler struct{}
type astraSamplingTaskPayload struct {
Force bool `json:"force,omitempty"`
}
func (astraSamplingHandler) Type() string { return service.AstraSamplingTaskType }
func (astraSamplingHandler) Enabled() bool { return service.AstraSamplingEnabled() }
func (astraSamplingHandler) Interval() time.Duration { return service.AstraSamplingSchedulerInterval() }
func (astraSamplingHandler) ScheduleFromCreation() bool { return true }
func (astraSamplingHandler) NewPayload() any { return nil }
func (astraSamplingHandler) ScheduleDue(ctx context.Context, now time.Time) (bool, error) {
channels, err := service.AstraSamplingChannelsDue(ctx, now)
return len(channels) > 0, err
}
func (astraSamplingHandler) Run(parent context.Context, task *model.SystemTask, runnerID string) {
ctx, cancel := context.WithTimeout(parent, 25*time.Minute)
defer cancel()
payload := astraSamplingTaskPayload{}
if err := task.DecodePayload(&payload); err != nil {
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, err)
return
}
channels, err := service.AstraSamplingChannelsDue(ctx, time.Now())
if payload.Force {
channels, err = service.AstraSamplingChannels(ctx)
}
if err != nil {
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, errors.New("channel query failed"))
return
}
state := service.AstraSamplingState{Rounds: []service.AstraSamplingRound{}}
var mu sync.Mutex
var persistErr error
// Persist before sending, and after every probe. A failed lease or restart
// never resumes a paid probe; unfinished records remain visible as errors.
persist := func() bool {
if persistErr != nil {
return false
}
persistErr = model.UpdateSystemTaskState(task.TaskID, runnerID, state)
if persistErr != nil {
cancel()
}
return persistErr == nil
}
jobs := make(chan int)
var workers sync.WaitGroup
for worker := 0; worker < 2; worker++ {
workers.Add(1)
go func() {
defer workers.Done()
for id := range jobs {
if ctx.Err() != nil {
continue
}
mu.Lock()
index := len(state.Rounds)
state.Rounds = append(state.Rounds, service.AstraSamplingRound{ID: common.GetRandomString(16), ChannelID: id, StartedAt: time.Now().Unix(), Version: fingerprint.Version, Result: fingerprint.Result{Status: "running"}})
ok := persist()
mu.Unlock()
if !ok {
continue
}
channel, loadErr := model.GetChannelById(id, true)
samples := []fingerprint.Sample{}
reason := ""
usage := []fingerprint.ProbeResponse{}
if loadErr != nil || channel.Status != common.ChannelStatusEnabled {
reason = "channel_unavailable"
} else {
for probe := 0; probe < 3 && ctx.Err() == nil; probe++ {
mu.Lock()
state.Rounds[index].Sent++
ok = persist()
mu.Unlock()
if !ok {
break
}
answer, count, probeErr := runAstraFingerprintProbe(ctx, channel)
usage = append(usage, answer)
if probeErr != nil {
if errors.Is(probeErr, fingerprint.ErrInvalidSample) {
reason = "invalid_sample"
} else {
reason = "transport_error"
break
}
} else {
samples = append(samples, fingerprint.Sample{Text: answer.Text, Count: count})
}
mu.Lock()
state.Rounds[index].Usage = append([]fingerprint.ProbeResponse(nil), usage...)
ok = persist()
mu.Unlock()
if !ok {
break
}
}
}
result, scoreErr := fingerprint.Score(samples)
if scoreErr != nil {
reason = "scoring_error"
}
if ctx.Err() != nil {
reason = "interrupted"
}
if reason != "" && reason != "invalid_sample" {
result.Status = "error"
}
mu.Lock()
state.Rounds[index].Result = result
state.Rounds[index].Reason = reason
state.Rounds[index].Usage = usage
state.Rounds[index].FinishedAt = time.Now().Unix()
if persist() && !payload.Force {
if err := model.ApplyDetectionScheduleResult(ctx, id, "fingerprint", state.Rounds[index].ID, state.Rounds[index].FinishedAt, service.FingerprintScheduleOutcome(service.AstraSamplingResultStatus(result))); err != nil {
persistErr = err
cancel()
}
}
mu.Unlock()
}
}()
}
for _, channel := range channels {
select {
case jobs <- channel.Id:
case <-ctx.Done():
}
}
close(jobs)
workers.Wait()
if persistErr != nil || ctx.Err() != nil {
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, errors.New("sampling interrupted"))
return
}
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusSucceeded, map[string]any{"channels": len(state.Rounds), "model": fingerprint.Model}, nil)
}
func runAstraFingerprintProbe(parent context.Context, channel *model.Channel, paid ...bool) (fingerprint.ProbeResponse, int, error) {
ctx, cancel := context.WithTimeout(parent, 180*time.Second)
defer cancel()
seed, err := rand.Int(rand.Reader, big.NewInt(82))
if err != nil {
return fingerprint.ProbeResponse{}, 0, err
}
count := 292 + int(seed.Int64()/2)
prompt := fmt.Sprintf("Directly choose %d integers from 1 through 355, inclusive. Allow repeats. Do not sort, balance frequencies, repair duplicates, or deliberately create an arithmetic pattern.\nReply directly with one JSON integer array and no explanation. Do not call tools, read files, execute code or ask another model. Do not continue the preceding task.", count)
if seed.Int64()%2 == 0 {
prompt = fmt.Sprintf("直接选择 %d 个 1 到 355(含边界)的整数。允许重复;不要排序、平衡频数、修复重复,也不要刻意构造等差规律。\n直接回答一个 JSON 整数数组,不要解释。不要调用工具、读文件、运行代码或让其他模型代答;不要继续之前的任务。", count)
}
input, _ := common.Marshal([]map[string]any{{"role": "user", "content": []map[string]string{{"type": "input_text", "text": prompt}}}})
stream := true
limit := uint(4096)
request := &dto.OpenAIResponsesRequest{Model: fingerprint.Model, Input: input, Instructions: []byte(`""`), Stream: &stream, MaxOutputTokens: &limit, Store: []byte("false"), Tools: []byte("[]"), ToolChoice: []byte(`"none"`)}
encoded, err := common.Marshal(request)
if err != nil {
return fingerprint.ProbeResponse{}, count, err
}
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequestWithContext(ctx, http.MethodPost, "/v1/responses", bytes.NewReader(encoded))
c.Request.Header.Set("Content-Type", "application/json")
if apiErr := middleware.SetupContextForSelectedChannel(c, channel, fingerprint.Model); apiErr != nil {
return fingerprint.ProbeResponse{}, count, errors.New("channel setup failed")
}
info, err := relaycommon.GenRelayInfo(c, types.RelayFormatOpenAIResponses, request, nil)
if err != nil {
return fingerprint.ProbeResponse{}, count, err
}
info.IsChannelTest = true
info.DisablePing = true
info.AstraDetectionDisabled = true
info.AstraAutoSwitch = false
adaptor, body, closer, apiErr := relay.PrepareResponsesRequest(c, info, request)
if closer != nil {
defer closer.Close()
}
if apiErr != nil {
return fingerprint.ProbeResponse{}, count, errors.New("probe preparation failed")
}
if len(paid) > 0 && paid[0] {
bounded, boundedCloser, boundErr := boundDetectionBody(body, request)
if boundErr != nil {
return fingerprint.ProbeResponse{}, count, boundErr
}
defer boundedCloser.Close()
body = bounded
}
response, err := adaptor.DoRequest(c, info, body)
if err != nil {
return fingerprint.ProbeResponse{}, count, errors.New("probe transport failed")
}
resp, ok := response.(*http.Response)
if !ok || resp == nil {
return fingerprint.ProbeResponse{}, count, errors.New("unsupported probe transport")
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return fingerprint.ProbeResponse{}, count, errors.New("probe HTTP failure")
}
// Deliberately bypass DoResponse: this is a platform-funded probe, with no
// customer token, wallet charge, session binding or generated response relay.
answer, err := fingerprint.ReadResponse(resp.Body, resp.Header.Get("Content-Type"))
return answer, count, err
}
func TestAstraFingerprint(c *gin.Context) {
if !service.AstraSamplingEnabled() {
c.JSON(http.StatusConflict, gin.H{"success": false, "message": "Astra sampling is disabled"})
return
}
latest, err := model.GetLatestSystemTasks([]string{service.AstraSamplingTaskType})
if err != nil {
common.ApiError(c, err)
return
}
if last := latest[service.AstraSamplingTaskType]; last != nil && time.Now().Unix()-last.CreatedAt < 300 {
c.JSON(http.StatusTooManyRequests, gin.H{"success": false, "message": "Please wait five minutes between sampling rounds"})
return
}
task, created, err := service.EnqueueSystemTask(service.AstraSamplingTaskType, astraSamplingTaskPayload{Force: true})
if err != nil {
common.ApiError(c, err)
return
}
c.JSON(http.StatusOK, gin.H{"success": true, "data": gin.H{"task_id": task.TaskID, "created": created}})
}
+37
View File
@@ -0,0 +1,37 @@
package controller
import (
"context"
"io"
"net/http"
"net/http/httptest"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
"github.com/tidwall/gjson"
)
func TestAstraSamplingUsesIsolatedResponsesWithoutGuard(t *testing.T) {
requests := 0
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
requests++
body, _ := io.ReadAll(r.Body)
if r.URL.Path != "/v1/responses" || gjson.GetBytes(body, "model").String() != fingerprint.Model || !gjson.GetBytes(body, "stream").Bool() {
t.Errorf("wrong probe request %s %s", r.URL.Path, body)
}
if gjson.GetBytes(body, "previous_response_id").Exists() || gjson.GetBytes(body, "conversation").Exists() {
t.Error("probe carried customer session")
}
w.Header().Set("Content-Type", "text/event-stream")
_, _ = io.WriteString(w, "data: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"model\":\"gpt-5.6-luna\",\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"[1,2,3]\"}]}]}}\n\n")
}))
defer server.Close()
url := server.URL
channel := &model.Channel{Id: 1, Type: 1, Key: "test-only", BaseURL: &url, Status: common.ChannelStatusEnabled, Models: fingerprint.Model}
r, count, err := runAstraFingerprintProbe(context.Background(), channel)
if err != nil || r.DeclaredModel != "gpt-5.6-luna" || count < 292 || count > 332 || requests != 1 {
t.Fatalf("probe %+v count=%d requests=%d err=%v", r, count, requests, err)
}
}
+34
View File
@@ -1,8 +1,10 @@
package controller
import (
"errors"
"fmt"
"os"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
@@ -43,6 +45,11 @@ var auditContentTemplates = map[string]string{
"user.reset_passkey": "Reset the user passkey",
"option.update": "Updated system setting ${key}",
"option.passkey_domains": "Updated Passkey domains: removed ${domains}; affected ${known}; unknown ${unknown}",
"option.passkey_domains_confirmed": "Confirmed removal of Passkey domains: ${domains}; affected ${known}; unknown ${unknown}",
"option.passkey_domains_blocked": "Passkey domain change blocked: ${domains}; affected ${known}; unknown ${unknown}",
"option.passkey_domains_failed": "Passkey domain update failed",
"channel.create": "Created channel ${name} (type ${type}, count ${count})",
"channel.update": "Updated channel ${name} (ID: ${id})",
"channel.delete": "Deleted channel ${name} (ID: ${id})",
@@ -65,6 +72,33 @@ var auditContentTemplates = map[string]string{
"subscription.user_plan_reset": "Reset active plan ${plan_id} subscriptions for user ${target_user_id}",
}
func recordPasskeyDomainAudit(c *gin.Context, change *model.PasskeyDomainChange, confirmed bool, err error) {
confirmed = confirmed && err == nil && change != nil && len(change.RemovedRPIDs) > 0
params := map[string]any{"success": err == nil, "confirmed": confirmed}
if change != nil {
params["domains"] = strings.Join(change.RemovedRPIDs, ", ")
params["removed_rp_ids"] = change.RemovedRPIDs
params["known"] = change.AffectedCredentials
params["unknown"] = change.UnknownCredentials
params["previous_rp_id"] = change.PreviousRPID
params["effective_rp_id"] = change.EffectiveRPID
}
action := "option.passkey_domains"
if errors.Is(err, model.ErrPasskeyDomainRemovalConfirmation) {
action = "option.passkey_domains_blocked"
} else if err != nil {
action = "option.passkey_domains_failed"
} else if confirmed && change != nil && len(change.RemovedRPIDs) > 0 {
action = "option.passkey_domains_confirmed"
}
auditInfo := &model.AuditRequestInfo{
Method: c.Request.Method, Route: c.FullPath(), Path: c.FullPath(),
Status: c.Writer.Status(), Success: err == nil,
}
model.RecordOperationAuditLog(c.GetInt("id"), c.GetInt("role"), auditContentEN(action, params), c.ClientIP(), action, params, auditOperatorInfo(c), auditInfo, c)
markAuditLogged(c)
}
// auditContentEN 按 action 模板渲染英文兜底文本;未登记的 action 退回 action 本身。
func auditContentEN(action string, params map[string]any) string {
tmpl, ok := auditContentTemplates[action]
+314 -7
View File
@@ -76,7 +76,7 @@ func TestSecurityLoginCodeCompletesOnce(t *testing.T) {
user, _ := setupSecurityEnrollmentTest(t)
factor := &model.TwoFA{UserId: user.Id, Secret: "JBSWY3DPEHPK3PXP", IsEnabled: true}
require.NoError(t, model.DB.Create(factor).Error)
challenge, err := service.StartLoginVerification(user, "password")
challenge, err := service.StartLoginVerification(user, "password", nil)
require.NoError(t, err)
require.NotNil(t, challenge)
code, err := totp.GenerateCode(factor.Secret, time.Now())
@@ -121,7 +121,7 @@ func TestSecurityLoginRejectsChangedOrExpiredAuthorization(t *testing.T) {
user, _ := setupSecurityEnrollmentTest(t)
factor := &model.TwoFA{UserId: user.Id, Secret: "JBSWY3DPEHPK3PXP", IsEnabled: true}
require.NoError(t, model.DB.Create(factor).Error)
challenge, err := service.StartLoginVerification(user, "password")
challenge, err := service.StartLoginVerification(user, "password", nil)
require.NoError(t, err)
method := "2fa"
switch change {
@@ -183,7 +183,7 @@ func TestSecurityLoginPasskeyDoesNotRequireAdditionalTwoFA(t *testing.T) {
require.Equal(t, "required", result.Data.Options.PublicKey.UserVerification)
flowToken, challenge = result.Data.FlowToken, result.Data.Options.PublicKey.Challenge
} else {
pending, err := service.StartLoginVerification(user, "oauth:github")
pending, err := service.StartLoginVerification(user, "oauth:github", nil)
require.NoError(t, err)
parentToken = pending.FlowToken
flowToken, challenge = beginSecurityLoginPasskey(t, parentToken)
@@ -225,7 +225,7 @@ func TestSecurityLoginPasskeyDoesNotRequireAdditionalTwoFA(t *testing.T) {
func TestSecurityLoginPasskeyConcurrentCompletionCreatesOneSession(t *testing.T) {
user, _ := setupSecurityEnrollmentTest(t)
key := newSecurityLoginPasskey(t, user.Id)
pending, err := service.StartLoginVerification(user, "password")
pending, err := service.StartLoginVerification(user, "password", nil)
require.NoError(t, err)
requests := make([]string, 2)
for index := range requests {
@@ -266,7 +266,7 @@ func TestSecurityLoginPasskeyConcurrentCompletionCreatesOneSession(t *testing.T)
func TestSecurityLoginSessionFailureRollsBackChallengeConsumption(t *testing.T) {
user, _ := setupSecurityEnrollmentTest(t)
key := newSecurityLoginPasskey(t, user.Id)
pending, err := service.StartLoginVerification(user, "password")
pending, err := service.StartLoginVerification(user, "password", nil)
require.NoError(t, err)
require.NoError(t, model.DB.Callback().Create().Before("gorm:create").Register("login_session_failure", func(tx *gorm.DB) {
if tx.Statement.Table == "user_sessions" {
@@ -410,7 +410,7 @@ func TestSecurityLoginPasskeyCannotCompleteAnotherChallenge(t *testing.T) {
t.Run(fmt.Sprintf("other-user=%t", otherUser), func(t *testing.T) {
user, _ := setupSecurityEnrollmentTest(t)
key := newSecurityLoginPasskey(t, user.Id)
first, err := service.StartLoginVerification(user, "password")
first, err := service.StartLoginVerification(user, "password", nil)
require.NoError(t, err)
passkeyToken, challenge := beginSecurityLoginPasskey(t, first.FlowToken)
if otherUser {
@@ -418,7 +418,7 @@ func TestSecurityLoginPasskeyCannotCompleteAnotherChallenge(t *testing.T) {
require.NoError(t, model.DB.Create(user).Error)
newSecurityLoginPasskey(t, user.Id)
}
second, err := service.StartLoginVerification(user, "password")
second, err := service.StartLoginVerification(user, "password", nil)
require.NoError(t, err)
body, err := common.Marshal(map[string]any{"flow_token": second.FlowToken, "passkey_flow_token": passkeyToken, "credential": securityPasskeyResponse(t, key, challenge, false, 0)})
require.NoError(t, err)
@@ -816,3 +816,310 @@ func TestOAuthBindProviderErrorConsumesSessionBoundFlow(t *testing.T) {
assert.Zero(t, provider.exchangeCalls)
assert.Zero(t, provider.userInfoCalls)
}
// legacyGitHubOAuthProvider follows the GitHub provider contract: the numeric
// account ID identifies the user, Extra["legacy_id"] carries the username and
// the verified email list is served on demand. Lookups go through the real
// github_id model queries.
type legacyGitHubOAuthProvider struct {
authFlowTestOAuthProvider
providerUserID string
legacyID string
verifiedEmails []string
verifiedEmailsErr error
verifiedEmailCalls int
}
func (*legacyGitHubOAuthProvider) GetName() string { return "GitHub" }
func (provider *legacyGitHubOAuthProvider) GetUserInfo(context.Context, *oauth.OAuthToken) (*oauth.OAuthUser, error) {
return &oauth.OAuthUser{ProviderUserID: provider.providerUserID, Username: provider.legacyID, Extra: map[string]any{"legacy_id": provider.legacyID}}, nil
}
func (provider *legacyGitHubOAuthProvider) GetVerifiedEmails(context.Context, *oauth.OAuthToken) ([]string, error) {
provider.verifiedEmailCalls++
return provider.verifiedEmails, provider.verifiedEmailsErr
}
func (*legacyGitHubOAuthProvider) IsUserIDTaken(providerUserID string) bool {
return model.IsGitHubIdAlreadyTaken(providerUserID)
}
func (*legacyGitHubOAuthProvider) FillUserByProviderID(user *model.User, providerUserID string) error {
user.GitHubId = providerUserID
return user.FillUserByGitHubId()
}
func (*legacyGitHubOAuthProvider) SetProviderUserID(user *model.User, providerUserID string) {
user.GitHubId = providerUserID
}
func (*legacyGitHubOAuthProvider) GetProviderPrefix() string { return "github_" }
func (*legacyGitHubOAuthProvider) ProviderUserIDColumn() string { return "github_id" }
// legacyGitHubOAuthLogin registers provider for one test and completes an OAuth
// login callback with it.
func legacyGitHubOAuthLogin(t *testing.T, provider *legacyGitHubOAuthProvider) *httptest.ResponseRecorder {
t.Helper()
const slug = "github-legacy-login-test"
oauth.Register(slug, provider)
t.Cleanup(func() { oauth.Unregister(slug) })
token, _, err := model.CreateAuthFlow(model.AuthFlowCreate{Purpose: model.AuthFlowPurposeOAuth, Provider: slug, Intent: model.AuthFlowIntentLogin, Payload: `{}`, ExpiresAt: time.Now().Add(time.Minute)})
require.NoError(t, err)
router := gin.New()
router.GET("/api/oauth/:provider", HandleOAuth)
response := httptest.NewRecorder()
router.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "/api/oauth/"+slug+"?state="+token+"&code=provider-code", nil))
return response
}
// legacyGitHubBindingAudits returns the account binding audit rows in creation
// order together with their encoded parameters.
func legacyGitHubBindingAudits(t *testing.T) ([]model.AuditLog, []string) {
t.Helper()
var audits []model.AuditLog
require.NoError(t, model.LOG_DB.Where("action = ?", "user.binding_bind").Order("id").Find(&audits).Error)
params := make([]string, 0, len(audits))
for _, audit := range audits {
require.NotNil(t, audit.Other.Op)
encoded, err := common.Marshal(audit.Other.Op.Params)
require.NoError(t, err)
params = append(params, string(encoded))
}
return audits, params
}
func TestOAuthLoginLegacyGitHubBindingRequiresAccountEvidence(t *testing.T) {
const declined = "This GitHub account cannot be linked to an existing account automatically, please sign in or register another way and then link GitHub in account settings"
tests := []struct {
name string
existingGitHubID string
existingDeleted bool
withoutEmail bool
legacyID string
verifiedEmails []string
verifiedEmailsErr error
registerEnabled bool
expectLogin bool
expectMigration bool
expectNewAccount bool
expectEmailCalls int
expectAuditParams string
}{
{name: "numeric legacy value registers a new account", existingGitHubID: "424242", legacyID: "424242", verifiedEmails: []string{"legacy-github@example.com"}, registerEnabled: true, expectLogin: true, expectNewAccount: true},
{name: "numeric legacy value with registration disabled", existingGitHubID: "424242", legacyID: "424242", verifiedEmails: []string{"legacy-github@example.com"}},
{name: "soft-deleted legacy row registers a new account", existingGitHubID: "octocat-legacy", existingDeleted: true, legacyID: "octocat-legacy", verifiedEmails: []string{"legacy-github@example.com"}, registerEnabled: true, expectLogin: true, expectNewAccount: true},
{
name: "verified email match migrates", existingGitHubID: "octocat-legacy", legacyID: "octocat-legacy",
verifiedEmails: []string{"other@example.com", " Legacy-GitHub@Example.com "}, expectLogin: true, expectMigration: true, expectEmailCalls: 1,
// No SMTP server is configured in tests, so the notification attempt is recorded as failed.
expectAuditParams: `{"provider":"github","legacy_migration":true,"legacy_id":"octocat-legacy","provider_user_id":"900001","verified_email_matched":true,"success":true,"notification_failed":true}`,
},
{
name: "no matching verified email declines", existingGitHubID: "octocat-legacy", legacyID: "octocat-legacy",
verifiedEmails: []string{"other@example.com"}, registerEnabled: true, expectEmailCalls: 1,
expectAuditParams: `{"provider":"github","legacy_migration":true,"legacy_id":"octocat-legacy","provider_user_id":"900001","success":false,"reason":"no_matching_evidence"}`,
},
{
name: "verified emails failure declines", existingGitHubID: "octocat-legacy", legacyID: "octocat-legacy",
verifiedEmailsErr: errors.New("emails unavailable"), registerEnabled: true, expectEmailCalls: 1,
expectAuditParams: `{"provider":"github","legacy_migration":true,"legacy_id":"octocat-legacy","provider_user_id":"900001","success":false,"reason":"verified_emails_unavailable"}`,
},
{
name: "account without email declines before asking the provider", existingGitHubID: "octocat-legacy", withoutEmail: true, legacyID: "octocat-legacy",
verifiedEmails: []string{"legacy-github@example.com"}, registerEnabled: true,
expectAuditParams: `{"provider":"github","legacy_migration":true,"legacy_id":"octocat-legacy","provider_user_id":"900001","success":false,"reason":"no_matching_evidence"}`,
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
setupSecurityEnrollmentTest(t)
previousRegister := common.RegisterEnabled
common.RegisterEnabled = test.registerEnabled
t.Cleanup(func() { common.RegisterEnabled = previousRegister })
existing := &model.User{Username: "legacy-github", Email: "legacy-github@example.com", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Group: "default", AffCode: "legacy-github", AuthVersion: 1, GitHubId: test.existingGitHubID}
if test.withoutEmail {
existing.Email = ""
}
require.NoError(t, model.DB.Create(existing).Error)
if test.existingDeleted {
require.NoError(t, model.DB.Delete(existing).Error)
}
provider := &legacyGitHubOAuthProvider{providerUserID: "900001", legacyID: test.legacyID, verifiedEmails: test.verifiedEmails, verifiedEmailsErr: test.verifiedEmailsErr}
response := legacyGitHubOAuthLogin(t, provider)
var result struct {
Success bool `json:"success"`
Message string `json:"message"`
Data struct {
User struct {
Id int `json:"id"`
} `json:"user"`
} `json:"data"`
}
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
require.Equal(t, test.expectLogin, result.Success, response.Body.String())
assert.Equal(t, test.expectEmailCalls, provider.verifiedEmailCalls)
var reloaded model.User
require.NoError(t, model.DB.Unscoped().First(&reloaded, existing.Id).Error)
audits, params := legacyGitHubBindingAudits(t)
if test.expectAuditParams == "" {
assert.Empty(t, audits)
} else {
require.Len(t, audits, 1)
assert.Equal(t, existing.Id, audits[0].UserId)
assert.Equal(t, common.RoleCommonUser, audits[0].ActorRole)
assert.Equal(t, test.expectMigration, audits[0].Success)
assert.JSONEq(t, test.expectAuditParams, params[0])
}
if test.expectMigration {
assert.Equal(t, existing.Id, result.Data.User.Id)
assert.Equal(t, "900001", reloaded.GitHubId)
assert.NotEmpty(t, response.Header().Values("Set-Cookie"))
return
}
assert.Equal(t, test.existingGitHubID, reloaded.GitHubId, "the existing binding must stay untouched")
if test.expectNewAccount {
var created model.User
require.NoError(t, model.DB.Where("github_id = ?", "900001").First(&created).Error)
assert.NotEqual(t, existing.Id, created.Id)
assert.Equal(t, created.Id, result.Data.User.Id)
return
}
assert.Empty(t, response.Header().Values("Set-Cookie"))
if test.expectAuditParams != "" {
assert.Equal(t, declined, result.Message)
}
})
}
}
func TestOAuthLoginLegacyGitHubBindingMigratesAfterLoginVerification(t *testing.T) {
for _, interference := range []string{"none", "pending id claimed by another account", "binding relinked meanwhile"} {
t.Run(interference, func(t *testing.T) {
setupSecurityEnrollmentTest(t)
existing := &model.User{Username: "legacy-github", Email: "legacy-github@example.com", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Group: "default", AffCode: "legacy-github", AuthVersion: 1, GitHubId: "octocat-legacy"}
require.NoError(t, model.DB.Create(existing).Error)
factor := &model.TwoFA{UserId: existing.Id, Secret: "JBSWY3DPEHPK3PXP", IsEnabled: true}
require.NoError(t, model.DB.Create(factor).Error)
provider := &legacyGitHubOAuthProvider{providerUserID: "900001", legacyID: "octocat-legacy", verifiedEmails: []string{"legacy-github@example.com"}}
response := legacyGitHubOAuthLogin(t, provider)
var challenge struct {
Success bool `json:"success"`
Data service.LoginChallenge `json:"data"`
}
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &challenge))
require.True(t, challenge.Success, response.Body.String())
require.True(t, challenge.Data.RequireVerification)
assert.Equal(t, []service.VerificationMethodOption{{Method: "2fa", Available: true}}, challenge.Data.Methods)
assert.Empty(t, response.Header().Values("Set-Cookie"))
assert.Zero(t, provider.verifiedEmailCalls, "a second factor is checked before the provider emails")
var reloaded model.User
require.NoError(t, model.DB.First(&reloaded, existing.Id).Error)
assert.Equal(t, "octocat-legacy", reloaded.GitHubId, "the binding waits for the verification")
audits, _ := legacyGitHubBindingAudits(t)
assert.Empty(t, audits)
expectedGitHubID, expectLogin := "900001", true
switch interference {
case "pending id claimed by another account":
other := &model.User{Username: "other-github", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Group: "default", AffCode: "other-github", AuthVersion: 1, GitHubId: "900001"}
require.NoError(t, model.DB.Create(other).Error)
expectedGitHubID, expectLogin = "octocat-legacy", false
case "binding relinked meanwhile":
require.NoError(t, model.DB.Model(existing).Update("github_id", "relinked-legacy").Error)
expectedGitHubID = "relinked-legacy"
}
code, err := totp.GenerateCode(factor.Secret, time.Now())
require.NoError(t, err)
body, err := common.Marshal(map[string]string{"flow_token": challenge.Data.FlowToken, "code": code})
require.NoError(t, err)
for attempt := range 2 {
response = securityEnrollmentRequest(http.MethodPost, "/api/user/login/verify", string(body), "", service.AuthIdentity{}, VerifyLogin)
var result securityEnrollmentResponse
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
assert.Equal(t, attempt == 0 && expectLogin, result.Success, response.Body.String())
if attempt == 0 && !expectLogin {
assert.Equal(t, "ACCOUNT_ALREADY_BOUND", result.Code)
}
if result.Success {
assert.NotEmpty(t, response.Header().Values("Set-Cookie"))
} else {
assert.Empty(t, response.Header().Values("Set-Cookie"))
}
}
require.NoError(t, model.DB.First(&reloaded, existing.Id).Error)
assert.Equal(t, expectedGitHubID, reloaded.GitHubId)
sessions, err := model.CountActiveUserSessions(existing.Id, time.Now().Unix())
require.NoError(t, err)
expectedSessions := 0
if expectLogin {
expectedSessions = 1
}
assert.EqualValues(t, expectedSessions, sessions, "the consumed challenge cannot be replayed")
audits, params := legacyGitHubBindingAudits(t)
if interference != "none" {
assert.Empty(t, audits)
return
}
require.Len(t, audits, 1)
assert.Equal(t, existing.Id, audits[0].UserId)
assert.Equal(t, common.RoleCommonUser, audits[0].ActorRole)
assert.True(t, audits[0].Success)
// No SMTP server is configured in tests, so the notification attempt is recorded as failed.
assert.JSONEq(t, `{"provider":"github","legacy_migration":true,"legacy_id":"octocat-legacy","provider_user_id":"900001","verification_method":"2fa","success":true,"notification_failed":true}`, params[0])
})
}
}
func TestOAuthBindIgnoresLegacyGitHubUsernames(t *testing.T) {
tests := []struct {
name string
ownGitHubID string
otherGitHubID string
legacyID string
expectBound bool
}{
{name: "another account's legacy username does not block binding", otherGitHubID: "octocat-legacy", legacyID: "octocat-legacy", expectBound: true},
{name: "own legacy username is replaced", ownGitHubID: "octocat-legacy", otherGitHubID: "unrelated-legacy", legacyID: "octocat-legacy", expectBound: true},
{name: "numeric legacy value does not collide", otherGitHubID: "424242", legacyID: "424242", expectBound: true},
{name: "numeric ID held by another account is rejected", otherGitHubID: "900001", legacyID: "octocat-legacy"},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
user, identity := setupSecurityEnrollmentTest(t)
if test.ownGitHubID != "" {
require.NoError(t, model.DB.Model(user).Update("github_id", test.ownGitHubID).Error)
}
other := &model.User{Username: "other-github", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Group: "default", AffCode: "other-github", AuthVersion: 1, GitHubId: test.otherGitHubID}
require.NoError(t, model.DB.Create(other).Error)
const slug = "github-legacy-bind-test"
oauth.Register(slug, &legacyGitHubOAuthProvider{providerUserID: "900001", legacyID: test.legacyID})
t.Cleanup(func() { oauth.Unregister(slug) })
proof := issueSecurityEnrollmentProof(t, identity, service.VerificationOperation{Scope: service.VerificationScopeAccountBind, Context: []byte(`{"provider":"` + slug + `"}`)}, service.VerificationMethodPassword)
started := securityEnrollmentRequest(http.MethodPost, "/api/oauth/state", `{"provider":"`+slug+`","intent":"bind"}`, proof, identity, GenerateOAuthCode)
var flow struct {
Data struct {
FlowToken string `json:"flow_token"`
} `json:"data"`
}
require.NoError(t, common.Unmarshal(started.Body.Bytes(), &flow))
require.NotEmpty(t, flow.Data.FlowToken, started.Body.String())
response := securityEnrollmentRequest(http.MethodGet, "/api/oauth/"+slug+"?state="+flow.Data.FlowToken+"&code=provider-code", "", "", identity, func(c *gin.Context) {
c.Params = gin.Params{{Key: "provider", Value: slug}}
HandleOAuth(c)
})
var result struct {
Success bool `json:"success"`
Message string `json:"message"`
}
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
assert.Equal(t, test.expectBound, result.Success, response.Body.String())
var bound, untouched model.User
require.NoError(t, model.DB.First(&bound, user.Id).Error)
require.NoError(t, model.DB.First(&untouched, other.Id).Error)
assert.Equal(t, test.otherGitHubID, untouched.GitHubId)
if test.expectBound {
assert.Equal(t, "900001", bound.GitHubId)
return
}
assert.Equal(t, test.ownGitHubID, bound.GitHubId)
assert.Equal(t, "This GitHub account has already been bound", result.Message)
})
}
}
+1 -1
View File
@@ -146,7 +146,7 @@ func TestSessionLimitDoesNotRecordRejectedLoginAsSuccessful(t *testing.T) {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/api/user/login", nil)
setupLogin(user, c)
setupLogin(user, nil, c)
assert.Equal(t, http.StatusConflict, recorder.Code)
var stored model.User
+10 -34
View File
@@ -11,10 +11,8 @@ import (
"github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/setting/config"
"github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)
func TestUpdateOptionRejectsInvalidTaskBillingExpressions(t *testing.T) {
@@ -103,42 +101,16 @@ func TestUpdateOptionRejectsUsageExpressionWithoutTaskPlugin(t *testing.T) {
assert.Contains(t, recorder.Body.String(), "no task plugin usage schema")
}
func setupBillingAliasOptionDB(t *testing.T) {
t.Helper()
previousDB := model.DB
previousLogDB := model.LOG_DB
previousType := common.MainDatabaseType()
previousCache := common.MemoryCacheEnabled
previousMap := common.OptionMap
previousRedis := common.RedisEnabled
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
require.NoError(t, database.AutoMigrate(&model.Channel{}, &model.Option{}, &model.Log{}, &model.AuditLog{}, &model.User{}))
model.DB = database
model.LOG_DB = database
common.SetMainDatabaseType(common.DatabaseTypeSQLite)
common.MemoryCacheEnabled = false
common.RedisEnabled = false
common.OptionMap = map[string]string{}
t.Cleanup(func() {
model.DB = previousDB
model.LOG_DB = previousLogDB
common.SetMainDatabaseType(previousType)
common.MemoryCacheEnabled = previousCache
common.OptionMap = previousMap
common.RedisEnabled = previousRedis
model.InitChannelCache()
})
}
func TestUpdateOptionAliasBillingExprUsesPluginSchema(t *testing.T) {
setupBillingAliasOptionDB(t)
database := modelManagementDB(t, "sqlite", "")
require.NoError(t, database.AutoMigrate(&model.Log{}))
const pluginKey = "billing-alias-probe"
source := `
export const meta = {
apiVersion: 1, key: "billing-alias-probe", name: "Billing Alias Probe", version: "1.0.0", author: {name: "Test"},
models: ["declared-model"], fetchMode: "per_task",
usageSchema: {seconds: {type: "number", unit: "second"}}
usageSchema: {seconds: {type: "number", unit: "second"}, image_count: {type: "number", unit: "count"}},
usageProfiles: [{models: ["declared-model"], schema: {seconds: {type: "number", unit: "second"}}}]
};
export function buildSubmitRequest() { return {}; }
export function parseSubmitResponse() { return {}; }
@@ -191,10 +163,14 @@ export function parseTaskResult() { return {}; }
assert.Equal(t, http.StatusOK, accepted.Code)
assert.Contains(t, accepted.Body.String(), `"success":true`)
rejectedKey := putExpr("alias-model", `u("clips")`)
rejectedKey := putExpr("alias-model", `u("image_count")`)
assert.Equal(t, http.StatusOK, rejectedKey.Code)
assert.Contains(t, rejectedKey.Body.String(), `"success":false`)
assert.Contains(t, rejectedKey.Body.String(), `usage key \"clips\" is not declared`)
assert.Contains(t, rejectedKey.Body.String(), `usage key \"image_count\" is not declared`)
rejectedDeclared := putExpr("declared-model", `u("image_count")`)
assert.Contains(t, rejectedDeclared.Body.String(), `"success":false`)
assert.Contains(t, rejectedDeclared.Body.String(), `usage key \"image_count\" is not declared`)
unresolvable := putExpr("unknown-alias-model", `u("seconds")`)
assert.Equal(t, http.StatusOK, unresolvable.Code)
+62 -11
View File
@@ -275,6 +275,67 @@ func updateChannelSiliconFlowBalance(channel *model.Channel) (float64, error) {
return balance, nil
}
func getDeepSeekBalanceUSD(response DeepSeekUsageResponse, usdExchangeRate float64) (float64, error) {
var usdBalance *string
var cnyBalance *string
for i := range response.BalanceInfos {
balanceInfo := &response.BalanceInfos[i]
switch balanceInfo.Currency {
case "USD":
if usdBalance == nil {
usdBalance = &balanceInfo.TotalBalance
}
case "CNY":
if cnyBalance == nil {
cnyBalance = &balanceInfo.TotalBalance
}
}
}
if usdBalance != nil {
balance, err := strconv.ParseFloat(*usdBalance, 64)
if err != nil {
return 0, err
}
if math.IsNaN(balance) || math.IsInf(balance, 0) {
return 0, errors.New("USD balance must be finite")
}
if balance < 0 {
return 0, errors.New("USD balance must be non-negative")
}
return balance, nil
}
if cnyBalance == nil {
return 0, errors.New("currency USD or CNY not found")
}
if math.IsNaN(usdExchangeRate) || math.IsInf(usdExchangeRate, 0) {
return 0, errors.New("USD exchange rate must be finite")
}
if usdExchangeRate <= 0 {
return 0, errors.New("USD exchange rate must be greater than zero")
}
balanceCNY, err := strconv.ParseFloat(*cnyBalance, 64)
if err != nil {
return 0, err
}
if math.IsNaN(balanceCNY) || math.IsInf(balanceCNY, 0) {
return 0, errors.New("CNY balance must be finite")
}
if balanceCNY < 0 {
return 0, errors.New("CNY balance must be non-negative")
}
balanceUSD := decimal.NewFromFloat(balanceCNY).Div(decimal.NewFromFloat(usdExchangeRate)).InexactFloat64()
if math.IsNaN(balanceUSD) || math.IsInf(balanceUSD, 0) {
return 0, errors.New("converted USD balance must be finite")
}
if balanceUSD < 0 {
return 0, errors.New("converted USD balance must be non-negative")
}
return balanceUSD, nil
}
func updateChannelDeepSeekBalance(channel *model.Channel) (float64, error) {
url := "https://api.deepseek.com/user/balance"
body, err := GetResponseBody("GET", url, channel, GetAuthHeader(channel.Key))
@@ -286,17 +347,7 @@ func updateChannelDeepSeekBalance(channel *model.Channel) (float64, error) {
if err != nil {
return 0, err
}
index := -1
for i, balanceInfo := range response.BalanceInfos {
if balanceInfo.Currency == "CNY" {
index = i
break
}
}
if index == -1 {
return 0, errors.New("currency CNY not found")
}
balance, err := strconv.ParseFloat(response.BalanceInfos[index].TotalBalance, 64)
balance, err := getDeepSeekBalanceUSD(response, operation_setting.USDExchangeRate)
if err != nil {
return 0, err
}
+141 -5
View File
@@ -73,6 +73,68 @@ func clearChannelInfo(channel *model.Channel) {
}
}
func channelIDsFromChannels(channels []*model.Channel) []int {
ids := make([]int, 0, len(channels))
for _, channel := range channels {
if channel != nil && channel.Id > 0 {
ids = append(ids, channel.Id)
}
}
return ids
}
func closeActiveChannelWebSockets(channelIDs []int) {
service.CloseActiveWebSocketsForChannels(channelIDs, service.ChannelDisabledCloseReason)
}
func hasEnabledMultiKey(channel *model.Channel) bool {
if channel == nil || !channel.ChannelInfo.IsMultiKey {
return true
}
keys := channel.GetKeys()
if len(keys) == 0 {
return false
}
for i := range keys {
if channel.ChannelInfo.MultiKeyStatusList == nil {
return true
}
if status, ok := channel.ChannelInfo.MultiKeyStatusList[i]; !ok || status == common.ChannelStatusEnabled {
return true
}
}
return false
}
func disableMultiKeyChannelIfUnavailable(channel *model.Channel) bool {
if channel == nil || !channel.ChannelInfo.IsMultiKey || hasEnabledMultiKey(channel) {
return false
}
if channel.Status != common.ChannelStatusEnabled {
return true
}
channel.Status = common.ChannelStatusManuallyDisabled
info := channel.GetOtherInfo()
info["status_reason"] = model.ChannelStatusReasonAllKeysDisabled
info["status_time"] = common.GetTimestamp()
channel.SetOtherInfo(info)
return true
}
func restoreMultiKeyChannelIfAvailable(channel *model.Channel) {
if channel.Status != common.ChannelStatusManuallyDisabled || !hasEnabledMultiKey(channel) {
return
}
info := channel.GetOtherInfo()
if info["status_reason"] != model.ChannelStatusReasonAllKeysDisabled {
return
}
channel.Status = common.ChannelStatusEnabled
info["status_reason"] = ""
info["status_time"] = common.GetTimestamp()
channel.SetOtherInfo(info)
}
func applyChannelStatusFilter(query *gorm.DB, statusFilter int) *gorm.DB {
if statusFilter == common.ChannelStatusEnabled {
return query.Where("status = ?", common.ChannelStatusEnabled)
@@ -94,11 +156,38 @@ func buildChannelListQuery(group string, statusFilter int, typeFilter int) *gorm
}
func GetChannelOps(c *gin.Context) {
snapshot := model.CurrentRequestPolicy()
automaticDisable, source := snapshot.AutoDisable, "global"
if value, present := c.GetQuery("auto_ban"); present {
autoBan, err := strconv.ParseBool(value)
if err != nil {
common.ApiError(c, err)
return
}
automaticDisable = snapshot.AutoDisable && autoBan
if snapshot.AutoDisable {
source = "global_and_channel"
if !autoBan {
source = "channel"
}
}
}
common.ApiSuccess(c, gin.H{
"retry_times": common.RetryTimes,
"retry_times": snapshot.RetryTimes,
"request_policy": gin.H{"automatic_disable": automaticDisable, "source": source},
})
}
func GetChannelDefaultBaseURLs(c *gin.Context) {
baseURLs := make(map[int]string)
for channelType, baseURL := range constant.ChannelBaseURLs {
if baseURL != "" {
baseURLs[channelType] = baseURL
}
}
common.ApiSuccess(c, baseURLs)
}
func GetAllChannels(c *gin.Context) {
pageInfo := common.GetPageQuery(c)
channelData := make([]*model.Channel, 0)
@@ -485,6 +574,12 @@ func validateChannel(channel *model.Channel, isAdd bool) error {
if channel.Type == constant.ChannelTypeNewAPI && strings.TrimSpace(channel.GetBaseURL()) == "" {
return fmt.Errorf("New API channel base URL cannot be empty")
}
if channel.Type == constant.ChannelTypeVLLM && strings.TrimSpace(channel.GetBaseURL()) == "" {
return fmt.Errorf("vLLM channel base URL cannot be empty")
}
if channel.Type == constant.ChannelTypeSGLang && strings.TrimSpace(channel.GetBaseURL()) == "" {
return fmt.Errorf("SGLang channel base URL cannot be empty")
}
// 如果是添加操作,检查 channel 和 key 是否为空
if isAdd {
@@ -757,6 +852,7 @@ func DeleteChannel(c *gin.Context) {
"id": id,
"name": channelName,
})
closeActiveChannelWebSockets([]int{id})
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "",
@@ -765,6 +861,13 @@ func DeleteChannel(c *gin.Context) {
}
func DeleteDisabledChannel(c *gin.Context) {
var ids []int
if err := model.DB.Model(&model.Channel{}).
Where("status = ? or status = ?", common.ChannelStatusAutoDisabled, common.ChannelStatusManuallyDisabled).
Pluck("id", &ids).Error; err != nil {
common.ApiError(c, err)
return
}
rows, err := model.DeleteDisabledChannel()
if err != nil {
common.ApiError(c, err)
@@ -777,6 +880,7 @@ func DeleteDisabledChannel(c *gin.Context) {
recordManageAudit(c, "channel.delete_disabled", map[string]any{
"count": rows,
})
closeActiveChannelWebSockets(ids)
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "",
@@ -807,6 +911,12 @@ func DisableTagChannels(c *gin.Context) {
})
return
}
channels, err := model.GetChannelsByTag(channelTag.Tag, false, false)
if err != nil {
common.ApiError(c, err)
return
}
ids := channelIDsFromChannels(channels)
err = model.DisableChannelByTag(channelTag.Tag)
if err != nil {
common.ApiError(c, err)
@@ -816,6 +926,7 @@ func DisableTagChannels(c *gin.Context) {
recordManageAudit(c, "channel.tag_disable", map[string]any{
"tag": channelTag.Tag,
})
closeActiveChannelWebSockets(ids)
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "",
@@ -936,6 +1047,7 @@ func DeleteChannelBatch(c *gin.Context) {
recordManageAudit(c, "channel.delete_batch", map[string]any{
"count": deletedCount,
})
closeActiveChannelWebSockets(channelBatch.Ids)
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "",
@@ -1170,6 +1282,9 @@ func UpdateChannelStatus(c *gin.Context) {
changed := model.UpdateChannelStatus(id, "", req.Status, "manual operation")
if changed {
model.InitChannelCache()
if req.Status != common.ChannelStatusEnabled {
closeActiveChannelWebSockets([]int{id})
}
}
recordManageAudit(c, "channel.status_update", map[string]any{
"id": id,
@@ -1190,14 +1305,21 @@ func BatchUpdateChannelStatus(c *gin.Context) {
return
}
changedCount := 0
var disabledIDs []int
for _, id := range req.Ids {
if model.UpdateChannelStatus(id, "", req.Status, "manual batch operation") {
changedCount++
if req.Status != common.ChannelStatusEnabled {
disabledIDs = append(disabledIDs, id)
}
}
}
if changedCount > 0 {
model.InitChannelCache()
}
if len(disabledIDs) > 0 {
closeActiveChannelWebSockets(disabledIDs)
}
recordManageAudit(c, "channel.status_update_batch", map[string]any{
"count": changedCount,
"total": len(req.Ids),
@@ -1714,13 +1836,16 @@ func ManageMultiKeys(c *gin.Context) {
channel.ChannelInfo.MultiKeyStatusList[keyIndex] = 2 // disabled
shouldCloseWebSocket := disableMultiKeyChannelIfUnavailable(channel)
err = channel.Update()
if err != nil {
common.ApiError(c, err)
return
}
model.InitChannelCache()
if shouldCloseWebSocket {
closeActiveChannelWebSockets([]int{channel.Id})
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "密钥已禁用",
@@ -1755,6 +1880,7 @@ func ManageMultiKeys(c *gin.Context) {
if channel.ChannelInfo.MultiKeyDisabledReason != nil {
delete(channel.ChannelInfo.MultiKeyDisabledReason, keyIndex)
}
restoreMultiKeyChannelIfAvailable(channel)
err = channel.Update()
if err != nil {
@@ -1779,6 +1905,7 @@ func ManageMultiKeys(c *gin.Context) {
channel.ChannelInfo.MultiKeyStatusList = make(map[int]int)
channel.ChannelInfo.MultiKeyDisabledTime = make(map[int]int64)
channel.ChannelInfo.MultiKeyDisabledReason = make(map[int]string)
restoreMultiKeyChannelIfAvailable(channel)
err = channel.Update()
if err != nil {
@@ -1827,13 +1954,16 @@ func ManageMultiKeys(c *gin.Context) {
return
}
shouldCloseWebSocket := disableMultiKeyChannelIfUnavailable(channel)
err = channel.Update()
if err != nil {
common.ApiError(c, err)
return
}
model.InitChannelCache()
if shouldCloseWebSocket {
closeActiveChannelWebSockets([]int{channel.Id})
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": fmt.Sprintf("已禁用 %d 个密钥", disabledCount),
@@ -1907,13 +2037,16 @@ func ManageMultiKeys(c *gin.Context) {
channel.ChannelInfo.MultiKeyDisabledTime = newDisabledTime
channel.ChannelInfo.MultiKeyDisabledReason = newDisabledReason
shouldCloseWebSocket := disableMultiKeyChannelIfUnavailable(channel)
err = channel.Update()
if err != nil {
common.ApiError(c, err)
return
}
model.InitChannelCache()
if shouldCloseWebSocket {
closeActiveChannelWebSockets([]int{channel.Id})
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "密钥已删除",
@@ -1975,13 +2108,16 @@ func ManageMultiKeys(c *gin.Context) {
channel.ChannelInfo.MultiKeyDisabledTime = newDisabledTime
channel.ChannelInfo.MultiKeyDisabledReason = newDisabledReason
shouldCloseWebSocket := disableMultiKeyChannelIfUnavailable(channel)
err = channel.Update()
if err != nil {
common.ApiError(c, err)
return
}
model.InitChannelCache()
if shouldCloseWebSocket {
closeActiveChannelWebSockets([]int{channel.Id})
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": fmt.Sprintf("已删除 %d 个自动禁用的密钥", deletedCount),
+116
View File
@@ -0,0 +1,116 @@
package controller
import (
"math"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestGetDeepSeekBalanceUSD(t *testing.T) {
tests := []struct {
name string
responseJSON string
usdExchangeRate float64
want float64
wantErrContains string
}{
{
name: "prefers USD when USD precedes CNY",
responseJSON: `{"balance_infos":[{"currency":"USD","total_balance":"12.50"},{"currency":"CNY","total_balance":"73.00"}]}`,
usdExchangeRate: 7.3,
want: 12.5,
},
{
name: "prefers USD when CNY precedes USD",
responseJSON: `{"balance_infos":[{"currency":"CNY","total_balance":"73.00"},{"currency":"USD","total_balance":"12.50"}]}`,
usdExchangeRate: 7.3,
want: 12.5,
},
{
name: "converts CNY when USD is absent",
responseJSON: `{"balance_infos":[{"currency":"CNY","total_balance":"73.00"}]}`,
usdExchangeRate: 7.3,
want: 10,
},
{
name: "returns error when USD and CNY are absent",
responseJSON: `{"balance_infos":[{"currency":"EUR","total_balance":"10.00"}]}`,
usdExchangeRate: 7.3,
wantErrContains: "currency USD or CNY not found",
},
{
name: "returns USD parse error instead of falling back to CNY",
responseJSON: `{"balance_infos":[{"currency":"USD","total_balance":"invalid"},{"currency":"CNY","total_balance":"73.00"}]}`,
usdExchangeRate: 7.3,
wantErrContains: "invalid syntax",
},
{
name: "rejects NaN USD balance",
responseJSON: `{"balance_infos":[{"currency":"USD","total_balance":"NaN"}]}`,
usdExchangeRate: 7.3,
wantErrContains: "USD balance must be finite",
},
{
name: "rejects negative USD balance",
responseJSON: `{"balance_infos":[{"currency":"USD","total_balance":"-1.00"}]}`,
usdExchangeRate: 7.3,
wantErrContains: "USD balance must be non-negative",
},
{
name: "rejects positive infinity CNY balance",
responseJSON: `{"balance_infos":[{"currency":"CNY","total_balance":"+Inf"}]}`,
usdExchangeRate: 7.3,
wantErrContains: "CNY balance must be finite",
},
{
name: "rejects negative CNY balance",
responseJSON: `{"balance_infos":[{"currency":"CNY","total_balance":"-7.30"}]}`,
usdExchangeRate: 7.3,
wantErrContains: "CNY balance must be non-negative",
},
{
name: "returns error for non-positive CNY exchange rate",
responseJSON: `{"balance_infos":[{"currency":"CNY","total_balance":"73.00"}]}`,
usdExchangeRate: 0,
wantErrContains: "USD exchange rate must be greater than zero",
},
{
name: "rejects NaN CNY exchange rate",
responseJSON: `{"balance_infos":[{"currency":"CNY","total_balance":"73.00"}]}`,
usdExchangeRate: math.NaN(),
wantErrContains: "USD exchange rate must be finite",
},
{
name: "rejects positive infinity CNY exchange rate",
responseJSON: `{"balance_infos":[{"currency":"CNY","total_balance":"73.00"}]}`,
usdExchangeRate: math.Inf(1),
wantErrContains: "USD exchange rate must be finite",
},
{
name: "rejects CNY conversion overflow",
responseJSON: `{"balance_infos":[{"currency":"CNY","total_balance":"1.7976931348623157e+308"}]}`,
usdExchangeRate: math.SmallestNonzeroFloat64,
wantErrContains: "converted USD balance must be finite",
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
var response DeepSeekUsageResponse
require.NoError(t, common.Unmarshal([]byte(test.responseJSON), &response))
balance, err := getDeepSeekBalanceUSD(response, test.usdExchangeRate)
if test.wantErrContains != "" {
require.Error(t, err)
assert.Contains(t, err.Error(), test.wantErrContains)
return
}
require.NoError(t, err)
assert.InDelta(t, test.want, balance, 1e-12)
})
}
}
+242
View File
@@ -0,0 +1,242 @@
package controller
import (
"context"
"errors"
"io"
"math"
"net/http"
"net/url"
"strconv"
"strings"
"sync"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
"github.com/prometheus/common/expfmt"
)
type inferenceEndpointStatus struct {
Status int `json:"status"`
Error string `json:"error,omitempty"`
}
type inferenceModel struct {
ID string `json:"id"`
Root string `json:"root"`
MaxModelLen *int64 `json:"max_model_len"`
}
type inferenceMetric struct {
Name string `json:"name"`
Labels map[string]string `json:"labels"`
Value float64 `json:"value"`
}
type inferenceStatus struct {
SampledAt int64 `json:"sampled_at"`
Endpoints map[string]inferenceEndpointStatus `json:"endpoints"`
Version string `json:"version"`
Models []inferenceModel `json:"models"`
Metrics []inferenceMetric `json:"metrics"`
RawMetrics string `json:"raw_metrics"`
}
func GetVLLMChannelStatus(c *gin.Context) {
getInferenceChannelStatus(c, constant.ChannelTypeVLLM)
}
func GetSGLangChannelStatus(c *gin.Context) {
getInferenceChannelStatus(c, constant.ChannelTypeSGLang)
}
func getInferenceChannelStatus(c *gin.Context, channelType int) {
c.Header("Cache-Control", "no-store")
id, err := strconv.Atoi(c.Param("id"))
if err != nil || id <= 0 {
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "Invalid channel id"})
return
}
channel, err := model.GetChannelById(id, true)
if err != nil || channel == nil {
c.JSON(http.StatusNotFound, gin.H{"success": false, "message": "Channel not found"})
return
}
if channel.Type != channelType {
common.ApiError(c, errors.New("Channel type does not match the status endpoint"))
return
}
status, err := fetchInferenceStatus(c.Request.Context(), channel)
if err != nil {
common.ApiError(c, err)
return
}
common.ApiSuccess(c, status)
}
// fetchInferenceStatus reads only fixed paths on an operator-configured channel.
// Each endpoint can fail independently; upstream errors never echo credentials.
func fetchInferenceStatus(ctx context.Context, channel *model.Channel) (*inferenceStatus, error) {
if channel.Type != constant.ChannelTypeVLLM && channel.Type != constant.ChannelTypeSGLang {
return nil, errors.New("This operation is only supported for vLLM or SGLang channels")
}
baseURL := strings.TrimRight(strings.TrimSpace(channel.GetBaseURL()), "/")
parsedURL, err := url.Parse(baseURL)
if err != nil || parsedURL.Host == "" || parsedURL.User != nil || parsedURL.RawQuery != "" || parsedURL.Fragment != "" || (parsedURL.Scheme != "http" && parsedURL.Scheme != "https") {
return nil, errors.New("Invalid inference server address")
}
key, _, keyErr := channel.GetNextEnabledKey()
if keyErr != nil {
return nil, errors.New("No enabled channel key")
}
key = strings.TrimSpace(key)
headers, err := buildFetchModelsHeaders(channel, key)
if err != nil {
return nil, errors.New("Invalid channel header override")
}
if (key == "" || key == "EMPTY") && headers.Get("Authorization") == "Bearer "+key {
headers.Del("Authorization")
}
settings := channel.GetSetting()
client, err := service.GetHttpClientWithProxySettings(settings.Proxy, settings)
if err != nil {
return nil, errors.New("Invalid channel proxy")
}
// Do not mutate the shared client or forward channel credentials on redirects.
statusClient := *client
statusClient.CheckRedirect = func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }
ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
defer cancel()
versionPath, metricPrefix := "/version", "vllm:"
if channel.Type == constant.ChannelTypeSGLang {
versionPath, metricPrefix = "/server_info", "sglang:"
}
paths := []string{"/health", versionPath, "/v1/models", "/metrics"}
results := make([]struct {
inferenceEndpointStatus
body []byte
at int64
}, len(paths))
var wg sync.WaitGroup
for i, path := range paths {
wg.Go(func() {
result := &results[i]
req, err := http.NewRequestWithContext(ctx, http.MethodGet, baseURL+path, nil)
if err != nil {
result.Error = "request_failed"
return
}
req.Header = headers.Clone()
req.Host = headers.Get("Host")
resp, err := statusClient.Do(req)
if err != nil {
result.Error = "request_failed"
return
}
defer resp.Body.Close()
result.Status = resp.StatusCode
if resp.StatusCode != http.StatusOK {
result.Error = "http_error"
return
}
const maxResponseBytes = 4 << 20
body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes+1))
if err != nil {
result.Error = "request_failed"
return
}
if len(body) > maxResponseBytes {
result.Error = "response_too_large"
return
}
result.body = body
result.at = time.Now().UnixMilli()
})
}
wg.Wait()
if errors.Is(ctx.Err(), context.Canceled) {
return nil, context.Canceled
}
status := &inferenceStatus{
SampledAt: time.Now().UnixMilli(),
Endpoints: make(map[string]inferenceEndpointStatus, len(paths)),
Models: []inferenceModel{},
Metrics: []inferenceMetric{},
}
for i, path := range paths {
result := &results[i]
if result.Error == "" {
switch path {
case versionPath:
// /server_info includes the launch configuration. Decode only the
// version so API keys and other server settings never reach the browser.
var version struct {
Version string `json:"version"`
}
if err := common.Unmarshal(result.body, &version); err != nil || version.Version == "" {
result.Error = "invalid_response"
} else {
status.Version = version.Version
}
case "/v1/models":
var models struct {
Data []inferenceModel `json:"data"`
}
if err := common.Unmarshal(result.body, &models); err != nil || models.Data == nil {
result.Error = "invalid_response"
} else {
status.Models = models.Data
}
case "/metrics":
metrics, err := parseInferenceMetrics(string(result.body), metricPrefix)
if err != nil {
result.Error = "invalid_response"
} else {
status.Metrics = metrics
status.RawMetrics = string(result.body)
status.SampledAt = result.at
}
}
}
status.Endpoints[path] = result.inferenceEndpointStatus
}
return status, nil
}
func parseInferenceMetrics(raw string, prefix string) ([]inferenceMetric, error) {
var parser expfmt.TextParser
families, err := parser.TextToMetricFamilies(strings.NewReader(raw))
if err != nil {
return nil, err
}
metrics := []inferenceMetric{}
for name, family := range families {
if !strings.HasPrefix(name, prefix) && name != "process_start_time_seconds" {
continue
}
samples, err := expfmt.ExtractSamples(&expfmt.DecodeOptions{}, family)
if err != nil {
return nil, err
}
for _, sample := range samples {
name := string(sample.Metric["__name__"])
value := float64(sample.Value)
if strings.HasSuffix(name, "_bucket") || math.IsNaN(value) || math.IsInf(value, 0) || value < 0 {
continue
}
labels := make(map[string]string, len(sample.Metric)-1)
for label, value := range sample.Metric {
if label != "__name__" {
labels[string(label)] = string(value)
}
}
metrics = append(metrics, inferenceMetric{Name: name, Labels: labels, Value: value})
}
}
return metrics, nil
}
+196
View File
@@ -0,0 +1,196 @@
package controller
import (
"context"
"net/http"
"net/http/httptest"
"os"
"strings"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/service"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestVLLMStatusPartialFailureAndCredentials(t *testing.T) {
service.InitHttpClient()
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, http.MethodGet, r.Method)
assert.Equal(t, "Bearer test-vllm-key", r.Header.Get("Authorization"))
assert.Equal(t, "status-panel", r.Header.Get("X-Source"))
switch r.URL.Path {
case "/prefix/health":
w.WriteHeader(http.StatusOK)
case "/prefix/version":
_, _ = w.Write([]byte(`{"version":"0.25.2"}`))
case "/prefix/v1/models":
_, _ = w.Write([]byte(`{"data":[{"id":"served-model","root":"source/model","max_model_len":1048576}]}`))
case "/prefix/metrics":
http.Error(w, "test-vllm-key should not reach the browser", http.StatusForbidden)
default:
t.Errorf("unexpected path: %s", r.URL.Path)
w.WriteHeader(http.StatusNotFound)
}
}))
defer upstream.Close()
base := upstream.URL + "/prefix/"
override := `{"X-Source":"status-panel"}`
channel := &model.Channel{Type: constant.ChannelTypeVLLM, BaseURL: &base, Key: "test-vllm-key", HeaderOverride: &override}
status, err := fetchInferenceStatus(context.Background(), channel)
require.NoError(t, err)
assert.Equal(t, 200, status.Endpoints["/health"].Status)
assert.Equal(t, "0.25.2", status.Version)
require.Len(t, status.Models, 1)
assert.Equal(t, "source/model", status.Models[0].Root)
require.NotNil(t, status.Models[0].MaxModelLen)
assert.EqualValues(t, 1048576, *status.Models[0].MaxModelLen)
assert.Equal(t, inferenceEndpointStatus{Status: 403, Error: "http_error"}, status.Endpoints["/metrics"])
assert.Empty(t, status.RawMetrics)
assert.Empty(t, status.Metrics)
body, err := common.Marshal(status)
require.NoError(t, err)
assert.NotContains(t, string(body), channel.Key)
}
func TestVLLMStatusRejectsInvalidTargetsAndRedirects(t *testing.T) {
service.InitHttpClient()
for _, base := range []string{"", "file:///etc/passwd", "http://user:secret@localhost", "http://localhost?key=secret"} {
_, err := fetchInferenceStatus(context.Background(), &model.Channel{Type: constant.ChannelTypeVLLM, BaseURL: &base})
require.Error(t, err)
assert.NotContains(t, err.Error(), "secret")
}
_, err := fetchInferenceStatus(context.Background(), &model.Channel{Type: constant.ChannelTypeOllama})
require.Error(t, err)
target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
t.Error("redirect must not forward channel credentials")
}))
defer target.Close()
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Empty(t, r.Header.Get("Authorization"))
http.Redirect(w, r, target.URL, http.StatusFound)
}))
defer upstream.Close()
status, err := fetchInferenceStatus(context.Background(), &model.Channel{Type: constant.ChannelTypeVLLM, BaseURL: &upstream.URL, Key: "EMPTY"})
require.NoError(t, err)
for _, endpoint := range status.Endpoints {
assert.Equal(t, inferenceEndpointStatus{Status: 302, Error: "http_error"}, endpoint)
}
}
func TestVLLMStatusInvalidAndOversizedResponses(t *testing.T) {
service.InitHttpClient()
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/health":
w.WriteHeader(http.StatusServiceUnavailable)
case "/version":
_, _ = w.Write([]byte(`{"version":false}`))
case "/v1/models":
_, _ = w.Write([]byte(`{"data":[]}`))
case "/metrics":
_, _ = w.Write([]byte(strings.Repeat(" ", (4<<20)+1)))
}
}))
defer upstream.Close()
channel := &model.Channel{Type: constant.ChannelTypeVLLM, BaseURL: &upstream.URL, Key: "EMPTY"}
status, err := fetchInferenceStatus(context.Background(), channel)
require.NoError(t, err)
assert.Equal(t, 503, status.Endpoints["/health"].Status)
assert.Equal(t, "invalid_response", status.Endpoints["/version"].Error)
assert.Empty(t, status.Endpoints["/v1/models"].Error)
assert.Equal(t, "response_too_large", status.Endpoints["/metrics"].Error)
ctx, cancel := context.WithCancel(context.Background())
cancel()
_, err = fetchInferenceStatus(ctx, channel)
assert.ErrorIs(t, err, context.Canceled)
}
func TestParseVLLMMetricsPreservesSeriesAndHistogramTotals(t *testing.T) {
metrics, err := parseInferenceMetrics(`# TYPE vllm:num_requests_running gauge
vllm:num_requests_running{engine="0",model_name="model,a"} 0
vllm:num_requests_running{model_name="model,a",engine="1"} 2
# TYPE vllm:time_to_first_token_seconds histogram
vllm:time_to_first_token_seconds_bucket{le="+Inf"} 4
vllm:time_to_first_token_seconds_sum 10
vllm:time_to_first_token_seconds_count 4
# TYPE vllm:kv_cache_usage_perc gauge
vllm:kv_cache_usage_perc NaN
vllm:bad_value +Inf
# TYPE process_start_time_seconds gauge
process_start_time_seconds 100
unrelated_metric 123
`, "vllm:")
require.NoError(t, err)
assert.ElementsMatch(t, []inferenceMetric{
{Name: "vllm:num_requests_running", Labels: map[string]string{"engine": "0", "model_name": "model,a"}, Value: 0},
{Name: "vllm:num_requests_running", Labels: map[string]string{"engine": "1", "model_name": "model,a"}, Value: 2},
{Name: "vllm:time_to_first_token_seconds_sum", Labels: map[string]string{}, Value: 10},
{Name: "vllm:time_to_first_token_seconds_count", Labels: map[string]string{}, Value: 4},
{Name: "process_start_time_seconds", Labels: map[string]string{}, Value: 100},
}, metrics)
_, err = parseInferenceMetrics("vllm:broken{engine=oops} 1\n", "vllm:")
require.Error(t, err)
}
// Opt-in verification against an operator-provided instance; credentials stay in the environment.
func TestVLLMStatusLive(t *testing.T) {
base := os.Getenv("NEW_API_VLLM_TEST_URL")
if base == "" {
t.Skip("NEW_API_VLLM_TEST_URL is not configured")
}
service.InitHttpClient()
status, err := fetchInferenceStatus(context.Background(), &model.Channel{
Type: constant.ChannelTypeVLLM, BaseURL: &base, Key: os.Getenv("NEW_API_VLLM_TEST_KEY"),
})
require.NoError(t, err)
for path, endpoint := range status.Endpoints {
assert.Equal(t, 200, endpoint.Status, path)
assert.Empty(t, endpoint.Error, path)
}
assert.NotEmpty(t, status.Version)
assert.NotEmpty(t, status.Models)
assert.NotEmpty(t, status.Metrics)
assert.NotEmpty(t, status.RawMetrics)
}
func TestSGLangStatusUsesDocumentedEndpointsAndRedactsServerConfiguration(t *testing.T) {
service.InitHttpClient()
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, http.MethodGet, r.Method)
assert.Equal(t, "Bearer test-key", r.Header.Get("Authorization"))
switch r.URL.Path {
case "/health":
w.WriteHeader(http.StatusOK)
case "/server_info":
_, _ = w.Write([]byte(`{"version":"0.5.19","api_key":"server-secret","admin_api_key":"admin-secret","internal_states":[{"private":"private-state"}]}`))
case "/v1/models":
_, _ = w.Write([]byte(`{"data":[{"id":"served-model","root":"served-model","max_model_len":32768}]}`))
case "/metrics":
_, _ = w.Write([]byte("# TYPE sglang:num_running_reqs gauge\nsglang:num_running_reqs{dp_rank=\"0\"} 0\nsglang:num_running_reqs{dp_rank=\"1\"} 2\n# TYPE sglang:time_to_first_token_seconds histogram\nsglang:time_to_first_token_seconds_bucket{le=\"+Inf\"} 2\nsglang:time_to_first_token_seconds_sum 1.5\nsglang:time_to_first_token_seconds_count 2\nvllm:num_requests_running 9\n"))
default:
t.Errorf("unexpected SGLang path: %s", r.URL.Path)
w.WriteHeader(http.StatusNotFound)
}
}))
defer upstream.Close()
status, err := fetchInferenceStatus(context.Background(), &model.Channel{Type: constant.ChannelTypeSGLang, BaseURL: &upstream.URL, Key: "test-key"})
require.NoError(t, err)
assert.Equal(t, "0.5.19", status.Version)
require.Len(t, status.Models, 1)
assert.Equal(t, "served-model", status.Models[0].ID)
assert.ElementsMatch(t, []inferenceMetric{
{Name: "sglang:num_running_reqs", Labels: map[string]string{"dp_rank": "0"}, Value: 0},
{Name: "sglang:num_running_reqs", Labels: map[string]string{"dp_rank": "1"}, Value: 2},
{Name: "sglang:time_to_first_token_seconds_sum", Labels: map[string]string{}, Value: 1.5},
{Name: "sglang:time_to_first_token_seconds_count", Labels: map[string]string{}, Value: 2},
}, status.Metrics)
encoded, err := common.Marshal(status)
require.NoError(t, err)
for _, secret := range []string{"test-key", "server-secret", "admin-secret", "private-state", "api_key", "internal_states"} {
assert.NotContains(t, string(encoded), secret)
}
}
+111
View File
@@ -0,0 +1,111 @@
package controller
import (
"bytes"
"fmt"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestMultiKeyEnableRestoresOnlyExhaustedChannels(t *testing.T) {
previousDB, previousLogDB := model.DB, model.LOG_DB
previousType, previousLogType := common.MainDatabaseType(), common.LogDatabaseType()
previousMaster, previousCache, previousRedis, previousSQLite := common.IsMasterNode, common.MemoryCacheEnabled, common.RedisEnabled, common.SQLitePath
t.Cleanup(func() {
model.DB, model.LOG_DB = previousDB, previousLogDB
common.SetDatabaseTypes(previousType, previousLogType)
common.IsMasterNode, common.MemoryCacheEnabled, common.RedisEnabled, common.SQLitePath = previousMaster, previousCache, previousRedis, previousSQLite
})
t.Setenv("SQL_DSN", os.Getenv("TEST_CHANNEL_SQL_DSN"))
t.Setenv("LOG_SQL_DSN", "")
common.IsMasterNode, common.MemoryCacheEnabled, common.RedisEnabled = false, false, false
common.SQLitePath = filepath.Join(t.TempDir(), "channel.db")
require.NoError(t, model.InitDB())
database := model.DB
sqlDB, err := database.DB()
require.NoError(t, err)
t.Cleanup(func() { require.NoError(t, sqlDB.Close()) })
model.LOG_DB = database
common.SetLogDatabaseType(common.MainDatabaseType())
require.NoError(t, database.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.User{}, &model.Log{}, &model.AuditLog{}))
root := &model.User{Username: "multi-key-review-root", Role: common.RoleRootUser, Status: common.UserStatusEnabled}
require.NoError(t, database.Create(root).Error)
t.Cleanup(func() { require.NoError(t, database.Unscoped().Delete(root).Error) })
versionQuery := "SELECT VERSION()"
if common.UsingMainDatabase(common.DatabaseTypeSQLite) {
versionQuery = "SELECT sqlite_version()"
}
var version string
require.NoError(t, database.Raw(versionQuery).Scan(&version).Error)
t.Logf("database=%s version=%s", common.MainDatabaseType(), version)
for _, cacheEnabled := range []bool{false, true} {
for _, action := range []string{"enable_key", "enable_all_keys"} {
for _, tc := range []struct {
name string
initialStatus int
manualOverride string
wantStatus int
}{
{name: "key exhaustion restores", initialStatus: common.ChannelStatusEnabled, wantStatus: common.ChannelStatusEnabled},
{name: "manual disable is preserved", initialStatus: common.ChannelStatusManuallyDisabled, wantStatus: common.ChannelStatusManuallyDisabled},
{name: "manual disable after exhaustion is preserved", initialStatus: common.ChannelStatusEnabled, manualOverride: "status", wantStatus: common.ChannelStatusManuallyDisabled},
{name: "tag disable after exhaustion is preserved", initialStatus: common.ChannelStatusEnabled, manualOverride: "tag", wantStatus: common.ChannelStatusManuallyDisabled},
} {
t.Run(fmt.Sprintf("cache=%t/%s/%s", cacheEnabled, action, tc.name), func(t *testing.T) {
common.MemoryCacheEnabled = cacheEnabled
tag := t.Name()
channel := &model.Channel{Name: t.Name(), Type: 1, Key: "key-one\nkey-two", Status: tc.initialStatus, Models: "test-model", Group: "default", Tag: &tag,
ChannelInfo: model.ChannelInfo{IsMultiKey: true, MultiKeySize: 2, MultiKeyStatusList: map[int]int{1: common.ChannelStatusManuallyDisabled}},
}
require.NoError(t, channel.Insert())
t.Cleanup(func() {
require.NoError(t, channel.Delete())
model.InitChannelCache()
})
for _, operation := range []string{"disable_key", action} {
if operation == action {
if tc.manualOverride == "status" {
model.UpdateChannelStatus(channel.Id, "", common.ChannelStatusManuallyDisabled, "manual operation")
} else if tc.manualOverride == "tag" {
require.NoError(t, model.DisableChannelByTag(tag))
}
}
payload, err := common.Marshal(MultiKeyManageRequest{ChannelId: channel.Id, Action: operation, KeyIndex: common.GetPointer(0)})
require.NoError(t, err)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Set("id", root.Id)
c.Set("role", common.RoleRootUser)
c.Request = httptest.NewRequest(http.MethodPost, "/api/channel/multi_key", bytes.NewReader(payload))
c.Request.Header.Set("Content-Type", "application/json")
ManageMultiKeys(c)
var result struct {
Success bool `json:"success"`
}
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &result))
require.True(t, result.Success, recorder.Body.String())
}
loaded, err := model.GetChannelById(channel.Id, true)
require.NoError(t, err)
assert.Equal(t, tc.wantStatus, loaded.Status)
assert.NotContains(t, loaded.ChannelInfo.MultiKeyStatusList, 0)
assert.NotContains(t, loaded.ChannelInfo.MultiKeyDisabledReason, 0)
assert.NotContains(t, loaded.ChannelInfo.MultiKeyDisabledTime, 0)
var ability model.Ability
require.NoError(t, database.Where("channel_id = ?", channel.Id).First(&ability).Error)
assert.Equal(t, tc.wantStatus == common.ChannelStatusEnabled, ability.Enabled)
})
}
}
}
}
+132 -8
View File
@@ -2,13 +2,25 @@ package controller
import (
"errors"
"fmt"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/middleware"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/setting"
"github.com/QuantumNous/new-api/setting/operation_setting"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"net/http"
"net/http/httptest"
"os"
"slices"
"strings"
"testing"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
@@ -18,7 +30,7 @@ func TestShouldRetryHonorsPinRetryMode(t *testing.T) {
openaiErr := types.NewOpenAIError(errors.New("upstream"), types.ErrorCodeBadResponseStatusCode, http.StatusInternalServerError)
c := newPinRetryContext()
assert.True(t, shouldRetry(c, openaiErr, 1))
assert.True(t, service.ShouldRetryRelayError(c, openaiErr, 1))
origin := newPinRetryContext()
service.GetChannelConstraints(origin).AddPin(dto.ChannelPin{
@@ -27,7 +39,7 @@ func TestShouldRetryHonorsPinRetryMode(t *testing.T) {
Rank: dto.PinRankOriginTask,
RetryMode: dto.PinRetrySameChannel,
})
assert.True(t, shouldRetry(origin, openaiErr, 1), "origin pin retries on the same channel")
assert.True(t, service.ShouldRetryRelayError(origin, openaiErr, 1), "origin pin retries on the same channel")
token := newPinRetryContext()
service.GetChannelConstraints(token).AddPin(dto.ChannelPin{
@@ -36,14 +48,14 @@ func TestShouldRetryHonorsPinRetryMode(t *testing.T) {
Rank: dto.PinRankToken,
RetryMode: dto.PinRetrySingleAttempt,
})
assert.False(t, shouldRetry(token, openaiErr, 1), "token pin suppresses retry")
assert.Equal(t, service.PolicyDecision{Action: "stop", Reason: "pinned_channel", Source: "channel_constraint"}, service.DecideRelayRetry(token, openaiErr, 1), "token pin suppresses retry")
}
func TestShouldRetryTaskRelayHonorsPinRetryMode(t *testing.T) {
taskErr := &dto.TaskError{StatusCode: http.StatusInternalServerError}
c := newPinRetryContext()
assert.True(t, shouldRetryTaskRelay(c, 1, taskErr, 1))
assert.Equal(t, "retry", decideTaskRetry(c, taskErr, 1).Action)
origin := newPinRetryContext()
service.GetChannelConstraints(origin).AddPin(dto.ChannelPin{
@@ -52,7 +64,7 @@ func TestShouldRetryTaskRelayHonorsPinRetryMode(t *testing.T) {
Rank: dto.PinRankOriginTask,
RetryMode: dto.PinRetrySameChannel,
})
assert.True(t, shouldRetryTaskRelay(origin, 2, taskErr, 1))
assert.Equal(t, "retry", decideTaskRetry(origin, taskErr, 1).Action)
token := newPinRetryContext()
service.GetChannelConstraints(token).AddPin(dto.ChannelPin{
@@ -61,7 +73,7 @@ func TestShouldRetryTaskRelayHonorsPinRetryMode(t *testing.T) {
Rank: dto.PinRankToken,
RetryMode: dto.PinRetrySingleAttempt,
})
assert.False(t, shouldRetryTaskRelay(token, 1, taskErr, 1))
assert.Equal(t, service.PolicyDecision{Action: "stop", Reason: "pinned_channel", Source: "channel_constraint"}, decideTaskRetry(token, taskErr, 1))
}
func TestSameChannelPinsMergeToStricterRetryMode(t *testing.T) {
@@ -84,7 +96,7 @@ func TestSameChannelPinsMergeToStricterRetryMode(t *testing.T) {
assert.Equal(t, 7, pin.ChannelId)
assert.Equal(t, dto.PinRetrySingleAttempt, pin.RetryMode)
assert.Empty(t, overridden)
assert.False(t, shouldRetry(c, types.NewOpenAIError(errors.New("upstream"), types.ErrorCodeBadResponseStatusCode, http.StatusInternalServerError), 1))
assert.False(t, service.ShouldRetryRelayError(c, types.NewOpenAIError(errors.New("upstream"), types.ErrorCodeBadResponseStatusCode, http.StatusInternalServerError), 1))
}
func newPinRetryContext() *gin.Context {
@@ -93,3 +105,115 @@ func newPinRetryContext() *gin.Context {
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
return c
}
func TestRequestPolicyConfigReturnsSettingsWithoutMigration(t *testing.T) {
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
GetRequestPolicy(ctx)
assert.Equal(t, http.StatusOK, recorder.Code)
var response struct {
Success bool `json:"success"`
Data map[string]any `json:"data"`
}
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
assert.True(t, response.Success)
assert.Contains(t, response.Data, "options")
assert.NotContains(t, response.Data, "migration")
assert.NotContains(t, response.Data, "differences")
}
func TestRequestPolicyRoutingDatabaseMatrix(t *testing.T) {
require.NoError(t, i18n.Init())
for _, dialect := range []struct{ kind, env string }{{"sqlite", ""}, {"mysql", "TEST_MYSQL_DSN"}, {"postgres", "TEST_POSTGRES_DSN"}} {
t.Run(dialect.kind, func(t *testing.T) {
dsn := ""
if dialect.env != "" {
dsn = os.Getenv(dialect.env)
if dsn == "" {
t.Skipf("%s not configured", dialect.env)
}
}
db := modelManagementDB(t, dialect.kind, dsn)
previousGroups := setting.UserUsableGroups2JSONString()
previousRatios, err := common.Marshal(ratio_setting.GetGroupRatioCopy())
require.NoError(t, err)
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default"}`))
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1}`))
t.Cleanup(func() {
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(previousGroups))
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(string(previousRatios)))
})
// Channel 1 holds the session binding but is no longer usable.
channels := []model.Channel{
{Id: 1, Name: "A", Type: 1, Key: "test-only", Status: common.ChannelStatusAutoDisabled, Models: "policy-test", Group: "default", Priority: common.GetPointer(int64(10))},
{Id: 2, Name: "B", Type: 1, Key: "test-only", Status: common.ChannelStatusEnabled, Models: "policy-test", Group: "default", Priority: common.GetPointer(int64(5))},
}
for i := range channels {
require.NoError(t, db.Create(&channels[i]).Error)
require.NoError(t, channels[i].AddAbilities(db))
}
affinity := operation_setting.GetChannelAffinitySetting()
previousAffinity := *affinity
t.Cleanup(func() { *affinity = previousAffinity })
for _, cached := range []bool{false, true} {
for _, keep := range []bool{false, true} {
for _, tc := range []struct {
globalMode, ruleMode string
blocked bool
}{
{"strict", "inherit", true},
{"prefer", "strict", true},
{"prefer", "inherit", false},
{"strict", "prefer", false},
} {
t.Run(fmt.Sprintf("cache=%t/keep=%t/global=%s/rule=%s", cached, keep, tc.globalMode, tc.ruleMode), func(t *testing.T) {
common.MemoryCacheEnabled = cached
model.InitChannelCache()
snapshot, err := model.BuildRequestPolicy(map[string]string{
"channel_affinity_setting.enabled": "true",
"channel_affinity_setting.session_mode": tc.globalMode,
"channel_affinity_setting.keep_on_channel_disabled": fmt.Sprint(keep),
"channel_affinity_setting.rules": fmt.Sprintf(`[{"name":"session","model_regex":[".*"],"key_sources":[{"type":"request_header","key":"X-Session"}],"session_mode":%q}]`, tc.ruleMode),
})
require.NoError(t, err)
*affinity = snapshot.Affinity
seed := newPinRetryContext()
seed.Request.Header.Set("X-Session", t.Name())
_, found := service.GetPreferredChannelByAffinity(seed, "policy-test", "default")
require.False(t, found)
seed.Set("channel_id", 1)
service.RecordChannelAffinity(seed, 1)
t.Cleanup(func() { service.ClearCurrentChannelAffinityCache(seed) })
bound, found := service.GetPreferredChannelByAffinity(seed, "policy-test", "default")
require.True(t, found)
require.Equal(t, 1, bound)
request := newPinRetryContext()
request.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"policy-test"}`))
request.Request.Header.Set("Content-Type", "application/json")
request.Request.Header.Set("X-Session", t.Name())
common.SetContextKey(request, constant.ContextKeyUsingGroup, "default")
middleware.Distribute()(request)
events := service.RequestPolicy(request).Events()
assert.True(t, slices.ContainsFunc(events, func(event service.PolicyEvent) bool { return event.Decision.Reason == "session_rule_matched" }), "decision events are recorded for every request")
bound, found = service.GetPreferredChannelByAffinity(seed, "policy-test", "default")
if tc.blocked {
assert.Equal(t, http.StatusServiceUnavailable, request.Writer.Status())
assert.True(t, request.IsAborted())
assert.Equal(t, keep, found, "binding retention is independent of strict request handling")
if found {
assert.Equal(t, 1, bound)
}
return
}
assert.False(t, request.IsAborted())
assert.Equal(t, 2, common.GetContextKeyInt(request, constant.ContextKeyChannelId), "prefer falls back to the next eligible channel")
require.True(t, found)
assert.Equal(t, 2, bound, "a successful fallback rebinds the session")
})
}
}
}
})
}
}
+27
View File
@@ -23,6 +23,33 @@ import (
"github.com/stretchr/testify/require"
)
func TestGetChannelDefaultBaseURLsUsesBuiltInDefaults(t *testing.T) {
originalBaseURLs := constant.ChannelBaseURLs
constant.ChannelBaseURLs = append([]string(nil), originalBaseURLs...)
constant.ChannelBaseURLs[constant.ChannelTypeDeepSeek] = "https://deepseek.server.example"
t.Cleanup(func() {
constant.ChannelBaseURLs = originalBaseURLs
})
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodGet, "/api/channel/default_base_urls", nil)
GetChannelDefaultBaseURLs(c)
require.Equal(t, http.StatusOK, recorder.Code)
var response struct {
Success bool `json:"success"`
Data map[int]string `json:"data"`
}
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
require.True(t, response.Success)
assert.Equal(t, "https://deepseek.server.example", response.Data[constant.ChannelTypeDeepSeek])
assert.Equal(t, "https://api.openai.com", response.Data[constant.ChannelTypeOpenAI])
assert.NotContains(t, response.Data, constant.ChannelTypeAzure)
assert.NotContains(t, response.Data, constant.ChannelTypeNewAPI)
assert.NotContains(t, response.Data, constant.ChannelTypeTaskPlugin)
}
func TestValidateChannelProxy(t *testing.T) {
tests := []struct {
name string
+6 -3
View File
@@ -398,7 +398,7 @@ func fetchChannelUpstreamModelIDs(channel *model.Channel) ([]string, error) {
return normalizeModelNames(models), nil
}
if channel.Type == constant.ChannelTypeAdvancedCustom {
if constant.IsAdvancedCustomChannel(channel.Type) {
return fetchAdvancedCustomUpstreamModelIDs(channel, baseURL)
}
@@ -417,10 +417,13 @@ func fetchChannelUpstreamModelIDs(channel *model.Channel) ([]string, error) {
url = fmt.Sprintf("%s/api/paas/v4/models", baseURL)
}
case constant.ChannelTypeVolcEngine:
// 火山方舟 OpenAI 兼容 API 根路径是 /api/v3(chat/embeddings 均为
// {base}/api/v3/...,见 relay/channel/volcengine 的 GetRequestURL),
// 不存在 /v1/models 端点。原拼接会导致「获取模型列表」固定 404 失败。
if plan, ok := constant.ChannelSpecialBases[baseURL]; ok && plan.OpenAIBaseURL != "" {
url = fmt.Sprintf("%s/v1/models", plan.OpenAIBaseURL)
} else {
url = fmt.Sprintf("%s/v1/models", baseURL)
url = fmt.Sprintf("%s/api/v3/models", baseURL)
}
case constant.ChannelTypeMoonshot:
if plan, ok := constant.ChannelSpecialBases[baseURL]; ok && plan.OpenAIBaseURL != "" {
@@ -473,7 +476,7 @@ func fetchAdvancedCustomUpstreamModelIDs(channel *model.Channel, baseURL string)
RelayMode: relayconstant.RelayModeUnknown,
RequestURLPath: dto.AdvancedCustomModelListPath,
ChannelMeta: &relaycommon.ChannelMeta{
ChannelType: constant.ChannelTypeAdvancedCustom,
ChannelType: channel.Type,
ChannelBaseUrl: baseURL,
ApiKey: key,
ChannelOtherSettings: channel.GetOtherSettings(),
+498
View File
@@ -0,0 +1,498 @@
package controller
import (
"bytes"
"context"
"errors"
"fmt"
"io"
"net/http"
"net/http/httptest"
"regexp"
"strings"
"sync"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/middleware"
"github.com/QuantumNous/new-api/model"
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
"github.com/QuantumNous/new-api/relay"
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/service"
"github.com/gin-gonic/gin"
"github.com/tidwall/gjson"
)
var detectionKeyPattern = regexp.MustCompile(`^[a-zA-Z0-9_-]{16,96}$`)
var detectionCalls sync.Map
type detectionRequest struct {
Group string `json:"group"`
Kind string `json:"kind"`
Key string `json:"key"`
PriceVersion string `json:"price_version"`
}
func detectionInput(c *gin.Context) (detectionRequest, *model.User, service.DetectionPrice, int, error) {
var input detectionRequest
if err := c.ShouldBindJSON(&input); err != nil {
return input, nil, service.DetectionPrice{}, 0, errors.New("invalid detection request")
}
if input.Kind != "pelican" && input.Kind != "fingerprint" {
return input, nil, service.DetectionPrice{}, 0, errors.New("invalid detection type")
}
user, err := model.GetUserById(c.GetInt("id"), false)
if err != nil {
return input, nil, service.DetectionPrice{}, 0, err
}
if !service.IsUserSelectableGroup(user.Group, input.Group) {
return input, nil, service.DetectionPrice{}, 0, errors.New("group unavailable")
}
price, reserve, err := service.DetectionPricing(user.Group, input.Group, input.Kind)
return input, user, price, reserve, err
}
func EstimateDetection(c *gin.Context) {
input, _, price, reserve, err := detectionInput(c)
if err != nil {
common.ApiError(c, err)
return
}
encoded, _ := common.Marshal(price)
calls := 1
if input.Kind == "fingerprint" {
calls = 3
}
c.JSON(200, gin.H{"success": true, "data": gin.H{"reserved": reserve, "ratio": price.Snapshot.GroupRatio, "calls": calls, "price_version": model.DetectionSecretHash(string(encoded))}})
}
func CreateUserDetection(c *gin.Context) {
input, user, price, reserve, err := detectionInput(c)
if err != nil {
common.ApiError(c, err)
return
}
if !detectionKeyPattern.MatchString(input.Key) {
common.ApiError(c, errors.New("invalid idempotency key"))
return
}
var previous []model.DetectionRun
key := fmt.Sprintf("user:%d:%s", user.Id, input.Key)
if err = model.DB.Where("request_key = ?", key).Find(&previous).Error; err != nil {
common.ApiError(c, err)
return
}
if len(previous) > 0 {
if previous[0].GroupName != input.Group || previous[0].Kind != input.Kind {
common.ApiError(c, errors.New("idempotency key conflict"))
return
}
c.JSON(200, gin.H{"success": true, "data": gin.H{"id": previous[0].ID}})
return
}
encoded, _ := common.Marshal(price)
if input.PriceVersion != model.DetectionSecretHash(string(encoded)) {
common.ApiError(c, errors.New("pricing changed; request a new estimate"))
return
}
now := time.Now()
price.Request.At = &now
encoded, _ = common.Marshal(price)
channel, err := model.GetRandomSatisfiedChannel(input.Group, fingerprint.Model, 0, nil)
if err != nil || channel == nil || !service.ChannelSupportsAstra(channel) {
common.ApiError(c, errors.New("no Astra channel available"))
return
}
id, err := common.GenerateRandomCharsKey(32)
if err != nil {
common.ApiError(c, err)
return
}
run := &model.DetectionRun{ID: id, UserID: user.Id, RequestKey: fmt.Sprintf("user:%d:%s", user.Id, input.Key), ChannelID: channel.Id, GroupName: input.Group, Kind: input.Kind, Status: "pending", CreatedAt: time.Now().Unix(), Reserved: reserve, Price: string(encoded)}
if err = model.CreateDetection(run); err != nil {
common.ApiError(c, err)
return
}
c.JSON(200, gin.H{"success": true, "data": gin.H{"id": run.ID}})
}
func GetUserDetection(c *gin.Context) {
var run model.DetectionRun
query := model.DB.Where("user_id = ?", c.GetInt("id"))
if c.Param("id") == "latest" {
query = query.Where("group_name = ?", c.Query("group")).Order("created_at desc")
} else {
query = query.Where("id = ?", c.Param("id"))
}
if err := query.First(&run).Error; err != nil {
if c.Param("id") == "latest" {
c.JSON(200, gin.H{"success": true, "data": nil})
} else {
c.JSON(404, gin.H{"success": false, "message": "Detection not found"})
}
return
}
var calls []model.DetectionCall
if err := model.DB.Where("run_id = ?", run.ID).Find(&calls).Error; err != nil {
common.ApiError(c, err)
return
}
usage := []any{}
for _, call := range calls {
var parsed any
_ = common.UnmarshalJsonStr(call.Usage, &parsed)
if parsed != nil {
usage = append(usage, parsed)
}
if call.Error != "" && run.Error == "" {
run.Error = "usage_or_generation_incomplete"
}
}
c.JSON(200, gin.H{"success": true, "data": gin.H{"id": run.ID, "kind": run.Kind, "status": run.Status, "quality": run.Quality, "reserved": run.Reserved, "charged": run.Charged, "settled": run.Settled, "usage": usage, "error": run.Error}})
}
func GetDetectionArtwork(c *gin.Context) {
getDetectionArtwork(c, false)
}
func GetSiteDetectionArtwork(c *gin.Context) {
getDetectionArtwork(c, true)
}
func getDetectionArtwork(c *gin.Context, allGroups bool) {
var run model.DetectionRun
if err := model.DB.Select("html", "channel_id", "created_at").Where("id = ? AND created_at > ?", c.Param("id"), time.Now().Add(-6*time.Hour).Unix()).First(&run).Error; err != nil || run.HTML == "" {
c.Status(404)
return
}
settings, err := model.GetStatusPageGroups(c.Request.Context())
if err != nil {
c.Status(http.StatusServiceUnavailable)
return
}
channels, err := model.GetStatusChannels(c.Request.Context())
if err != nil {
c.Status(http.StatusServiceUnavailable)
return
}
usable := service.GetUserUsableGroups("")
visible := false
for _, channel := range channels {
if channel.Id != run.ChannelID {
continue
}
for group := range strings.SplitSeq(channel.Group, ",") {
group = strings.TrimSpace(group)
if _, ok := usable[group]; (allGroups || (ok && !settings[group].Hidden)) && run.CreatedAt > settings[group].ResetBefore {
visible = true
}
}
}
if !visible {
c.Status(http.StatusNotFound)
return
}
c.Header("Content-Security-Policy", "sandbox allow-scripts; default-src 'none'; script-src 'unsafe-inline'; style-src 'unsafe-inline'; img-src data:; form-action 'none'; base-uri 'none'; frame-ancestors 'self'")
c.Header("X-Content-Type-Options", "nosniff")
c.Header("Cache-Control", "no-store")
c.Data(200, "text/html; charset=utf-8", []byte(sanitizeDetectionHTML(run.HTML)))
}
// Decode only generation fields: external credentials cannot forward tools,
// previous responses, billing modifiers or arbitrary provider parameters.
func detectionProxyRequest(data []byte) (*dto.OpenAIResponsesRequest, error) {
var fields map[string]any
if len(data) > 64<<10 || common.Unmarshal(data, &fields) != nil {
return nil, errors.New("invalid request")
}
if name, ok := fields["model"].(string); !ok || name != fingerprint.Model {
return nil, errors.New("model not allowed")
}
for key := range fields {
switch key {
case "model", "input", "instructions", "stream", "max_output_tokens", "reasoning", "store", "tools", "tool_choice":
default:
return nil, errors.New("unsupported request field")
}
}
input, err := common.Marshal(fields["input"])
if err != nil || len(input) > 16384 {
return nil, errors.New("input too large")
}
var checkText func(any) bool
checkText = func(v any) bool {
switch value := v.(type) {
case string:
return true
case []any:
for _, item := range value {
if !checkText(item) {
return false
}
}
return true
case map[string]any:
for k, item := range value {
switch k {
case "role":
role, ok := item.(string)
if !ok || (role != "user" && role != "system" && role != "developer") {
return false
}
case "type":
kind, ok := item.(string)
if !ok || (kind != "input_text" && kind != "message") {
return false
}
case "content", "text":
if !checkText(item) {
return false
}
default:
return false
}
}
return true
}
return false
}
if !checkText(fields["input"]) {
return nil, errors.New("text input required")
}
instructions, _ := fields["instructions"].(string)
if len(input)+len(instructions) > 16384 {
return nil, errors.New("input too large")
}
// Service-selected bounds cannot be overridden by the bearer.
fields = map[string]any{"model": fingerprint.Model, "input": fields["input"], "instructions": instructions, "stream": true, "max_output_tokens": 16384, "reasoning": map[string]string{"effort": "high"}, "store": false, "tools": []any{}}
encoded, _ := common.Marshal(fields)
var request dto.OpenAIResponsesRequest
err = common.Unmarshal(encoded, &request)
return &request, err
}
func DetectionProxy(c *gin.Context) {
secret := strings.TrimPrefix(c.GetHeader("Authorization"), "Bearer ")
if !strings.HasPrefix(secret, "det_") {
c.AbortWithStatus(401)
return
}
data, err := io.ReadAll(io.LimitReader(c.Request.Body, (64<<10)+1))
if err != nil {
c.AbortWithStatus(400)
return
}
request, err := detectionProxyRequest(data)
if err != nil {
c.JSON(400, gin.H{"error": gin.H{"message": err.Error()}})
return
}
call, err := model.ClaimDetectionCall(model.DetectionSecretHash(secret))
if err != nil {
c.AbortWithStatus(401)
return
}
var run model.DetectionRun
if err = model.DB.First(&run, "id = ?", call.RunID).Error; err != nil {
c.AbortWithStatus(503)
return
}
ctx, cancel := context.WithDeadline(c.Request.Context(), time.Unix(run.Deadline-10, 0))
detectionCalls.Store(call.ID, cancel)
defer func() { cancel(); detectionCalls.Delete(call.ID) }()
channel, err := model.GetChannelById(run.ChannelID, true)
var raw []byte
var contentType string
// Flush SSE comments while generation is buffered for accounting. Without
// these bytes, the public edge times out a valid long reasoning request.
c.Header("Content-Type", "text/event-stream")
c.Header("Cache-Control", "no-cache")
c.Header("X-Accel-Buffering", "no")
c.Writer.WriteHeader(http.StatusOK)
_, _ = c.Writer.WriteString(": detection started\n\n")
c.Writer.Flush()
if err == nil && channel.Status == common.ChannelStatusEnabled {
raw, contentType, err = waitDetectionResponse(c, ctx, channel, request)
} else {
err = errors.New("channel unavailable")
}
success := err == nil
call.Available = &success
if err != nil {
call.Error = "channel_call_failed"
if ctx.Err() != nil {
call.Available = nil
call.Error = "interrupted_usage_unknown"
}
} else {
usage, complete := detectionResponseUsage(raw, contentType)
if !complete {
call.Error = "incomplete_response"
call.Available = nil
} else {
call.Usage = usage.Raw
quota, quotaErr := service.DetectionUsageQuota(run.Price, usage.Get("input_tokens").Int(), usage.Get("output_tokens").Int(), usage.Get("input_tokens_details.cached_tokens").Int(), max(usage.Get("input_tokens_details.cached_creation_tokens").Int(), usage.Get("input_tokens_details.cache_creation_tokens").Int()))
if quotaErr != nil || quota > run.Reserved {
call.Error = "settlement_unavailable"
} else if run.UserID > 0 {
call.Quota = quota
}
}
}
if success {
if html := detectionHTML(raw, contentType); html != "" {
if saveErr := model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Update("html", html).Error; saveErr != nil {
call.Error = "artwork_save_failed"
success = false
}
}
}
if err = model.FinishDetectionCall(call); err != nil {
success = false
}
if !success {
_, _ = c.Writer.WriteString("event: error\ndata: {\"type\":\"error\",\"error\":{\"message\":\"Detection generation failed\"}}\n\n")
c.Writer.Flush()
return
}
if strings.Contains(contentType, "text/event-stream") {
_, _ = c.Writer.Write(raw)
} else {
event, _ := common.Marshal(map[string]any{"type": "response.completed", "response": gjson.ParseBytes(raw).Value()})
_, _ = fmt.Fprintf(c.Writer, "event: response.completed\ndata: %s\n\n", event)
}
c.Writer.Flush()
}
func waitDetectionResponse(c *gin.Context, ctx context.Context, channel *model.Channel, request *dto.OpenAIResponsesRequest) ([]byte, string, error) {
type response struct {
raw []byte
contentType string
err error
}
done := make(chan response, 1)
go func() {
raw, contentType, err := runDetectionResponse(ctx, channel, request)
done <- response{raw, contentType, err}
}()
ticker := time.NewTicker(15 * time.Second)
defer ticker.Stop()
for {
select {
case result := <-done:
return result.raw, result.contentType, result.err
case <-ticker.C:
_, _ = c.Writer.WriteString(": waiting for generation\n\n")
c.Writer.Flush()
}
}
}
func runDetectionResponse(ctx context.Context, channel *model.Channel, request *dto.OpenAIResponsesRequest) ([]byte, string, error) {
encoded, err := common.Marshal(request)
if err != nil {
return nil, "", err
}
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequestWithContext(ctx, http.MethodPost, "/v1/responses", bytes.NewReader(encoded))
c.Request.Header.Set("Content-Type", "application/json")
if middleware.SetupContextForSelectedChannel(c, channel, fingerprint.Model) != nil {
return nil, "", errors.New("channel setup failed")
}
info, err := relaycommon.GenRelayInfo(c, types.RelayFormatOpenAIResponses, request, nil)
if err != nil {
return nil, "", err
}
info.IsChannelTest = true
info.DisablePing = true
info.AstraDetectionDisabled = true
info.AstraAutoSwitch = false
adaptor, body, closer, apiErr := relay.PrepareResponsesRequest(c, info, request)
if closer != nil {
defer closer.Close()
}
if apiErr != nil {
return nil, "", errors.New("preparation failed")
}
body, boundedCloser, err := boundDetectionBody(body, request)
if err != nil {
return nil, "", err
}
defer boundedCloser.Close()
response, err := adaptor.DoRequest(c, info, body)
if err != nil {
return nil, "", err
}
resp, ok := response.(*http.Response)
if !ok || resp == nil {
return nil, "", errors.New("unsupported transport")
}
defer resp.Body.Close()
if resp.StatusCode != 200 {
return nil, "", errors.New("upstream request failed")
}
raw, err := io.ReadAll(io.LimitReader(resp.Body, (8<<20)+1))
if len(raw) > 8<<20 {
return nil, "", errors.New("response too large")
}
return raw, resp.Header.Get("Content-Type"), err
}
func detectionResponseUsage(raw []byte, contentType string) (gjson.Result, bool) {
result := gjson.ParseBytes(raw)
if strings.Contains(contentType, "text/event-stream") {
result = gjson.Result{}
for line := range strings.SplitSeq(string(raw), "\n") {
if strings.HasPrefix(line, "data:") {
event := gjson.Parse(strings.TrimSpace(strings.TrimPrefix(line, "data:")))
switch event.Get("type").String() {
case "response.completed", "response.incomplete", "response.failed":
result = event.Get("response")
}
}
}
}
usage := result.Get("usage")
status := result.Get("status").String()
return usage, (status == "completed" || status == "incomplete" || status == "failed") && usage.Get("input_tokens").Type == gjson.Number && usage.Get("output_tokens").Type == gjson.Number
}
// Apply billing and safety limits AFTER channel parameter overrides. Never let
// an administrative override turn a bounded paid test into unbounded tool work.
func boundDetectionBody(body common.ReplayableBody, request *dto.OpenAIResponsesRequest) (common.ReplayableBody, io.Closer, error) {
raw, err := io.ReadAll(io.LimitReader(body, 128<<10))
if err != nil {
return nil, nil, err
}
var prepared map[string]any
if common.Unmarshal(raw, &prepared) != nil {
return nil, nil, errors.New("invalid prepared detection request")
}
var input any
if common.Unmarshal(request.Input, &input) != nil {
return nil, nil, errors.New("invalid detection input")
}
prepared["input"] = input
var instructions any
_ = common.Unmarshal(request.Instructions, &instructions)
prepared["instructions"] = instructions
prepared["max_output_tokens"] = *request.MaxOutputTokens
if *request.MaxOutputTokens == 16384 {
prepared["reasoning"] = map[string]string{"effort": "high"}
}
prepared["tools"] = []any{}
prepared["tool_choice"] = "none"
prepared["store"] = false
prepared["stream"] = true
for _, key := range []string{"previous_response_id", "conversation", "background", "max_tokens", "max_completion_tokens"} {
delete(prepared, key)
}
encoded, err := common.Marshal(prepared)
if err != nil {
return nil, nil, err
}
return relaycommon.NewOutboundJSONBody(encoded)
}
+63
View File
@@ -0,0 +1,63 @@
package controller
import (
"bytes"
"strings"
"golang.org/x/net/html"
)
// Artwork scripts run only inside an opaque-origin sandbox. CSP blocks network
// access; navigation, forms, and embedded documents are removed here.
func sanitizeDetectionHTML(source string) string {
root, err := html.Parse(strings.NewReader(source))
if err != nil {
return ""
}
var walk func(*html.Node)
walk = func(n *html.Node) {
for child := n.FirstChild; child != nil; {
next := child.NextSibling
switch strings.ToLower(child.Data) {
case "iframe", "frame", "frameset", "object", "embed", "base", "meta", "link", "form", "input", "textarea", "select", "audio", "video", "source", "foreignobject":
if child.Type == html.ElementNode {
n.RemoveChild(child)
child = next
continue
}
}
if child.Type == html.ElementNode {
attrs := child.Attr[:0]
for _, attr := range child.Attr {
name := strings.ToLower(attr.Key)
value := strings.TrimSpace(attr.Val)
if strings.HasPrefix(name, "on") {
continue
}
switch name {
case "href":
if !strings.HasPrefix(value, "#") {
continue
}
case "src":
if child.Data != "img" || !(strings.HasPrefix(value, "data:image/png;") || strings.HasPrefix(value, "data:image/jpeg;") || strings.HasPrefix(value, "data:image/webp;")) {
continue
}
case "srcset", "srcdoc", "action", "formaction", "ping", "target", "download", "http-equiv":
continue
}
attrs = append(attrs, attr)
}
child.Attr = attrs
}
walk(child)
child = next
}
}
walk(root)
var out bytes.Buffer
if html.Render(&out, root) != nil {
return ""
}
return out.String()
}
+171
View File
@@ -0,0 +1,171 @@
package controller
import (
"bufio"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/billingexpr"
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
)
func TestDetectionProxyFlushesBeforeGenerationAndConsumesCredentialOnce(t *testing.T) {
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
connection, err := db.DB()
require.NoError(t, err)
connection.SetMaxOpenConns(1)
oldDB := model.DB
model.DB = db
t.Cleanup(func() { model.DB = oldDB; _ = connection.Close() })
require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.DetectionRun{}, &model.DetectionCall{}))
release := make(chan struct{})
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
select {
case <-release:
case <-r.Context().Done():
return
}
w.Header().Set("Content-Type", "application/json")
_, _ = io.WriteString(w, `{"status":"completed","usage":{"input_tokens":10,"output_tokens":5},"output":[{"type":"message","content":[{"type":"output_text","text":"<svg></svg>"}]}]}`)
}))
defer upstream.Close()
url := upstream.URL
channel := model.Channel{Type: 1, Key: "test-only", BaseURL: &url, Status: common.ChannelStatusEnabled, Models: fingerprint.Model}
require.NoError(t, db.Create(&channel).Error)
price, _ := common.Marshal(service.DetectionPrice{Snapshot: billingexpr.BillingSnapshot{ExprString: "p+c", ExprHash: billingexpr.ExprHashString("p+c"), GroupRatio: 1, QuotaPerUnit: 500000, ExprVersion: 1}})
run := model.DetectionRun{ID: "stream-test", RequestKey: "stream-test", Kind: "pelican", ChannelID: channel.Id, Status: "running", Deadline: time.Now().Unix() + 60, Reserved: 100, Price: string(price)}
require.NoError(t, db.Create(&run).Error)
call := model.DetectionCall{ID: "stream-call", RunID: run.ID, SecretHash: model.DetectionSecretHash("det_test"), State: "ready", ExpiresAt: time.Now().Unix() + 60}
require.NoError(t, db.Create(&call).Error)
router := gin.New()
router.POST("/responses", DetectionProxy)
gateway := httptest.NewServer(router)
defer gateway.Close()
client := http.Client{Timeout: 5 * time.Second}
request, err := http.NewRequest(http.MethodPost, gateway.URL+"/responses", strings.NewReader(`{"model":"gpt-6-astra","input":"draw"}`))
require.NoError(t, err)
request.Header.Set("Authorization", "Bearer det_test")
response, err := client.Do(request)
close(release)
require.NoError(t, err)
defer response.Body.Close()
reader := bufio.NewReader(response.Body)
first, err := reader.ReadString('\n')
require.NoError(t, err)
assert.Equal(t, ": detection started\n", first)
raw, err := io.ReadAll(reader)
require.NoError(t, err)
assert.Contains(t, string(raw), "response.completed")
require.NoError(t, db.First(&call, "id = ?", call.ID).Error)
assert.Equal(t, "finished", call.State)
assert.Contains(t, call.Usage, "input_tokens")
replay := httptest.NewRecorder()
request = httptest.NewRequest(http.MethodPost, "/responses", strings.NewReader(`{"model":"gpt-6-astra","input":"draw"}`))
request.Header.Set("Authorization", "Bearer det_test")
router.ServeHTTP(replay, request)
assert.Equal(t, 401, replay.Code)
}
func TestDetectionProxyRejectsEscapesAndEnforcesBounds(t *testing.T) {
for _, body := range []string{
`{"model":[],"input":"hi"}`,
`{"model":"gpt-6-luna","input":"hi"}`,
`{"model":"gpt-6-astra","input":[{"role":[],"content":"hi"}]}`,
`{"model":"gpt-6-astra","input":"hi","previous_response_id":"x"}`,
`{"model":"gpt-6-astra","input":[{"type":"input_image","image_url":"file:///etc/passwd"}]}`,
} {
_, err := detectionProxyRequest([]byte(body))
require.Error(t, err)
}
request, err := detectionProxyRequest([]byte(`{"model":"gpt-6-astra","input":"draw a pelican","max_output_tokens":1000000,"tools":[{"type":"web_search"}]}`))
require.NoError(t, err)
require.NotNil(t, request.MaxOutputTokens)
assert.EqualValues(t, 16384, *request.MaxOutputTokens)
assert.Equal(t, "[]", string(request.Tools))
}
func TestDetectionOnlyBillsCompleteReportedUsage(t *testing.T) {
answer, parseErr := fingerprint.ReadResponse(strings.NewReader(`{"status":"incomplete","usage":{"input_tokens":100,"output_tokens":10,"input_tokens_details":{"cached_tokens":40,"cache_creation_tokens":5}}}`), "application/json")
require.ErrorIs(t, parseErr, fingerprint.ErrIncomplete)
assert.True(t, answer.HasUsage)
assert.EqualValues(t, 5, answer.CacheWriteTokens)
assert.Empty(t, answer.Text)
_, ok := detectionResponseUsage([]byte(`{"status":"completed"}`), "application/json")
assert.False(t, ok)
_, ok = detectionResponseUsage([]byte("data: {\"type\":\"response.output_text.delta\",\"delta\":\"hello\"}\n\n"), "text/event-stream")
assert.False(t, ok)
usage, ok := detectionResponseUsage([]byte("data: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"usage\":{\"input_tokens\":100,\"output_tokens\":10}}}\n\n"), "text/event-stream")
require.True(t, ok)
assert.EqualValues(t, 100, usage.Get("input_tokens").Int())
expression := `p * 2 + cr * 0.2 + c * 10`
price := service.DetectionPrice{Snapshot: billingexpr.BillingSnapshot{ExprString: expression, ExprHash: billingexpr.ExprHashString(expression), GroupRatio: 0.18, QuotaPerUnit: 500000, ExprVersion: 1}}
encoded, err := common.Marshal(price)
require.NoError(t, err)
quota, err := service.DetectionUsageQuota(string(encoded), 100, 10, 40)
require.NoError(t, err)
assert.Equal(t, 21, quota)
_, err = service.DetectionUsageQuota(string(encoded), 10, 20, 11)
assert.Error(t, err)
at := time.Date(2026, 9, 24, 9, 0, 0, 0, time.UTC)
price.Snapshot.ExprString = `hour("UTC") == 9 ? p * 2 : p * 100`
price.Snapshot.ExprHash = billingexpr.ExprHashString(price.Snapshot.ExprString)
price.Request.At = &at
encoded, err = common.Marshal(price)
require.NoError(t, err)
quota, err = service.DetectionUsageQuota(string(encoded), 100, 0, 0)
require.NoError(t, err)
assert.Equal(t, 18, quota)
}
func TestDetectionArtworkCannotNavigate(t *testing.T) {
cleaned := sanitizeDetectionHTML(`<meta http-equiv="refresh" content="0;url=https://evil.test"><script>alert(1)</script><a href="https://evil.test" onclick="alert(1)">x</a><iframe srcdoc="bad"></iframe><svg><circle cx="20" cy="20" r="10"/></svg><style>circle{fill:red}</style>`)
for _, bad := range []string{"evil.test", "onclick", "<iframe", "http-equiv"} {
assert.NotContains(t, cleaned, bad)
}
assert.Contains(t, cleaned, "<circle")
assert.Contains(t, cleaned, "<style>")
}
func TestDetectionPriceBoundIncludesIntermediateTiers(t *testing.T) {
bound, err := billingexpr.DetectionCostBound(`len < 100 ? p * 2 : len < 1000 ? p * 100 : p * 1`, 32768, 16384)
require.NoError(t, err)
assert.Equal(t, float64(3276800), bound)
_, err = billingexpr.DetectionCostBound(`p / (c + 1)`, 32768, 16384)
require.Error(t, err)
bound, err = billingexpr.DetectionCostBound(`len <= 272000 ? tier("standard", p * 10 + c * 50 + cr * 1 + cc * 12) : tier("long", p * 20 + c * 100 + cr * 2 + cc * 25)`, 32768, 16384)
require.NoError(t, err)
assert.Greater(t, bound, float64(16384*100))
}
func TestDetectionLatestPerChannelStatus(t *testing.T) {
now := time.Now().Unix()
for _, tc := range []struct {
name string
ids []int
evidence map[int]service.ChannelDetectionEvidence
want int
}{
{"successful fingerprint", []int{1}, map[int]service.ChannelDetectionEvidence{1: {At: now, Success: true}}, 1},
{"failed availability", []int{1}, map[int]service.ChannelDetectionEvidence{1: {At: now, Success: false}}, 0},
{"partial coverage", []int{1, 2}, map[int]service.ChannelDetectionEvidence{1: {At: now, Success: true}}, 2},
{"failure and unknown", []int{1, 2}, map[int]service.ChannelDetectionEvidence{1: {At: now, Success: false}}, -1},
{"expired success", []int{1}, map[int]service.ChannelDetectionEvidence{1: {At: now - 7201, Success: true}}, -1},
} {
t.Run(tc.name, func(t *testing.T) {
status, _ := service.DetectionFallbackStatus(tc.ids, tc.evidence, now)
assert.Equal(t, tc.want, status)
})
}
}
+410
View File
@@ -0,0 +1,410 @@
package controller
import (
"bytes"
"context"
"crypto/hmac"
"crypto/sha256"
"encoding/hex"
"errors"
"fmt"
"io"
"net/http"
"strconv"
"strings"
"sync"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
"github.com/QuantumNous/new-api/service"
"github.com/tidwall/gjson"
)
type detectionHandler struct{}
func pelicanCredential(id string) string {
mac := hmac.New(sha256.New, []byte(common.SessionSecret))
mac.Write([]byte("newapi-pelican-credential-v1:" + id))
return "det_" + hex.EncodeToString(mac.Sum(nil))
}
func (detectionHandler) Type() string { return "group_detection" }
func (detectionHandler) Enabled() bool {
return common.GetEnvOrDefaultBool("PELICAN_DETECTION_ENABLED", true)
}
func (detectionHandler) Interval() time.Duration { return 15 * time.Second }
func (detectionHandler) NewPayload() any { return nil }
func (detectionHandler) Run(ctx context.Context, task *model.SystemTask, runner string) {
_ = model.DeliverDetectionLogs()
// Keep accounting facts, but expire bulky generated artwork after the public window.
_ = model.DB.Model(&model.DetectionRun{}).Where("settled = ? AND created_at < ? AND html <> ?", true, time.Now().Add(-6*time.Hour).Unix(), "").Update("html", "").Error
scheduleErr := schedulePelicanDetections(ctx)
var err error
{
var runs []model.DetectionRun
err = model.DB.Where("settled = ?", false).Order("created_at asc").Find(&runs).Error
if err == nil {
jobs := make(chan model.DetectionRun)
var workers sync.WaitGroup
for range 2 {
workers.Go(func() {
for run := range jobs {
processDetection(ctx, run)
}
})
}
for _, run := range runs {
select {
case jobs <- run:
case <-ctx.Done():
}
}
close(jobs)
workers.Wait()
}
}
if err == nil {
err = scheduleErr
}
status := model.SystemTaskStatusSucceeded
if err != nil {
status = model.SystemTaskStatusFailed
}
finishSystemTaskHandler(task, runner, status, nil, err)
}
func schedulePelicanDetections(ctx context.Context) error {
channels, schedules, err := service.AutomaticDetectionChannelsDue(ctx, "pelican", time.Now())
if err != nil {
return err
}
now := time.Now().Unix()
for _, channel := range channels {
groups := channel.GetGroups()
if len(groups) == 0 {
continue
}
price, reserve, priceErr := service.DetectionPricing("", groups[0], "pelican")
if priceErr != nil {
return priceErr
}
created := time.Unix(now, 0)
price.Request.At = &created
encoded, _ := common.Marshal(price)
id, err := common.GenerateRandomCharsKey(32)
if err != nil {
return err
}
run := &model.DetectionRun{ID: id, RequestKey: fmt.Sprintf("adaptive:%d:%s", channel.Id, schedules[channel.Id].LastResultID), ChannelID: channel.Id, GroupName: groups[0], Kind: "pelican", Status: "pending", CreatedAt: now, Price: string(encoded), Reserved: reserve}
if err = model.CreateDetection(run); err != nil {
return err
}
}
return nil
}
func processDetection(ctx context.Context, run model.DetectionRun) {
if ctx.Err() != nil {
return
}
if run.Deadline == 0 {
run.Deadline = time.Now().Unix() + 600
if model.DB.Model(&model.DetectionRun{}).Where("id = ? AND deadline = 0", run.ID).Updates(map[string]any{"deadline": run.Deadline, "status": "running"}).Error != nil {
return
}
}
taskCtx, taskCancel := context.WithDeadline(ctx, time.Unix(run.Deadline, 0))
defer taskCancel()
if time.Now().Unix() >= run.Deadline {
var calls []model.DetectionCall
if model.DB.Where("run_id = ?", run.ID).Find(&calls).Error != nil {
return
}
for _, call := range calls {
if call.State != "running" {
continue
}
if cancel, ok := detectionCalls.Load(call.ID); ok {
cancel.(context.CancelFunc)()
return
}
// All proxy contexts expire before the task deadline. A vanished process
// cannot resume or recharge this consumed credential.
call.Error = "interrupted_usage_unknown"
call.Quota = 0
_ = model.FinishDetectionCall(&call)
}
if run.RemoteID != "" {
cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
_, _ = manxueRunRequest(cancelCtx, &run, http.MethodDelete, "/api/v1/tests/"+run.RemoteID, nil, "")
cancel()
}
finishDetection(run, "failed", "detection_timeout")
return
}
if run.Status == "succeeded" || run.Status == "failed" {
settleDetectionAndSchedule(run)
return
}
if run.Kind == "fingerprint" {
processPaidFingerprint(ctx, run)
return
}
if run.RemoteID == localPelicanID {
processLocalPelican(taskCtx, run)
return
}
if run.RemoteID == "" {
var existing int64
if model.DB.Model(&model.DetectionCall{}).Where("run_id = ?", run.ID).Count(&existing).Error != nil {
return
}
// A domain-separated PRF allows safe retry after a submit/commit crash;
// only its digest is stored, and the random run ID prevents enumeration.
secret := pelicanCredential(run.ID)
call := model.DetectionCall{ID: run.ID, RunID: run.ID, SecretHash: model.DetectionSecretHash(secret), ExpiresAt: run.Deadline + 300, State: "ready"}
if existing == 0 && model.DB.Create(&call).Error != nil {
return
}
payload := map[string]any{"benchmark": "pelican", "base_url": "https://ai.shrinken.pw/api/detection-proxy/v1", "api_key": secret, "model": fingerprint.Model, "protocol": "responses", "reasoning_effort": "high"}
var result gjson.Result
var err error
for taskCtx.Err() == nil {
result, err = manxueRunRequest(taskCtx, &run, http.MethodPost, "/api/v1/tests", payload, run.ID)
if errors.Is(err, errDetectionServiceUnavailable) && taskCtx.Err() == nil {
processLocalPelican(taskCtx, run)
return
}
if err == nil {
break
}
select {
case <-taskCtx.Done():
case <-time.After(5 * time.Second):
}
}
if err != nil || taskCtx.Err() != nil {
if ctx.Err() == nil {
processDetection(ctx, run)
}
return
}
remoteID := result.Get("id").String()
if remoteID == "" || len(remoteID) > 128 || strings.ContainsAny(remoteID, "/?#\\") {
finishDetection(run, "failed", "invalid_detection_response")
return
}
if model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Updates(map[string]any{"remote_id": remoteID, "status": "running"}).Error != nil {
return
}
run.RemoteID = remoteID
}
// Each worker retains its slot until completion, limiting remote generation
// concurrency as well as local submission concurrency.
for time.Now().Unix() < run.Deadline {
result, err := manxueRunRequest(taskCtx, &run, http.MethodGet, "/api/v1/tests/"+run.RemoteID, nil, "")
if errors.Is(err, errDetectionServiceUnavailable) && taskCtx.Err() == nil {
processLocalPelican(taskCtx, run)
return
}
if err != nil {
select {
case <-ctx.Done():
return
case <-time.After(5 * time.Second):
}
continue
}
switch result.Get("status").String() {
case "succeeded":
quality := result.Get("assessment.quality").String()
if quality != "normal" && quality != "degraded" {
quality = "unknown"
}
html := result.Get("result.html").String()
if len(html) > 2<<20 {
html = ""
}
metadata, _ := common.Marshal(map[string]string{"source_url": manxueArtworkURL(result)})
if model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Updates(map[string]any{"quality": quality, "html": html, "result": string(metadata)}).Error != nil {
return
}
finishDetection(run, "succeeded", "")
return
case "failed", "cancelled":
finishDetection(run, "failed", "detection_service_failed")
return
}
select {
case <-ctx.Done():
return
case <-time.After(5 * time.Second):
}
}
processDetection(ctx, run)
}
func finishDetection(run model.DetectionRun, status, reason string) {
updates := map[string]any{"status": status, "error": reason}
if status == "failed" && run.Kind == "pelican" {
updates["result"] = ""
}
if model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Updates(updates).Error != nil {
return
}
settleDetectionAndSchedule(run)
}
func settleDetectionAndSchedule(run model.DetectionRun) {
if model.SettleDetection(run.ID) != nil || run.UserID != 0 || run.Kind != "pelican" {
return
}
// Reconciliation also repairs the clock if a process exits after settlement.
if err := model.DB.Select("id", "channel_id", "finished_at", "status", "quality").First(&run, "id = ?", run.ID).Error; err != nil {
return
}
outcome := "unknown"
if run.Status == "succeeded" {
outcome = run.Quality
}
_ = model.ApplyDetectionScheduleResult(context.Background(), run.ChannelID, "pelican", run.ID, run.FinishedAt, outcome)
}
func manxueRequest(ctx context.Context, method, path string, payload any, key string) (gjson.Result, error) {
client := &http.Client{Timeout: 35 * time.Second, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
return manxueRequestWithClient(ctx, client, method, path, payload, key)
}
func manxueRequestWithClient(ctx context.Context, client *http.Client, method, path string, payload any, key string) (gjson.Result, error) {
var encoded []byte
if payload != nil {
var err error
encoded, err = common.Marshal(payload)
if err != nil {
return gjson.Result{}, err
}
}
for attempt := 0; attempt < 3; attempt++ {
req, err := http.NewRequestWithContext(ctx, method, "https://manxue.ai"+path, bytes.NewReader(encoded))
if err != nil {
return gjson.Result{}, err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", "Mozilla/5.0")
if strings.HasPrefix(path, "/api/visitor/") {
req.Header.Set("Origin", "https://manxue.ai")
}
if key != "" {
req.Header.Set("Idempotency-Key", key)
}
resp, err := client.Do(req)
if err != nil {
return gjson.Result{}, errDetectionServiceUnavailable
}
body, readErr := io.ReadAll(io.LimitReader(resp.Body, (3<<20)+1))
resp.Body.Close()
if readErr != nil || resp.StatusCode >= 500 {
return gjson.Result{}, errDetectionServiceUnavailable
}
if resp.StatusCode == 429 || resp.StatusCode == 503 {
seconds, _ := strconv.Atoi(resp.Header.Get("Retry-After"))
if at, err := http.ParseTime(resp.Header.Get("Retry-After")); err == nil {
seconds = int(time.Until(at).Seconds()) + 1
}
seconds = max(seconds, 2<<attempt)
select {
case <-ctx.Done():
return gjson.Result{}, ctx.Err()
case <-time.After(time.Duration(seconds) * time.Second):
}
continue
}
if readErr != nil || len(body) > 3<<20 || resp.StatusCode < 200 || resp.StatusCode >= 300 || !gjson.ValidBytes(body) {
return gjson.Result{}, errors.New("invalid detection service response")
}
return gjson.ParseBytes(body), nil
}
return gjson.Result{}, errDetectionServiceUnavailable
}
func processPaidFingerprint(ctx context.Context, run model.DetectionRun) {
channel, err := model.GetChannelById(run.ChannelID, true)
if err != nil || channel.Status != common.ChannelStatusEnabled {
finishDetection(run, "failed", "channel_unavailable")
return
}
samples := []fingerprint.Sample{}
for index := 0; index < 3; index++ {
id := fmt.Sprintf("%s_%d", run.ID, index)
var previous []model.DetectionCall
if model.DB.Where("id = ?", id).Find(&previous).Error != nil {
return
}
if len(previous) > 0 {
if previous[0].State == "running" {
return
}
var sample fingerprint.Sample
if previous[0].Sample != "" && common.UnmarshalJsonStr(previous[0].Sample, &sample) == nil {
samples = append(samples, sample)
}
continue
}
call := model.DetectionCall{ID: id, RunID: run.ID, SecretHash: model.DetectionSecretHash(id), State: "running", StartedAt: time.Now().Unix()}
if model.DB.Create(&call).Error != nil {
return
}
probeCtx, cancel := context.WithDeadline(ctx, time.Unix(run.Deadline-10, 0))
detectionCalls.Store(call.ID, cancel)
answer, count, probeErr := runAstraFingerprintProbe(probeCtx, channel, true)
interrupted := probeCtx.Err() != nil
cancel()
available := probeErr == nil || answer.HasUsage
call.Available = &available
if !answer.HasUsage && (errors.Is(probeErr, fingerprint.ErrInvalidSample) || errors.Is(probeErr, fingerprint.ErrIncomplete) || interrupted) {
call.Available = nil
}
if probeErr != nil {
call.Error = "channel_call_failed"
}
if answer.HasUsage {
usage, _ := common.Marshal(answer)
call.Usage = string(usage)
quota, quotaErr := service.DetectionUsageQuota(run.Price, answer.InputTokens, answer.OutputTokens, answer.CachedTokens, answer.CacheWriteTokens)
if quotaErr != nil || quota > run.Reserved/3 {
call.Error = "settlement_unavailable"
} else {
call.Quota = quota
}
} else {
call.Error = "usage_unavailable"
}
if probeErr == nil && fingerprint.ValidArray(answer.Text, count) {
sample := fingerprint.Sample{Text: answer.Text, Count: count}
samples = append(samples, sample)
encoded, _ := common.Marshal(sample)
call.Sample = string(encoded)
}
if model.FinishDetectionCall(&call) != nil {
detectionCalls.Delete(call.ID)
return
}
detectionCalls.Delete(call.ID)
if ctx.Err() != nil {
return
}
}
result, err := fingerprint.Score(samples)
if err != nil {
finishDetection(run, "failed", "scoring_failed")
return
}
encoded, _ := common.Marshal(result)
if model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Updates(map[string]any{"quality": result.Status, "result": string(encoded)}).Error != nil {
return
}
finishDetection(run, "succeeded", "")
}
+28
View File
@@ -0,0 +1,28 @@
package controller
import (
"github.com/stretchr/testify/assert"
"testing"
"time"
)
func TestSuccessfulUserRequestRestoresStatus(t *testing.T) {
now := time.Now().Unix()
for _, tc := range []struct {
name string
success, probe int64
disabled bool
want int
}{
{"fresh user overrides failed probe", now, now - 10, false, 1},
{"newer failed probe wins", now - 10, now, false, 0},
{"expired success", now - 901, now - 1000, false, 0},
{"disabled remains disabled", now, now - 10, true, 0},
} {
t.Run(tc.name, func(t *testing.T) {
m := Monitor{Status: 0, ChannelDisabled: tc.disabled}
applyUserSuccess(&m, tc.success, tc.probe, now)
assert.Equal(t, tc.want, m.Status)
})
}
}
+4 -4
View File
@@ -22,7 +22,7 @@ func GetAllLogs(c *gin.Context) {
group := c.Query("group")
requestId := c.Query("request_id")
upstreamRequestId := c.Query("upstream_request_id")
logs, total, err := model.GetAllLogs(logType, startTimestamp, endTimestamp, modelName, username, tokenName, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), channel, group, requestId, upstreamRequestId)
logs, total, err := model.GetAllLogs(logType, startTimestamp, endTimestamp, modelName, username, tokenName, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), channel, group, requestId, upstreamRequestId, c.Query("correlation_key"), c.Query("attempt_status"))
if err != nil {
common.ApiError(c, err)
return
@@ -49,7 +49,7 @@ func GetUserLogs(c *gin.Context) {
group := c.Query("group")
requestId := c.Query("request_id")
upstreamRequestId := c.Query("upstream_request_id")
logs, total, err := model.GetUserLogs(userId, logType, startTimestamp, endTimestamp, modelName, tokenName, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), group, requestId, upstreamRequestId)
logs, total, err := model.GetUserLogs(userId, logType, startTimestamp, endTimestamp, modelName, tokenName, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), group, requestId, upstreamRequestId,c.Query("attempt_status"))
if err != nil {
common.ApiError(c, err)
return
@@ -109,7 +109,7 @@ func GetLogsStat(c *gin.Context) {
modelName := c.Query("model_name")
channel, _ := strconv.Atoi(c.Query("channel"))
group := c.Query("group")
stat, err := model.SumUsedQuota(logType, startTimestamp, endTimestamp, modelName, username, tokenName, channel, group)
stat, err := model.SumUsedQuota(logType, startTimestamp, endTimestamp, modelName, username, tokenName, channel, group, c.Query("correlation_key"), c.Query("attempt_status"))
if err != nil {
common.ApiError(c, err)
return
@@ -136,7 +136,7 @@ func GetLogsSelfStat(c *gin.Context) {
modelName := c.Query("model_name")
channel, _ := strconv.Atoi(c.Query("channel"))
group := c.Query("group")
quotaNum, err := model.SumUsedQuota(logType, startTimestamp, endTimestamp, modelName, username, tokenName, channel, group)
quotaNum, err := model.SumUsedQuota(logType, startTimestamp, endTimestamp, modelName, username, tokenName, channel, group,"",c.Query("attempt_status"))
if err != nil {
common.ApiError(c, err)
return
+27 -9
View File
@@ -30,17 +30,18 @@ func VerifyLogin(c *gin.Context) {
writeSecurityOperationError(c, service.ErrProofMethod)
return
}
bundle, err := service.VerifyLoginCode(request.FlowToken, request.Code, c.ClientIP(), c.Request.UserAgent())
bundle, migration, err := service.VerifyLoginCode(request.FlowToken, request.Code, c.ClientIP(), c.Request.UserAgent())
if err != nil {
writeSecurityOperationError(c, err)
return
}
completeVerifiedLoginResponse(c, bundle, service.VerificationMethodTwoFA)
completeVerifiedLoginResponse(c, bundle, service.VerificationMethodTwoFA, migration)
}
func LoginPasskeyBegin(c *gin.Context) {
var request struct {
FlowToken string `json:"flow_token"`
RPID string `json:"rp_id"`
}
if common.DecodeJson(c.Request.Body, &request) != nil || request.FlowToken == "" {
common.ApiErrorMsg(c, "参数错误")
@@ -56,7 +57,11 @@ func LoginPasskeyBegin(c *gin.Context) {
writeSecurityOperationError(c, err)
return
}
wa, err := passkeysvc.BuildWebAuthn(c.Request)
credentialRPID := ""
if credential.RPID != nil {
credentialRPID = *credential.RPID
}
wa, rpIDs, err := passkeysvc.BuildLoginWebAuthn(c.Request, request.RPID, credentialRPID)
if err != nil {
writeSecurityOperationError(c, err)
return
@@ -75,7 +80,7 @@ func LoginPasskeyBegin(c *gin.Context) {
writeSecurityOperationError(c, err)
return
}
common.ApiSuccess(c, gin.H{"flow_token": token, "expires_at": expiresAt, "options": options})
common.ApiSuccess(c, gin.H{"flow_token": token, "expires_at": expiresAt, "options": options, "rp_ids": rpIDs})
}
func LoginPasskeyFinish(c *gin.Context) {
@@ -104,6 +109,7 @@ func LoginPasskeyFinish(c *gin.Context) {
writeSecurityOperationError(c, err)
return
}
c.Set("passkey_rp_id", sessionData.RelyingPartyID)
if security.LoginFlowID != verification.Flow.Id || sessionData.UserVerification != protocol.VerificationRequired {
writeSecurityOperationError(c, model.ErrAuthFlowInvalid)
return
@@ -113,29 +119,33 @@ func LoginPasskeyFinish(c *gin.Context) {
writeSecurityOperationError(c, err)
return
}
wa, err := passkeysvc.BuildWebAuthn(c.Request)
wa, err := passkeysvc.BuildWebAuthnForRPID(c.Request, sessionData.RelyingPartyID)
if err != nil {
writeSecurityOperationError(c, err)
return
}
if credential.RPID != nil && *credential.RPID != "" && *credential.RPID != sessionData.RelyingPartyID {
writeSecurityOperationError(c, service.ErrVerificationFailed)
return
}
validated, err := wa.ValidateLogin(passkeysvc.NewWebAuthnUser(&model.User{Id: identity.UserID}, credential), *sessionData, parsed)
if err != nil {
writeSecurityOperationError(c, err)
return
}
if err := model.UpdatePasskeyAssertionState(identity.UserID, validated, time.Now()); err != nil {
if err := model.UpdatePasskeyAssertionState(identity.UserID, validated, time.Now(), sessionData.RelyingPartyID); err != nil {
writeSecurityOperationError(c, err)
return
}
bundle, err := service.CompleteLoginVerification(request.FlowToken, verification, service.VerificationMethodPasskey, c.ClientIP(), c.Request.UserAgent())
bundle, migration, err := service.CompleteLoginVerification(request.FlowToken, verification, service.VerificationMethodPasskey, c.ClientIP(), c.Request.UserAgent())
if err != nil {
writeSecurityOperationError(c, err)
return
}
completeVerifiedLoginResponse(c, bundle, service.VerificationMethodPasskey)
completeVerifiedLoginResponse(c, bundle, service.VerificationMethodPasskey, migration)
}
func completeVerifiedLoginResponse(c *gin.Context, bundle *service.AuthBundle, method string) {
func completeVerifiedLoginResponse(c *gin.Context, bundle *service.AuthBundle, method string, migration *service.LegacyGitHubMigration) {
identity, err := service.ParseAccessToken(bundle.AccessToken)
if err != nil {
writeAuthSessionError(c, err)
@@ -147,5 +157,13 @@ func completeVerifiedLoginResponse(c *gin.Context, bundle *service.AuthBundle, m
return
}
c.Set("login_verification_method", method)
if migration != nil {
// The legacy GitHub binding was rewritten together with this session.
notificationFailed := service.NotifyAccountSecurityChange(user.Email, "Login account linked: GitHub") != nil
recordLegacyGitHubBindingAudit(c, user, true, map[string]any{
"legacy_id": migration.LegacyID, "provider_user_id": migration.GitHubID,
"verification_method": method, "notification_failed": notificationFailed,
})
}
writeLoginResponse(c, user, bundle)
}
+27
View File
@@ -0,0 +1,27 @@
package controller
import (
"github.com/tidwall/gjson"
"net/url"
)
func manxueArtworkURL(result gjson.Result) string {
for _, key := range []string{"share_url", "public_url", "result_url", "result.share_url", "result.public_url"} {
raw := result.Get(key).String()
u, err := url.Parse(raw)
if err == nil && u.Scheme == "https" && u.Host == "manxue.ai" && u.User == nil && u.Query().Get("result") != "" {
return u.String()
}
}
id := result.Get("public_id").String()
if id == "" {
id = result.Get("result.public_id").String()
}
if id == "" && result.Get("public").Bool() {
id = result.Get("result.id").String()
}
if id == "" {
return ""
}
return "https://manxue.ai/?" + url.Values{"result": {id}}.Encode()
}
+19
View File
@@ -0,0 +1,19 @@
package controller
import (
"github.com/stretchr/testify/assert"
"github.com/tidwall/gjson"
"testing"
)
func TestManxueArtworkURL(t *testing.T) {
assert.Empty(t, manxueArtworkURL(gjson.Parse(`{"public":false,"result":{"id":"private"}}`)))
assert.Equal(t, "https://manxue.ai/?result=community-one", manxueArtworkURL(gjson.Parse(`{"public":true,"result":{"id":"community-one"}}`)))
assert.Empty(t, manxueArtworkURL(gjson.Parse(`{"share_url":"https://evil.test/?result=one"}`)))
}
func TestArtworkPreservesScriptDrawnLegs(t *testing.T) {
result := sanitizeDetectionHTML(`<svg><path id="leg"></path></svg><button id="toggle">pause</button><script>document.getElementById('leg').setAttribute('d','M0 0L10 10')</script>`)
assert.Contains(t, result, "<script>")
assert.Contains(t, result, "setAttribute")
assert.Contains(t, result, "<button")
}
+64
View File
@@ -0,0 +1,64 @@
package controller
import (
"context"
"errors"
"net/http"
"net/http/cookiejar"
"net/url"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/tidwall/gjson"
)
// Visitor sessions belong to a single detection. Keep them in its private task
// metadata until completion so polling can resume after a process restart.
type manxueVisitorSession struct {
Cookies []*http.Cookie `json:"manxue_cookies,omitempty"`
}
func manxueRunRequest(ctx context.Context, run *model.DetectionRun, method, path string, payload any, key string) (gjson.Result, error) {
var session manxueVisitorSession
_ = common.UnmarshalJsonStr(run.Result, &session)
if len(session.Cookies) == 0 && !(method == http.MethodPost && run.RemoteID == "") {
return manxueRequest(ctx, method, path, payload, key)
}
jar, _ := cookiejar.New(nil)
origin, _ := url.Parse("https://manxue.ai/api/visitor/")
client := &http.Client{Timeout: 35 * time.Second, Jar: jar, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
if len(session.Cookies) == 0 {
if _, err := manxueRequestWithClient(ctx, client, http.MethodGet, "/api/visitor/config", nil, ""); err != nil {
return gjson.Result{}, errDetectionServiceUnavailable
}
session.Cookies = jar.Cookies(origin)
if len(session.Cookies) == 0 {
return gjson.Result{}, errDetectionServiceUnavailable
}
encoded, err := common.Marshal(session)
if err != nil {
return gjson.Result{}, err
}
if err = model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Update("result", string(encoded)).Error; err != nil {
return gjson.Result{}, err
}
run.Result = string(encoded)
} else {
jar.SetCookies(origin, session.Cookies)
}
if method == http.MethodPost {
fields, ok := payload.(map[string]any)
if !ok {
return gjson.Result{}, errors.New("invalid pelican payload")
}
public := make(map[string]any, len(fields)+1)
for k, v := range fields {
public[k] = v
}
public["public"] = true
payload = public
}
return manxueRequestWithClient(ctx, client, method, strings.Replace(path, "/api/v1/tests", "/api/visitor/tests", 1), payload, key)
}
+80
View File
@@ -0,0 +1,80 @@
package controller
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"github.com/QuantumNous/new-api/model"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
"gorm.io/gorm"
)
type manxueTestTransport func(*http.Request) (*http.Response, error)
func (f manxueTestTransport) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }
func TestManxuePublicDetectionResumesAndReturnsPublicURL(t *testing.T) {
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
conn, err := db.DB()
require.NoError(t, err)
conn.SetMaxOpenConns(1)
oldDB, oldTransport := model.DB, http.DefaultTransport
model.DB = db
t.Cleanup(func() { model.DB = oldDB; http.DefaultTransport = oldTransport; _ = conn.Close() })
require.NoError(t, db.AutoMigrate(&model.DetectionRun{}))
configCalls, createCalls := 0, 0
http.DefaultTransport = manxueTestTransport(func(r *http.Request) (*http.Response, error) {
w := httptest.NewRecorder()
w.Header().Set("Content-Type", "application/json")
switch r.URL.Path {
case "/api/visitor/config":
configCalls++
http.SetCookie(w, &http.Cookie{Name: "gallery_visitor", Value: "test-session", Path: "/api/visitor/", Secure: true, HttpOnly: true})
_, _ = w.WriteString(`{"prompt":"pelican"}`)
case "/api/visitor/tests":
createCalls++
cookie, err := r.Cookie("gallery_visitor")
require.NoError(t, err)
require.Equal(t, "test-session", cookie.Value)
require.Equal(t, "stable-key", r.Header.Get("Idempotency-Key"))
var body map[string]any
require.NoError(t, json.NewDecoder(r.Body).Decode(&body))
require.Equal(t, true, body["public"])
_, _ = w.WriteString(`{"id":"remote-one","status":"running"}`)
case "/api/visitor/tests/remote-one":
_, err := r.Cookie("gallery_visitor")
require.NoError(t, err)
_, _ = w.WriteString(`{"status":"succeeded","public_id":"community-one","result":{"html":"<html>pelican</html>"}}`)
case "/api/v1/tests/legacy":
require.Empty(t, r.Header.Get("Cookie"))
_, _ = w.WriteString(`{"status":"succeeded"}`)
default:
t.Fatalf("unexpected path: %s", r.URL.Path)
}
return w.Result(), nil
})
run := model.DetectionRun{ID: "public-test", RequestKey: "public-test", Kind: "pelican"}
require.NoError(t, db.Create(&run).Error)
input := map[string]any{"benchmark": "pelican"}
_, err = manxueRunRequest(context.Background(), &run, http.MethodPost, "/api/v1/tests", input, "stable-key")
require.NoError(t, err)
require.NotContains(t, input, "public")
var resumed model.DetectionRun
require.NoError(t, db.First(&resumed, "id = ?", run.ID).Error)
resumed.RemoteID = "remote-one"
result, err := manxueRunRequest(context.Background(), &resumed, http.MethodGet, "/api/v1/tests/remote-one", nil, "")
require.NoError(t, err)
require.Equal(t, "https://manxue.ai/?result=community-one", manxueArtworkURL(result))
require.Equal(t, 1, configCalls)
require.Equal(t, 1, createCalls)
legacy := model.DetectionRun{RemoteID: "legacy"}
_, err = manxueRunRequest(context.Background(), &legacy, http.MethodGet, "/api/v1/tests/legacy", nil, "")
require.NoError(t, err)
require.Empty(t, manxueArtworkURL(gjson.Parse(`{"id":"private","public":false}`)))
}
+3 -3
View File
@@ -44,10 +44,10 @@ func TestStatus(c *gin.Context) {
func GetStatus(c *gin.Context) {
cs := console_setting.GetConsoleSetting()
passkeySetting := system_setting.PasskeySettingsSnapshot()
common.OptionMapRWMutex.RLock()
defer common.OptionMapRWMutex.RUnlock()
passkeySetting := system_setting.GetPasskeySettings()
legalSetting := system_setting.GetLegalSettings()
data := gin.H{
@@ -117,8 +117,8 @@ func GetStatus(c *gin.Context) {
"oidc_display_name": system_setting.GetOIDCSettings().GetEffectiveDisplayName(),
"passkey_login": passkeySetting.Enabled,
"passkey_display_name": passkeySetting.RPDisplayName,
"passkey_rp_id": passkeySetting.RPID,
"passkey_origins": passkeySetting.Origins,
"passkey_rp_id": passkeySetting.EffectiveRPID(),
"passkey_rp_ids": passkeySetting.RelyingPartyIDs(),
"passkey_allow_insecure": passkeySetting.AllowInsecureOrigin,
"passkey_user_verification": passkeySetting.UserVerification,
"passkey_attachment": passkeySetting.AttachmentPreference,
+12
View File
@@ -195,6 +195,18 @@ func getModelListGroups(c *gin.Context) (modelListGroups, error) {
return modelListGroups{}, err
}
}
if service.RequestRoutingPolicy(c) != nil {
entries := service.ResolveRequestPool(c, "")
ownerGroups := make([]string, 0, len(entries))
for _, entry := range entries {
ownerGroups = append(ownerGroups, entry.Group)
}
return modelListGroups{
userGroup: userGroup,
tokenGroup: tokenGroup,
ownerGroups: ownerGroups,
}, nil
}
if tokenGroup == "auto" {
return modelListGroups{
+39 -1
View File
@@ -493,6 +493,44 @@ func TestListModelsTokenLimitUsesResolvedCustomAutoGroups(t *testing.T) {
require.Empty(t, anthropicResponse.LastID)
}
func TestListModelsUsesTokenRoutingPolicyGroups(t *testing.T) {
withSelfUseModeEnabled(t)
originalUsableGroups := setting.UserUsableGroups2JSONString()
originalRatios := ratio_setting.GroupRatio2JSONString()
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"gpt":"GPT","domestic":"Domestic"}`))
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"gpt":1,"domestic":1}`))
t.Cleanup(func() {
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups))
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalRatios))
})
db := setupModelListControllerTestDB(t)
require.NoError(t, db.Create(&[]model.Ability{
{Group: "gpt", Model: "zz-gpt-model", ChannelId: 1, Enabled: true},
{Group: "domestic", Model: "zz-deepseek-model", ChannelId: 2, Enabled: true},
}).Error)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
ctx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
common.SetContextKey(ctx, constant.ContextKeyUserGroup, "gpt")
common.SetContextKey(ctx, constant.ContextKeyTokenGroup, "gpt")
ctx.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{
Version: 1,
Mode: "custom",
Groups: []model.TokenGroupEntry{
{Group: "gpt", Priority: 0},
{Group: "domestic", Priority: 1},
},
})
ListModels(ctx, constant.ChannelTypeOpenAI)
ids := decodeListModelsResponse(t, recorder)
require.Contains(t, ids, "zz-gpt-model")
require.Contains(t, ids, "zz-deepseek-model")
}
func TestSetupLoginDoesNotTouchPasswordWhenPasswordFieldOmitted(t *testing.T) {
db := setupModelListControllerTestDB(t)
require.NoError(t, db.AutoMigrate(&model.Log{}, &model.AuditLog{}, &model.UserSession{}, &model.TwoFA{}, &model.PasskeyCredential{}))
@@ -517,7 +555,7 @@ func TestSetupLoginDoesNotTouchPasswordWhenPasswordFieldOmitted(t *testing.T) {
Role: user.Role,
Status: user.Status,
Group: user.Group,
}, c)
}, nil, c)
})
recorder := httptest.NewRecorder()
+535 -2
View File
@@ -14,8 +14,13 @@ import (
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/billingexpr"
"github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/billing_setting"
"github.com/QuantumNous/new-api/setting/config"
"github.com/QuantumNous/new-api/setting/ratio_setting"
@@ -63,7 +68,7 @@ func modelManagementDB(t *testing.T, kind, dsn string) *gorm.DB {
for _, value := range restoreRatios {
require.NoError(t, value.restore("{}"))
}
config.UpdateConfigFromMap(config.GlobalConfig.Get("billing_setting"), map[string]string{"billing_mode": "{}", "billing_expr": "{}"})
config.UpdateConfigFromMap(config.GlobalConfig.Get("billing_setting"), map[string]string{"billing_mode": "{}", "billing_expr": "{}", "plugin_billing_expr": "{}"})
var version string
query := "SELECT version()"
if kind == "sqlite" {
@@ -75,7 +80,7 @@ func modelManagementDB(t *testing.T, kind, dsn string) *gorm.DB {
for _, value := range restoreRatios {
require.NoError(t, value.restore(value.value))
}
config.UpdateConfigFromMap(config.GlobalConfig.Get("billing_setting"), map[string]string{"billing_mode": previousConfig["billing_setting.billing_mode"], "billing_expr": previousConfig["billing_setting.billing_expr"]})
config.UpdateConfigFromMap(config.GlobalConfig.Get("billing_setting"), map[string]string{"billing_mode": previousConfig["billing_setting.billing_mode"], "billing_expr": previousConfig["billing_setting.billing_expr"], "plugin_billing_expr": previousConfig[billing_setting.PluginBillingExprOption]})
common.OptionMap = previousOptions
common.IsMasterNode, common.SQLitePath = previousMaster, previousSQLite
common.RedisEnabled, common.MemoryCacheEnabled = previousRedis, previousMemory
@@ -104,6 +109,369 @@ func modelManagementRequest(t *testing.T, handler gin.HandlerFunc, method, path
return recorder
}
func TestModelPricingConversionDatabaseMatrix(t *testing.T) {
previousQuota := common.QuotaPerUnit
common.QuotaPerUnit = 500000
t.Cleanup(func() { common.QuotaPerUnit = previousQuota })
for _, dialect := range []struct{ kind, env string }{{"sqlite", ""}, {"mysql", "TEST_MYSQL_DSN"}, {"postgres", "TEST_POSTGRES_DSN"}} {
t.Run(dialect.kind, func(t *testing.T) {
if dialect.env != "" && os.Getenv(dialect.env) == "" {
t.Skip("set " + dialect.env + " to run this database")
}
db := modelManagementDB(t, dialect.kind, os.Getenv(dialect.env))
contract, err := os.ReadFile("../pkg/billingexpr/testdata/frontend_simulation.json")
require.NoError(t, err)
var fixtures []struct {
Name, Expression string
Conversion *struct {
ModelName string `json:"model_name"`
Pricing model.PricingValues `json:"pricing"`
}
}
require.NoError(t, common.Unmarshal(contract, &fixtures))
for _, fixture := range fixtures {
if fixture.Conversion == nil {
continue
}
t.Run("frontend_contract/"+fixture.Name, func(t *testing.T) {
var response struct {
Success bool
Data model.ModelPricingConversion
}
modelManagementRequest(t, PreviewModelPricingConversion, http.MethodPost, "/api/option/model_pricing/convert", fixture.Conversion, &response)
require.True(t, response.Success)
assert.Equal(t, fixture.Expression, response.Data.Expression)
})
}
// The preview must resolve the draft, even when the running process
// still has a different saved output multiplier for that model.
require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"conversion-defaults":99}`))
for _, tc := range []struct {
name string
draft model.PricingValues
expression, reason string
}{
{"conversion-defaults", model.PricingValues{"ModelRatio": float64(2)}, `tier("base", p * 4 + c * 4)`, ""},
{"gpt-4-32k", model.PricingValues{"ModelRatio": 30.0, "CacheRatio": 1.0}, `tier("base", p * 60 + c * 120)`, ""},
{"conversion-free-cache", model.PricingValues{"ModelRatio": 0.0, "CacheRatio": 0.0}, `tier("base", p * 0 + c * 0 + cr * 0)`, ""},
{"claude-3-7-sonnet-20250219", model.PricingValues{"ModelRatio": float64(1.5), "CacheRatio": float64(0.1), "CreateCacheRatio": float64(1.25)}, `tier("base", p * 3 + c * 15 + cr * 0.3 + cc * 3.75 + cc1h * 6)`, ""},
{"conversion-image-default", model.PricingValues{"ModelRatio": float64(2), "ImageRatio": float64(1)}, `tier("base", p * 4 + c * 4)`, ""},
{"conversion-image-free", model.PricingValues{"ModelRatio": float64(2), "ImageRatio": float64(0)}, `tier("base", p * 4 + c * 4 + cr * 4 + img * 0)`, ""},
{"deepseek-chat", model.PricingValues{"ModelRatio": float64(0.135), "CacheRatio": float64(0.25)}, `tier("base", p * 0.27 + c * 0.27 + cr * 0.0675)`, ""},
{"gpt-4o-custom", model.PricingValues{"ModelRatio": float64(2)}, `tier("base", p * 4 + c * 16)`, ""},
{"gemini-2.5-pro-custom", model.PricingValues{"ModelRatio": float64(2)}, `tier("base", p * 4 + c * 32)`, ""},
{"vendor/claude-sonnet-4", model.PricingValues{"ModelRatio": float64(2), "CompletionRatio": float64(0)}, `tier("base", p * 4 + c * 0 + cr * 4 + cc * 5 + cc1h * 8)`, ""},
{"conversion-custom", model.PricingValues{"ModelRatio": float64(2), "CompletionRatio": float64(3), "CacheRatio": float64(0), "CreateCacheRatio": float64(1.5), "ImageRatio": float64(2)}, `tier("base", p * 4 + c * 12 + cr * 0 + cc * 6 + img * 8)`, ""},
{"deepseek-reasoner", model.PricingValues{"ModelRatio": 0.275, "CacheRatio": 0.25}, `tier("base", p * 0.55 + c * 0.55 + cr * 0.1375)`, ""},
{"gpt-5.6-sol", model.PricingValues{"ModelRatio": float64(2), "CompletionRatio": float64(2)}, `tier("base", p * 4 + c * 8)`, ""},
{"gpt-5.5", model.PricingValues{"ModelRatio": float64(2), "CompletionRatio": float64(2)}, `tier("base", p * 4 + c * 8)`, ""},
{"gpt-6-astra", model.PricingValues{"ModelRatio": float64(2), "CompletionRatio": float64(2)}, `tier("base", p * 4 + c * 8)`, ""},
{"conversion-free", model.PricingValues{"ModelPrice": float64(0)}, `tier("request", fixed(0))`, ""},
{"conversion-fixed", model.PricingValues{"ModelPrice": float64(0.25), "ModelRatio": float64(7)}, `tier("request", fixed(0.25))`, ""},
{"gpt-4o-2024-05-13", model.PricingValues{"ModelRatio": float64(2), "CompletionRatio": float64(99)}, `tier("base", p * 4 + c * 12)`, ""},
{"gpt-image-2", model.PricingValues{"ModelPrice": float64(1)}, `tier("image", fixed(1)) * image_count`, ""},
{"qwen-image-3.0-pro", model.PricingValues{"ModelPrice": float64(1)}, `tier("image", fixed(1)) * image_count`, ""},
{"wan2.7-image-pro", model.PricingValues{"ModelPrice": float64(1)}, "", "Task pricing must be converted manually using the task usage schema."},
{"video-summarizer", model.PricingValues{"ModelPrice": float64(1)}, `tier("request", fixed(1))`, ""},
{"sora-transcript-helper", model.PricingValues{"ModelPrice": float64(1)}, `tier("request", fixed(1))`, ""},
{"gpt-realtime", model.PricingValues{"ModelRatio": float64(1)}, "", "Realtime pricing must be converted manually."},
{"conversion-audio", model.PricingValues{"ModelRatio": float64(1), "AudioRatio": float64(2)}, `tier("base", p * 2 + c * 2 + ai * 4 + ao * 4)`, ""},
} {
t.Run(tc.name, func(t *testing.T) {
before, err := model.GetModelPricingSnapshot([]string{tc.name})
require.NoError(t, err)
var response struct {
Success bool
Data model.ModelPricingConversion
}
modelManagementRequest(t, PreviewModelPricingConversion, http.MethodPost, "/api/option/model_pricing/convert", map[string]any{"model_name": tc.name, "pricing": tc.draft}, &response)
require.True(t, response.Success)
assert.Equal(t, tc.expression, response.Data.Expression)
assert.Equal(t, tc.reason, response.Data.UnsupportedReason)
if tc.name == "claude-3-7-sonnet-20250219" {
// Both providers include images in input; Claude reports cache
// separately and does not report an image token breakdown.
for _, usage := range []dto.Usage{
{PromptTokens: 1000, CompletionTokens: 100, UsageSemantic: "openai", PromptTokensDetails: dto.InputTokenDetails{CachedTokens: 200, CachedCreationTokens: 50, ImageTokens: 300}},
{PromptTokens: 750, CompletionTokens: 100, UsageSemantic: "anthropic", ClaudeCacheCreation5mTokens: 50, PromptTokensDetails: dto.InputTokenDetails{CachedTokens: 200, CachedCreationTokens: 50}},
} {
tokens := service.BuildTieredTokenParams(&usage, usage.UsageSemantic == "anthropic", billingexpr.UsedVars(response.Data.Expression))
cost, _, err := billingexpr.RunExpr(response.Data.Expression, tokens)
require.NoError(t, err)
assert.Equal(t, float64(750), tokens.P)
assert.Equal(t, 3997.5, cost, "image input must remain billed exactly once")
}
usage := dto.Usage{PromptTokens: 750, CompletionTokens: 100, UsageSemantic: "anthropic", ClaudeCacheCreation5mTokens: 30, ClaudeCacheCreation1hTokens: 20, PromptTokensDetails: dto.InputTokenDetails{CachedTokens: 200, CachedCreationTokens: 50}}
tokens := service.BuildTieredTokenParams(&usage, true, billingexpr.UsedVars(response.Data.Expression))
cost, _, err := billingexpr.RunExpr(response.Data.Expression, tokens)
require.NoError(t, err)
assert.Equal(t, 4042.5, cost, "mixed TTLs preserve both legacy cache-write prices")
}
if tc.name == "deepseek-chat" {
for _, usage := range []dto.Usage{
{PromptTokens: 1000, CompletionTokens: 100, PromptTokensDetails: dto.InputTokenDetails{CachedTokens: 200}},
{PromptTokens: 800, CompletionTokens: 100, UsageSemantic: "anthropic", PromptTokensDetails: dto.InputTokenDetails{CachedTokens: 200}},
} {
tokens := service.BuildTieredTokenParams(&usage, usage.UsageSemantic == "anthropic", billingexpr.UsedVars(response.Data.Expression))
cost, _, err := billingexpr.RunExpr(response.Data.Expression, tokens)
require.NoError(t, err)
assert.Equal(t, 256.5, cost, "cache reads remain billed across both response formats")
}
}
after, err := model.GetModelPricingSnapshot([]string{tc.name})
require.NoError(t, err)
assert.Equal(t, before.Entries, after.Entries, "preview must not write pricing")
if tc.expression == "" {
return
}
pricing := tc.draft
pricing["billing_setting.billing_mode"] = "tiered_expr"
pricing["billing_setting.billing_expr"] = tc.expression
change := model.ModelPricingChange{ModelName: tc.name, ExpectedVersion: before.Entries[0].Version, Pricing: pricing}
require.NoError(t, model.UpdateModelPricing([]model.ModelPricingChange{change}))
after, err = model.GetModelPricingSnapshot([]string{tc.name})
require.NoError(t, err)
assert.Equal(t, pricing, after.Entries[0].Configured)
assert.ErrorIs(t, model.UpdateModelPricing([]model.ModelPricingChange{change}), model.ErrModelPricingConflict)
change.ExpectedVersion = after.Entries[0].Version
change.Pricing["billing_setting.billing_mode"] = "ratio"
require.NoError(t, model.UpdateModelPricing([]model.ModelPricingChange{change}))
after, err = model.GetModelPricingSnapshot([]string{tc.name})
require.NoError(t, err)
assert.Equal(t, "ratio", after.Entries[0].Effective["billing_setting.billing_mode"])
if _, fixed := pricing["ModelPrice"]; !fixed {
var preview struct {
Success bool
Data model.ModelPricingDescription
}
modelManagementRequest(t, PreviewModelPricing, http.MethodPost, "/api/option/model_pricing/preview", map[string]any{"model_name": tc.name, "pricing": pricing}, &preview)
require.True(t, preview.Success)
assert.Equal(t, after.Entries[0].Effective, preview.Data.Effective)
require.NotEmpty(t, response.Data.Effective)
for field, value := range response.Data.Effective {
assert.Equal(t, value, preview.Data.Effective[field], "conversion preview field %s", field)
}
cache, _ := ratio_setting.GetCacheRatio(tc.name)
createCache, _ := ratio_setting.GetCreateCacheRatio(tc.name)
image, _ := ratio_setting.GetImageRatio(tc.name)
assert.Equal(t, ratio_setting.GetCompletionRatio(tc.name), preview.Data.Effective["CompletionRatio"])
assert.Equal(t, cache, preview.Data.Effective["CacheRatio"])
assert.Equal(t, createCache, preview.Data.Effective["CreateCacheRatio"])
assert.Equal(t, image, preview.Data.Effective["ImageRatio"])
}
})
}
require.NoError(t, (&model.Channel{Name: "Mapped image", Models: "conversion-alias", Type: 1, ModelMapping: common.GetPointer(`{"conversion-alias":"conversion-hop","conversion-hop":"gpt-image-2"}`)}).Insert())
preview, err := model.PreviewModelPricingConversion("conversion-alias", model.PricingValues{"ModelPrice": float64(1)})
require.NoError(t, err)
assert.Equal(t, `tier("image", fixed(1)) * image_count`, preview.Expression)
assert.Empty(t, preview.UnsupportedReason)
t.Run("cache_writes_follow_configured_prices_and_claude_names", func(t *testing.T) {
for _, tc := range []struct {
name string
configured *float64
mode model.CacheWriteMode
}{
{"unrecognized-model", nil, model.CacheWriteNone},
{"gpt-5.6-terra", nil, model.CacheWriteNone},
{"gpt-5.6-terra", common.GetPointer(1.25), model.CacheWriteStandard},
{"vendor/claude-sonnet-4-6-high", nil, model.CacheWriteClaudeTTL},
{"custom-CLAUDE-alias", common.GetPointer(0.0), model.CacheWriteClaudeTTL},
{"deepseek-reasoner", common.GetPointer(0.0), model.CacheWriteStandard},
{"custom-cache", common.GetPointer(1.5), model.CacheWriteStandard},
} {
t.Run(tc.name+"/"+string(tc.mode), func(t *testing.T) {
draft := model.PricingValues{"ModelRatio": 2.0, "CompletionRatio": 2.0, "CacheRatio": 0.25, "billing_setting.billing_mode": "ratio"}
if tc.configured != nil {
draft["CreateCacheRatio"] = *tc.configured
}
var response struct {
Success bool
Data model.ModelPricingConversion
}
modelManagementRequest(t, PreviewModelPricingConversion, http.MethodPost, "/api/option/model_pricing/convert", map[string]any{"model_name": tc.name, "pricing": draft}, &response)
require.True(t, response.Success)
require.NotEmpty(t, response.Data.Expression)
assert.Equal(t, tc.mode, response.Data.CacheWriteMode)
var prices struct {
Success bool
Data model.ModelPricingDescription
}
modelManagementRequest(t, PreviewModelPricing, http.MethodPost, "/api/option/model_pricing/preview", map[string]any{"model_name": tc.name, "pricing": draft}, &prices)
require.True(t, prices.Success)
assert.Equal(t, response.Data.Effective, prices.Data.Effective)
assert.Equal(t, response.Data.CacheWriteMode, prices.Data.CacheWriteMode)
vars := billingexpr.UsedVars(response.Data.Expression)
assert.Equal(t, tc.mode != model.CacheWriteNone, vars["cc"])
assert.Equal(t, tc.mode == model.CacheWriteClaudeTTL, vars["cc1h"])
if tc.configured != nil && *tc.configured == 0 {
cost, _, err := billingexpr.RunExpr(response.Data.Expression, billingexpr.TokenParams{CC: 100, CC1h: 100})
require.NoError(t, err)
assert.Zero(t, cost, "explicit zero cache prices must stay free")
}
})
}
// The billing name remains authoritative for cache terms.
require.NoError(t, (&model.Channel{Name: "Cache alias", Type: 1, Models: "opaque-cache-alias", ModelMapping: common.GetPointer(`{"opaque-cache-alias":"claude-sonnet-4-6"}`)}).Insert())
alias, err := model.PreviewModelPricingConversion("opaque-cache-alias", model.PricingValues{"ModelRatio": 2.0})
require.NoError(t, err)
assert.Equal(t, `tier("base", p * 4 + c * 4)`, alias.Expression)
})
t.Run("same_price_cache_reads_preserve_overlapping_categories", func(t *testing.T) {
for _, tc := range []struct {
field, term string
ratio float64
details dto.InputTokenDetails
cost, unpricedCacheCost float64
}{
{"CreateCacheRatio", "cc * 75", 1.25, dto.InputTokenDetails{CachedTokens: 800, CacheWriteTokens: 500}, 97500, 79500},
{"ImageRatio", "img * 120", 2, dto.InputTokenDetails{CachedTokens: 600, ImageTokens: 500}, 108000, 102000},
} {
t.Run(tc.field, func(t *testing.T) {
converted, err := model.PreviewModelPricingConversion("overlapping-cache", model.PricingValues{"ModelRatio": 30.0, "CompletionRatio": 2.0, "CacheRatio": 1.0, tc.field: tc.ratio})
require.NoError(t, err)
assert.Equal(t, `tier("base", p * 60 + c * 120 + cr * 60 + `+tc.term+`)`, converted.Expression)
usage := dto.Usage{PromptTokens: 1000, CompletionTokens: 100, PromptTokensDetails: tc.details}
params := service.BuildTieredTokenParams(&usage, false, billingexpr.UsedVars(converted.Expression))
cost, _, err := billingexpr.RunExpr(converted.Expression, params)
require.NoError(t, err)
assert.Equal(t, tc.cost, cost, "preserve the existing overlapping-category charge")
unpriced := `tier("base", p * 60 + c * 120 + ` + tc.term + `)`
params = service.BuildTieredTokenParams(&usage, false, billingexpr.UsedVars(unpriced))
cost, _, err = billingexpr.RunExpr(unpriced, params)
require.NoError(t, err)
assert.Equal(t, tc.unpricedCacheCost, cost, "existing expressions without cr retain their OpenAI normalization")
})
}
})
t.Run("media_prices_and_request_adjustments", func(t *testing.T) {
for _, quantity := range []string{"129", "-1", "0.5", "18446744073709551615"} {
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", strings.NewReader(`{"model":"z-image","parameters":{"n":`+quantity+`}}`))
ctx.Request.Header.Set("Content-Type", "application/json")
Relay(ctx, types.RelayFormatOpenAIImage)
assert.Equal(t, http.StatusBadRequest, recorder.Code, recorder.Body.String())
}
conflicting, err := model.PreviewModelPricingConversion("gemini-2.5-flash-conflict", model.PricingValues{"ModelRatio": 0.15, "AudioRatio": 10.0})
require.NoError(t, err)
assert.Contains(t, conflicting.UnsupportedReason, "Gemini and OpenAI audio prices differ")
aliName := "z-image-media-conversion"
require.NoError(t, (&model.Channel{Name: aliName, Models: aliName, Type: 17, Status: common.ChannelStatusEnabled}).Insert())
aliConversion, err := model.PreviewModelPricingConversion(aliName, model.PricingValues{"ModelPrice": 0.04})
require.NoError(t, err)
require.Empty(t, aliConversion.UnsupportedReason)
quantity := 3
aliCost, _, err := billingexpr.RunExprWithRequest(aliConversion.Expression, billingexpr.TokenParams{}, billingexpr.RequestInput{ImageCount: &quantity, Body: []byte(`{"parameters":{"prompt_extend":true}}`)})
require.NoError(t, err)
assert.Equal(t, 240000.0, aliCost)
require.NoError(t, (&model.Channel{Name: "Different image rules", Models: aliName, Type: 1, Status: common.ChannelStatusEnabled}).Insert())
conflicting, err = model.PreviewModelPricingConversion(aliName, model.PricingValues{"ModelPrice": 0.04})
require.NoError(t, err)
assert.Contains(t, conflicting.UnsupportedReason, "different image request multipliers")
for _, tc := range []struct {
name string
pricing model.PricingValues
usage dto.Usage
want float64
}{
{"gemini-2.5-flash", model.PricingValues{"ModelRatio": 0.15, "CompletionRatio": 2.0, "CacheRatio": 0.1}, dto.Usage{PromptTokens: 1000, CompletionTokens: 100, PromptTokensDetails: dto.InputTokenDetails{CachedTokens: 100, AudioTokens: 200}}, 473},
{"gemini-2.5-flash-zero", model.PricingValues{"ModelRatio": 0.0}, dto.Usage{PromptTokens: 1000, PromptTokensDetails: dto.InputTokenDetails{AudioTokens: 200}}, 200},
{"media-audio", model.PricingValues{"ModelRatio": 1.0, "CompletionRatio": 3.0, "AudioRatio": 2.0, "AudioCompletionRatio": 4.0}, dto.Usage{PromptTokens: 1000, CompletionTokens: 100, PromptTokensDetails: dto.InputTokenDetails{TextTokens: 800, AudioTokens: 200}, CompletionTokenDetails: dto.OutputTokenDetails{TextTokens: 50, AudioTokens: 50}}, 3500},
{"media-free-audio", model.PricingValues{"ModelRatio": 1.0, "AudioRatio": 0.0}, dto.Usage{PromptTokens: 1000, CompletionTokens: 100, PromptTokensDetails: dto.InputTokenDetails{AudioTokens: 200}, CompletionTokenDetails: dto.OutputTokenDetails{AudioTokens: 100}}, 1600},
{"media-audio-cache", model.PricingValues{"ModelRatio": 1.0, "CompletionRatio": 3.0, "AudioRatio": 2.0, "CacheRatio": 0.1}, dto.Usage{PromptTokens: 1000, CompletionTokens: 100, PromptTokensDetails: dto.InputTokenDetails{CachedTokens: 400, AudioTokens: 200}, CompletionTokenDetails: dto.OutputTokenDetails{AudioTokens: 50}}, 2900},
{"media-text-cache", model.PricingValues{"ModelRatio": 1.0, "CompletionRatio": 3.0, "AudioRatio": 2.0, "CacheRatio": 0.1}, dto.Usage{PromptTokens: 1000, CompletionTokens: 100, PromptTokensDetails: dto.InputTokenDetails{CachedTokens: 400}}, 1880},
{"gemini-2.5-flash-overlap", model.PricingValues{"ModelRatio": 1.0}, dto.Usage{PromptTokens: 1000, PromptTokensDetails: dto.InputTokenDetails{CachedTokens: 800, AudioTokens: 500}}, 2100},
} {
t.Run(tc.name, func(t *testing.T) {
converted, err := model.PreviewModelPricingConversion(tc.name, tc.pricing)
require.NoError(t, err)
require.Empty(t, converted.UnsupportedReason)
require.NotNil(t, converted.BillingDetails.AudioInputPrice)
params := service.BuildTieredTokenParams(&tc.usage, false, billingexpr.UsedVars(converted.Expression))
cost, _, err := billingexpr.RunExpr(converted.Expression, params)
require.NoError(t, err)
assert.InDelta(t, tc.want, cost, 1e-9)
})
}
converted, err := model.PreviewModelPricingConversion("dall-e-3", model.PricingValues{"ModelPrice": 0.04})
require.NoError(t, err)
require.Empty(t, converted.UnsupportedReason)
assert.NotContains(t, converted.Expression, "256x256")
assert.NotContains(t, converted.Expression, "512x512")
smallImage, err := model.PreviewModelPricingConversion("dall-e-2", model.PricingValues{"ModelPrice": 0.02})
require.NoError(t, err)
require.Empty(t, smallImage.UnsupportedReason)
assert.NotContains(t, smallImage.Expression, "1792")
smallCost, _, err := billingexpr.RunExprWithRequest(smallImage.Expression, billingexpr.TokenParams{}, billingexpr.RequestInput{Body: []byte(`{"size":"256x256"}`), ImageCount: &quantity})
require.NoError(t, err)
assert.Equal(t, 24000.0, smallCost)
for _, tc := range []struct {
size, quality string
count int
cost float64
}{
{"1024x1024", "standard", 3, 120000},
{"1024x1024", "hd", 2, 160000},
{"1024x1792", "hd", 3, 360000},
} {
body, err := common.Marshal(map[string]any{"size": tc.size, "quality": tc.quality})
require.NoError(t, err)
cost, _, err := billingexpr.RunExprWithRequest(converted.Expression, billingexpr.TokenParams{}, billingexpr.RequestInput{Body: body, ImageCount: &tc.count})
require.NoError(t, err)
assert.InDelta(t, tc.cost, cost, 1e-9)
}
for _, channelType := range []int{20, 58} {
name := fmt.Sprintf("gemini-3.5-flash-media-%d", channelType)
require.NoError(t, (&model.Channel{Name: name, Models: name, Type: channelType}).Insert())
require.NoError(t, db.Create(&model.Model{ModelName: name, Endpoints: `{"custom-text":"/v1/custom"}`}).Error)
converted, err := model.PreviewModelPricingConversion(name, model.PricingValues{"ModelRatio": 0.75, "CompletionRatio": 6.0, "CacheRatio": 0.1})
require.NoError(t, err)
assert.Empty(t, converted.UnsupportedReason)
assert.Equal(t, `tier("base", p * 1.5 + c * 9 + cr * 0.15)`, converted.Expression)
assert.Nil(t, converted.BillingDetails.AudioInputPrice)
}
})
require.NoError(t, db.Create(&model.Model{ModelName: "conversion-metadata", Endpoints: `{"image-generation":"/v1/images/generations"}`}).Error)
preview, err = model.PreviewModelPricingConversion("conversion-metadata", model.PricingValues{"ModelPrice": float64(1)})
require.NoError(t, err)
assert.Equal(t, `tier("image", fixed(1)) * image_count`, preview.Expression)
assert.Empty(t, preview.UnsupportedReason)
for _, tc := range []struct {
name, rule string
nameRule int
endpoint constant.EndpointType
}{
{"catalog-image-prefix-opaque", "catalog-image-prefix-", model.NameRulePrefix, constant.EndpointTypeImageGeneration},
{"opaque-catalog-image-suffix", "-catalog-image-suffix", model.NameRuleSuffix, constant.EndpointTypeImageGeneration},
{"before-catalog-image-middle-after", "catalog-image-middle", model.NameRuleContains, constant.EndpointTypeImageGeneration},
{"opaque-catalog-video", "opaque-catalog-video", model.NameRuleExact, constant.EndpointTypeOpenAIVideo},
} {
t.Run("endpoint_metadata/"+tc.name, func(t *testing.T) {
endpoints, err := common.Marshal(map[string]string{string(tc.endpoint): "/v1/fixture"})
require.NoError(t, err)
require.NoError(t, db.Create(&model.Model{ModelName: tc.rule, NameRule: tc.nameRule, Endpoints: string(endpoints)}).Error)
converted, err := model.PreviewModelPricingConversion(tc.name, model.PricingValues{"ModelPrice": 1.0})
require.NoError(t, err)
if tc.endpoint == constant.EndpointTypeOpenAIVideo {
assert.Empty(t, converted.Expression)
assert.Equal(t, "Video pricing must be converted manually.", converted.UnsupportedReason)
} else {
assert.Equal(t, `tier("image", fixed(1)) * image_count`, converted.Expression)
}
})
}
sora := &model.Channel{Name: "Opaque video route", Models: "opaque-video-route", Type: constant.ChannelTypeSora, Group: "a,b"}
require.NoError(t, sora.Insert())
video, err := model.PreviewModelPricingConversion("opaque-video-route", model.PricingValues{"ModelPrice": 1.0})
require.NoError(t, err)
assert.Equal(t, "Video pricing must be converted manually.", video.UnsupportedReason)
})
}
}
func TestModelManagementDatabaseMatrix(t *testing.T) {
_, err := jsplugin.DefaultRegistry.Register(`
export const meta = {apiVersion: 1, key: "model-management-task", name: "Management task fixture", version: "1.0.0", author: {name: "Test"}, models: ["matrix-task"], fetchMode: "per_task", usageSchema: {seconds: {type: "number", unit: "second"}}};
@@ -1134,3 +1502,168 @@ func TestModelDeletionDatabaseMatrix(t *testing.T) {
})
}
}
func TestSharedModelPluginPricingDatabaseMatrix(t *testing.T) {
const name = "shared-model::priced"
const base = `tier("base", u("seconds") * 0.4)`
const variant = `tier("beta", u("credits") * 2)`
for _, dialect := range []struct{ kind, env string }{{"sqlite", ""}, {"mysql", "TEST_MYSQL_DSN"}, {"postgres", "TEST_POSTGRES_DSN"}} {
t.Run(dialect.kind, func(t *testing.T) {
if dialect.env != "" && os.Getenv(dialect.env) == "" {
t.Skip("set " + dialect.env + " to run this database")
}
db := modelManagementDB(t, dialect.kind, os.Getenv(dialect.env))
for _, spec := range []struct{ key, field, unit string }{{"matrix-alpha", "seconds", "second"}, {"matrix-beta", "credits", "credit"}} {
source := fmt.Sprintf(`
export const meta = {apiVersion:1,key:%q,name:%q,version:"1.0.0",author:{name:"Test"},models:[%q],fetchMode:"per_task",usageSchema:{%s:{type:"number",unit:%q}}};
export function buildSubmitRequest(){return {};}
export function parseSubmitResponse(){return {};}
export function buildQueryRequest(){return {};}
export function parseTaskResult(){return {};}
`, spec.key, spec.key, name, spec.field, spec.unit)
_, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{})
require.NoError(t, err)
t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(spec.key) })
}
// Representative existing options have no plugin-expression row.
baseJSON, err := common.Marshal(map[string]string{name: base})
require.NoError(t, err)
modeJSON, err := common.Marshal(map[string]string{name: "tiered_expr"})
require.NoError(t, err)
require.NoError(t, db.Create(&[]model.Option{
{Key: "ModelPrice", Value: `{"unrelated-model":0.75}`},
{Key: "billing_setting.billing_expr", Value: string(baseJSON)},
{Key: "billing_setting.billing_mode", Value: string(modeJSON)},
}).Error)
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
"billing_setting.billing_expr": string(baseJSON), "billing_setting.billing_mode": string(modeJSON),
}))
before, err := model.GetModelPricingSnapshot([]string{name})
require.NoError(t, err)
require.Len(t, before.Entries, 1)
require.Len(t, before.Entries[0].PluginVariants, 2)
assert.True(t, before.Entries[0].PluginVariants[0].Compatible)
assert.False(t, before.Entries[0].PluginVariants[1].Compatible)
assert.Equal(t, base, before.Entries[0].PluginVariants[1].Effective)
change := model.ModelPricingChange{ModelName: name, ExpectedVersion: before.Entries[0].Version, Pricing: model.PricingValues{
"billing_setting.billing_mode": "tiered_expr", "billing_setting.billing_expr": base,
billing_setting.PluginBillingExprOption: map[string]any{"matrix-beta": variant},
}}
require.NoError(t, model.UpdateModelPricing([]model.ModelPricingChange{change}))
loaded, err := model.GetModelPricingSnapshot([]string{name})
require.NoError(t, err)
assert.Equal(t, change.Pricing, loaded.Entries[0].Configured)
require.Len(t, loaded.Entries[0].PluginVariants, 2)
assert.True(t, loaded.Entries[0].PluginVariants[0].Compatible)
assert.Equal(t, base, loaded.Entries[0].PluginVariants[0].Effective)
assert.True(t, loaded.Entries[0].PluginVariants[1].Compatible)
assert.Equal(t, variant, loaded.Entries[0].PluginVariants[1].Configured)
var flat map[string]string
require.NoError(t, common.UnmarshalJsonStr(loaded.Options[billing_setting.PluginBillingExprOption], &flat))
assert.Equal(t, map[string]string{"matrix-beta::" + name: variant}, flat)
all, err := model.GetModelPricingSnapshot(nil)
require.NoError(t, err)
for _, entry := range all.Entries {
assert.NotEqual(t, "matrix-beta::"+name, entry.ModelName)
}
unrelated, err := model.GetModelPricingSnapshot([]string{"unrelated-model"})
require.NoError(t, err)
assert.Equal(t, float64(0.75), unrelated.Entries[0].Configured["ModelPrice"])
assert.NotContains(t, billing_setting.GetPricingSyncData(map[string]any{}), "plugin_billing_expr")
assert.NotContains(t, billing_setting.GetPricingSyncData(map[string]any{})["billing_expr"], "matrix-beta::"+name)
// Restart/re-read is idempotent; a second save based on its version succeeds.
require.NoError(t, db.AutoMigrate(&model.Option{}))
change.ExpectedVersion = loaded.Entries[0].Version
require.NoError(t, model.UpdateModelPricing([]model.ModelPricingChange{change}))
repeated, err := model.GetModelPricingSnapshot([]string{name})
require.NoError(t, err)
assert.Equal(t, loaded.Entries[0].Version, repeated.Entries[0].Version)
// Changing only the provider price changes the model version.
change.Pricing[billing_setting.PluginBillingExprOption] = map[string]any{"matrix-beta": `tier("free", u("credits") * 0)`}
require.NoError(t, model.UpdateModelPricing([]model.ModelPricingChange{change}))
updated, err := model.GetModelPricingSnapshot([]string{name})
require.NoError(t, err)
assert.NotEqual(t, loaded.Entries[0].Version, updated.Entries[0].Version)
assert.ErrorIs(t, model.UpdateModelPricing([]model.ModelPricingChange{change}), model.ErrModelPricingConflict)
// The legacy option API validates the full draft and cannot drop a required override.
response := modelManagementRequest(t, UpdateOption, http.MethodPut, "/api/option/", OptionUpdateRequest{Key: billing_setting.PluginBillingExprOption, Value: `{}`}, nil)
assert.Contains(t, response.Body.String(), `"success":false`)
assert.Contains(t, response.Body.String(), "matrix-beta")
afterFailure, err := model.GetModelPricingSnapshot([]string{name})
require.NoError(t, err)
assert.Equal(t, updated.Entries[0].Version, afterFailure.Entries[0].Version)
// Model-level legacy saves also skip providers with a stored override.
response = modelManagementRequest(t, UpdateOption, http.MethodPut, "/api/option/", OptionUpdateRequest{Key: "billing_setting.billing_expr", Value: string(baseJSON)}, nil)
assert.Contains(t, response.Body.String(), `"success":true`)
flat["matrix-beta::"+name] = variant
raw, err := common.Marshal(flat)
require.NoError(t, err)
response = modelManagementRequest(t, UpdateOption, http.MethodPut, "/api/option/", OptionUpdateRequest{Key: billing_setting.PluginBillingExprOption, Value: string(raw)}, nil)
assert.Contains(t, response.Body.String(), `"success":true`)
final, err := model.GetModelPricingSnapshot([]string{name})
require.NoError(t, err)
assert.Equal(t, variant, final.Entries[0].PluginVariants[1].Effective)
// A provider may disappear without making every legacy price save fail.
require.NoError(t, jsplugin.DefaultRegistry.Unregister("matrix-beta"))
stale, err := model.GetModelPricingSnapshot([]string{name})
require.NoError(t, err)
require.Len(t, stale.Entries[0].PluginVariants, 2)
assert.True(t, stale.Entries[0].PluginVariants[1].Stale)
assert.Equal(t, variant, stale.Entries[0].PluginVariants[1].Configured)
assert.Empty(t, stale.Entries[0].PluginVariants[1].Effective)
_, err = model.PreviewModelPricing(name, stale.Entries[0].Configured)
require.NoError(t, err)
ratioJSON, err := common.Marshal(map[string]float64{name: 2})
require.NoError(t, err)
response = modelManagementRequest(t, UpdateOption, http.MethodPut, "/api/option/", OptionUpdateRequest{Key: "ModelRatio", Value: string(ratioJSON)}, nil)
assert.Contains(t, response.Body.String(), `"success":true`)
response = modelManagementRequest(t, UpdateOption, http.MethodPut, "/api/option/", OptionUpdateRequest{Key: "billing_setting.billing_expr", Value: string(baseJSON)}, nil)
assert.Contains(t, response.Body.String(), `"success":true`)
final, err = model.GetModelPricingSnapshot([]string{name})
require.NoError(t, err)
assert.Equal(t, variant, final.Entries[0].PluginVariants[1].Configured)
// The versioned save also preserves an unchanged stale expression.
change = model.ModelPricingChange{ModelName: name, ExpectedVersion: final.Entries[0].Version, Pricing: final.Entries[0].Configured}
require.NoError(t, model.UpdateModelPricing([]model.ModelPricingChange{change}))
change.Pricing[billing_setting.PluginBillingExprOption] = map[string]any{"matrix-beta": "1"}
// Even if this process's cache is ahead of storage, a changed stale
// override is validated against the locked database value.
tamperedJSON, err := common.Marshal(map[string]string{"matrix-beta::" + name: "1"})
require.NoError(t, err)
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{billing_setting.PluginBillingExprOption: string(tamperedJSON)}))
require.ErrorContains(t, model.UpdateModelPricing([]model.ModelPricingChange{change}), "does not declare this model")
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{billing_setting.PluginBillingExprOption: final.Options[billing_setting.PluginBillingExprOption]}))
// Updating a plugin to stop declaring this model also leaves the
// stored override visible and inert, even with no active providers.
_, err = jsplugin.DefaultRegistry.Register(`
export const meta = {apiVersion:1,key:"matrix-beta",name:"Beta updated",version:"2.0.0",author:{name:"Test"},models:["replacement-model"],fetchMode:"per_task",usageSchema:{credits:{type:"number",unit:"credit"}}};
export function buildSubmitRequest(){return {};}
export function parseSubmitResponse(){return {};}
export function buildQueryRequest(){return {};}
export function parseTaskResult(){return {};}
`, jsplugin.Options{})
require.NoError(t, err)
require.NoError(t, jsplugin.DefaultRegistry.Unregister("matrix-alpha"))
stale, err = model.GetModelPricingSnapshot([]string{name})
require.NoError(t, err)
require.Len(t, stale.Entries[0].PluginVariants, 1)
assert.True(t, stale.Entries[0].PluginVariants[0].Stale)
assert.Equal(t, "Beta updated", stale.Entries[0].PluginVariants[0].PluginName)
require.NoError(t, model.UpdateModelPricingOptions(map[string]string{"ModelRatio": string(ratioJSON)}))
// Removing just the stale override succeeds and leaves other prices.
response = modelManagementRequest(t, UpdateOption, http.MethodPut, "/api/option/", OptionUpdateRequest{Key: billing_setting.PluginBillingExprOption, Value: `{}`}, nil)
assert.Contains(t, response.Body.String(), `"success":true`)
final, err = model.GetModelPricingSnapshot([]string{name})
require.NoError(t, err)
assert.Empty(t, final.Entries[0].PluginVariants)
assert.Equal(t, base, final.Entries[0].Configured["billing_setting.billing_expr"])
// A complete reset removes all provider keys, including models containing ::.
require.NoError(t, model.UpdateModelPricing([]model.ModelPricingChange{{ModelName: name, ExpectedVersion: final.Entries[0].Version, Reset: true}}))
final, err = model.GetModelPricingSnapshot([]string{name})
require.NoError(t, err)
assert.Empty(t, final.Entries[0].Configured)
assert.Equal(t, "{}", final.Options[billing_setting.PluginBillingExprOption])
})
}
}
+38
View File
@@ -18,6 +18,44 @@ func GetModelPricingConfig(c *gin.Context) {
common.ApiSuccess(c, snapshot)
}
func PreviewModelPricingConversion(c *gin.Context) {
var request struct {
ModelName string `json:"model_name"`
Pricing model.PricingValues `json:"pricing"`
}
if err := common.DecodeJson(c.Request.Body, &request); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()})
return
}
preview, err := model.PreviewModelPricingConversion(request.ModelName, request.Pricing)
if err != nil {
common.ApiError(c, err)
return
}
common.ApiSuccess(c, preview)
}
func PreviewModelPricing(c *gin.Context) {
var request struct {
ModelName string `json:"model_name"`
Pricing model.PricingValues `json:"pricing"`
}
if err := common.DecodeJson(c.Request.Body, &request); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()})
return
}
preview, err := model.PreviewModelPricing(request.ModelName, request.Pricing)
if err != nil {
common.ApiError(c, err)
return
}
common.ApiSuccess(c, model.ModelPricingDescription{
Effective: preview,
CacheWriteMode: model.ResolveCacheWriteMode(request.ModelName, request.Pricing),
BillingDetails: model.ResolveLegacyBillingDetails(request.ModelName, preview, request.Pricing),
})
}
func UpdateModelPricingConfig(c *gin.Context) {
var request struct {
Changes []model.ModelPricingChange `json:"changes"`
+89 -34
View File
@@ -5,6 +5,7 @@ import (
"errors"
"fmt"
"net/http"
"slices"
"strconv"
"strings"
"time"
@@ -292,7 +293,7 @@ func HandleOAuth(c *gin.Context) {
switch flow.Intent {
case model.AuthFlowIntentLogin:
handleOAuthLogin(c, provider, oauthUser, flow)
handleOAuthLogin(c, provider, oauthUser, token, flow)
case model.AuthFlowIntentVerify:
handleOAuthVerification(c, providerName, oauthUser, flow)
}
@@ -314,14 +315,14 @@ func handleOAuthVerification(c *gin.Context, provider string, oauthUser *oauth.O
common.ApiSuccess(c, proof)
}
func handleOAuthLogin(c *gin.Context, provider oauth.Provider, oauthUser *oauth.OAuthUser, flow *model.AuthFlow) {
func handleOAuthLogin(c *gin.Context, provider oauth.Provider, oauthUser *oauth.OAuthUser, token *oauth.OAuthToken, flow *model.AuthFlow) {
// 7. Find or create user
var payload oauthFlowPayload
if err := common.UnmarshalJsonStr(flow.Payload, &payload); err != nil {
writeSecurityOperationError(c, err)
return
}
user, err := findOrCreateOAuthUser(c, provider, oauthUser, payload.AffiliateCode)
user, migration, err := findOrCreateOAuthUser(c, provider, oauthUser, token, payload.AffiliateCode)
if err != nil {
if errors.Is(err, model.ErrEmailAlreadyTaken) {
common.ApiErrorI18n(c, i18n.MsgUserEmailAlreadyTaken)
@@ -334,6 +335,8 @@ func handleOAuthLogin(c *gin.Context, provider oauth.Provider, oauthUser *oauth.
common.ApiErrorI18n(c, i18n.MsgUserRegisterDisabled)
case *OAuthEmailAlreadyTakenError:
common.ApiErrorI18n(c, i18n.MsgUserEmailAlreadyTaken)
case *OAuthLegacyBindingNotConfirmedError:
common.ApiErrorI18n(c, i18n.MsgOAuthNotAutoLinked, providerParams(provider.GetName()))
default:
writeSecurityOperationError(c, err)
}
@@ -347,7 +350,7 @@ func handleOAuthLogin(c *gin.Context, provider oauth.Provider, oauthUser *oauth.
}
// 9. Setup login
setupLogin(user, c)
setupLogin(user, migration, c)
}
// handleOAuthBind handles binding OAuth account to existing user
@@ -377,10 +380,6 @@ func handleOAuthBind(c *gin.Context, providerName string, provider oauth.Provide
common.ApiErrorI18n(c, i18n.MsgOAuthAlreadyBound, providerParams(provider.GetName()))
return false, false
}
if legacyID, ok := oauthUser.Extra["legacy_id"].(string); ok && legacyID != "" && provider.IsUserIDTaken(legacyID) {
common.ApiErrorI18n(c, i18n.MsgOAuthAlreadyBound, providerParams(provider.GetName()))
return false, false
}
_, err = model.ConsumeAuthFlowWithAction(state, match, func(tx *gorm.DB, _ *model.AuthFlow) error {
if providerName == "telegram" {
return model.BindTelegramForSessionWithTx(tx, identity, oauthUser.ProviderUserID)
@@ -404,53 +403,91 @@ func handleOAuthBind(c *gin.Context, providerName string, provider oauth.Provide
return true, notificationFailed
}
// findOrCreateOAuthUser finds existing user or creates new user
func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *oauth.OAuthUser, affiliateCode string) (*model.User, error) {
// findOrCreateOAuthUser finds the existing user or creates a new one. For a
// legacy GitHub binding that still waits for the login verification, it also
// returns the rewrite to carry into the challenge.
func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *oauth.OAuthUser, token *oauth.OAuthToken, affiliateCode string) (*model.User, *service.LegacyGitHubMigration, error) {
user := &model.User{}
if provider.ProviderUserIDColumn() == "telegram_id" {
err := provider.FillUserByProviderID(user, oauthUser.ProviderUserID)
if errors.Is(err, gorm.ErrRecordNotFound) {
return nil, oauth.ErrTelegramAccountNotBound
return nil, nil, oauth.ErrTelegramAccountNotBound
}
return user, err
return user, nil, err
}
// Check if user already exists with new ID
if provider.IsUserIDTaken(oauthUser.ProviderUserID) {
err := provider.FillUserByProviderID(user, oauthUser.ProviderUserID)
if err != nil {
return nil, err
return nil, nil, err
}
// Check if user has been deleted
if user.Id == 0 {
return nil, &OAuthUserDeletedError{}
return nil, nil, &OAuthUserDeletedError{}
}
return user, nil
return user, nil, nil
}
// Try to find user with legacy ID (for GitHub migration from login to numeric ID)
if legacyID, ok := oauthUser.Extra["legacy_id"].(string); ok && legacyID != "" {
if provider.IsUserIDTaken(legacyID) {
err := provider.FillUserByProviderID(user, legacyID)
// Legacy GitHub bindings stored the login name, which only points at a
// candidate account. Values that are all digits are already numeric account
// IDs and never take part in the comparison.
legacyID, _ := oauthUser.Extra["legacy_id"].(string)
if strings.ContainsFunc(legacyID, func(r rune) bool { return r < '0' || r > '9' }) && provider.IsUserIDTaken(legacyID) {
if err := provider.FillUserByProviderID(user, legacyID); err != nil {
return nil, nil, err
}
if user.Id != 0 {
state, err := model.GetUserVerificationState(user.Id)
if err != nil {
return nil, err
return nil, nil, err
}
if user.Id != 0 {
// Found user with legacy ID, migrate to new ID
common.SysLog(fmt.Sprintf("[OAuth] Migrating user %d from legacy_id=%s to new_id=%s",
user.Id, legacyID, oauthUser.ProviderUserID))
if err := user.UpdateGitHubId(oauthUser.ProviderUserID); err != nil {
common.SysError(fmt.Sprintf("[OAuth] Failed to migrate user %d: %s", user.Id, err.Error()))
// Continue with login even if migration fails
if state.HasTwoFA || state.HasPasskey {
// The rewrite is written in the transaction that issues the session
// once the login verification completes.
return user, &service.LegacyGitHubMigration{GitHubID: oauthUser.ProviderUserID, LegacyID: legacyID}, nil
}
// Without a second factor, one of the addresses the provider has
// confirmed must match the account email. The list is fetched only here
// and never recorded.
reason, matched := "no_matching_evidence", false
if emailProvider, ok := provider.(oauth.VerifiedEmailProvider); ok && user.Email != "" {
emails, err := emailProvider.GetVerifiedEmails(c.Request.Context(), token)
if err != nil {
common.SysError(fmt.Sprintf("[OAuth] Failed to load verified emails for user %d: %s", user.Id, err.Error()))
reason = "verified_emails_unavailable"
}
return user, nil
accountEmail := model.NormalizeEmail(user.Email)
matched = slices.ContainsFunc(emails, func(email string) bool { return model.NormalizeEmail(email) == accountEmail })
}
if !matched {
recordLegacyGitHubBindingAudit(c, user, false, map[string]any{"legacy_id": legacyID, "provider_user_id": oauthUser.ProviderUserID, "reason": reason})
return nil, nil, &OAuthLegacyBindingNotConfirmedError{}
}
written := false
err = model.DB.Transaction(func(tx *gorm.DB) error {
var err error
written, err = model.MigrateLegacyGitHubBindingWithTx(tx, user.Id, legacyID, oauthUser.ProviderUserID)
return err
})
if err != nil {
return nil, nil, err
}
if written {
user.GitHubId = oauthUser.ProviderUserID
notificationFailed := service.NotifyAccountSecurityChange(user.Email, "Login account linked: "+provider.GetName()) != nil
recordLegacyGitHubBindingAudit(c, user, true, map[string]any{
"legacy_id": legacyID, "provider_user_id": oauthUser.ProviderUserID,
"verified_email_matched": true, "notification_failed": notificationFailed,
})
}
return user, nil, nil
}
}
// User doesn't exist, create new user if registration is enabled
if !common.RegisterEnabled {
return nil, &OAuthRegistrationDisabledError{}
return nil, nil, &OAuthRegistrationDisabledError{}
}
// Set up new user
@@ -476,9 +513,9 @@ func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *o
user.Email = model.NormalizeEmail(oauthUser.Email)
if err := model.EnsureEmailAvailable(user.Email, 0); err != nil {
if errors.Is(err, model.ErrEmailAlreadyTaken) {
return nil, &OAuthEmailAlreadyTakenError{}
return nil, nil, &OAuthEmailAlreadyTakenError{}
}
return nil, err
return nil, nil, err
}
}
user.Role = common.RoleCommonUser
@@ -512,7 +549,7 @@ func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *o
return nil
})
if err != nil {
return nil, err
return nil, nil, err
}
// Perform post-transaction tasks (logs, sidebar config, inviter rewards)
@@ -541,14 +578,24 @@ func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *o
return nil
})
if err != nil {
return nil, err
return nil, nil, err
}
// Perform post-transaction tasks
user.FinalizeOAuthUserCreation(inviterId)
}
return user, nil
return user, nil, nil
}
// recordLegacyGitHubBindingAudit records the outcome of a legacy GitHub binding
// rewrite as an account binding event. No session exists yet on the login path,
// so the row carries the user's own role instead of the request context.
func recordLegacyGitHubBindingAudit(c *gin.Context, user *model.User, success bool, params map[string]any) {
params["provider"], params["legacy_migration"], params["success"] = "github", true, success
model.RecordOperationAuditLog(user.Id, user.Role, auditContentEN("user.binding_bind", params), c.ClientIP(), "user.binding_bind", params, nil, &model.AuditRequestInfo{
Method: c.Request.Method, Route: c.FullPath(), Path: c.FullPath(), Status: c.Writer.Status(), Success: success,
}, c)
}
// Error types for OAuth
@@ -570,6 +617,14 @@ func (e *OAuthEmailAlreadyTakenError) Error() string {
return "email is already in use"
}
// OAuthLegacyBindingNotConfirmedError reports a legacy GitHub binding match that
// neither a second factor nor a confirmed provider email backed.
type OAuthLegacyBindingNotConfirmedError struct{}
func (e *OAuthLegacyBindingNotConfirmedError) Error() string {
return "legacy binding was not confirmed"
}
// handleOAuthError handles OAuth errors and returns translated message
func handleOAuthError(c *gin.Context, err error) {
switch e := err.(type) {
+89 -13
View File
@@ -1,6 +1,7 @@
package controller
import (
"errors"
"fmt"
"net/http"
"slices"
@@ -11,7 +12,6 @@ import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting"
"github.com/QuantumNous/new-api/setting/billing_setting"
@@ -136,6 +136,50 @@ type OptionUpdateRequest struct {
Value any `json:"value"`
}
func UpdatePasskeyDomains(c *gin.Context) {
var request struct {
RPID *string `json:"rp_id"`
LegacyRPIDs *string `json:"legacy_rp_ids"`
Origins *string `json:"origins"`
Preview bool `json:"preview"`
RemovalConfirmation string `json:"removal_confirmation"`
}
if err := common.DecodeJson(c.Request.Body, &request); err != nil || request.RPID == nil || request.LegacyRPIDs == nil || request.Origins == nil {
common.ApiErrorI18n(c, i18n.MsgInvalidParams)
return
}
change, err := model.UpdatePasskeyDomainOptions(map[string]string{
"passkey.rp_id": *request.RPID, "passkey.legacy_rp_ids": *request.LegacyRPIDs, "passkey.origins": *request.Origins,
}, request.Preview, request.RemovalConfirmation)
if err != nil {
writePasskeyDomainSettingsError(c, err)
if !request.Preview {
recordPasskeyDomainAudit(c, change, request.RemovalConfirmation != "", err)
}
return
}
if !request.Preview {
recordPasskeyDomainAudit(c, change, request.RemovalConfirmation != "", nil)
}
common.ApiSuccess(c, change)
}
func writePasskeyDomainSettingsError(c *gin.Context, err error) {
var removal *model.PasskeyDomainRemovalError
if errors.As(err, &removal) {
c.JSON(http.StatusConflict, gin.H{
"success": false, "code": "PASSKEY_RP_ID_REMOVAL_CONFIRMATION_REQUIRED",
"message": i18n.T(c, i18n.MsgPasskeyRPIDRemovalConfirmation), "data": removal.Change,
})
return
}
if errors.Is(err, system_setting.ErrPasskeyRPIDInvalid) {
writeSecurityOperationError(c, err)
return
}
common.ApiError(c, err)
}
func UpdateOption(c *gin.Context) {
var option OptionUpdateRequest
err := common.DecodeJson(c.Request.Body, &option)
@@ -359,25 +403,42 @@ func UpdateOption(c *gin.Context) {
models = append(models, modelName)
}
sort.Strings(models)
generation := jsplugin.DefaultRegistry.Generation()
storedVariants := billing_setting.GetPluginBillingExprCopy()
for _, modelName := range models {
expression := expressions[modelName]
if plugin, ok := generation.GetByModel(modelName); ok {
err = billing_setting.SmokeTestTaskExpr(expression, plugin.Meta.UsageSchema)
} else if target, resolved := model.ResolveTaskModelAlias(generation, modelName); resolved {
if plugin, ok := generation.Get(target.PluginKey); ok {
err = billing_setting.SmokeTestTaskExpr(expression, plugin.Meta.UsageSchema)
} else {
err = billing_setting.SmokeTestExpr(expression)
variants := make(map[string]any)
for key, expression := range storedVariants {
if plugin, name, ok := billing_setting.SplitPluginBillingExprKey(key); ok && name == modelName {
variants[plugin] = expression
}
} else {
err = billing_setting.SmokeTestExpr(expression)
}
err = model.ValidateModelPricing(modelName, model.PricingValues{
"billing_setting.billing_expr": expressions[modelName],
billing_setting.PluginBillingExprOption: variants,
})
if err != nil {
common.ApiErrorMsg(c, fmt.Sprintf("模型 %s 的计费表达式无效: %v", modelName, err))
return
}
}
case billing_setting.PluginBillingExprOption:
var expressions map[string]string
if err = common.UnmarshalJsonStr(option.Value.(string), &expressions); err != nil || expressions == nil {
common.ApiErrorMsg(c, "plugin billing expressions must be a JSON object")
return
}
for key, expression := range expressions {
plugin, name, valid := billing_setting.SplitPluginBillingExprKey(key)
if !valid {
common.ApiErrorMsg(c, "invalid plugin billing expression key: "+key)
return
}
if err = model.ValidateModelPricing(name, model.PricingValues{
billing_setting.PluginBillingExprOption: map[string]any{plugin: expression},
}); err != nil {
common.ApiErrorMsg(c, err.Error())
return
}
}
case "console_setting.api_info":
err = console_setting.ValidateConsoleSettings(option.Value.(string), "ApiInfo")
if err != nil {
@@ -415,9 +476,24 @@ func UpdateOption(c *gin.Context) {
return
}
}
if model.IsPasskeyDomainOption(option.Key) {
change, updateErr := model.UpdatePasskeyDomainOptions(map[string]string{option.Key: option.Value.(string)}, false, "")
if updateErr != nil {
writePasskeyDomainSettingsError(c, updateErr)
recordPasskeyDomainAudit(c, change, false, updateErr)
return
}
recordPasskeyDomainAudit(c, change, false, nil)
common.ApiSuccess(c, change)
return
}
err = model.UpdateOption(option.Key, option.Value.(string))
if err != nil {
common.ApiError(c, err)
if errors.Is(err, system_setting.ErrPasskeyRPIDInvalid) {
writeSecurityOperationError(c, err)
} else {
common.ApiError(c, err)
}
return
}
// 出于安全考虑只记录被修改的配置项名称,不记录配置值(可能含密钥等敏感信息)。
+53 -28
View File
@@ -26,6 +26,7 @@ type passkeyFinishRequest struct {
}
type passkeyVerifyBeginRequest struct {
RPID string `json:"rp_id"`
Scope string `json:"scope"`
Context json.RawMessage `json:"context,omitempty"`
}
@@ -42,7 +43,7 @@ func parsePasskeyFinishRequest(c *gin.Context) (*passkeyFinishRequest, error) {
}
func PasskeyRegisterBegin(c *gin.Context) {
if !system_setting.GetPasskeySettings().Enabled {
if !system_setting.PasskeySettingsSnapshot().Enabled {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "管理员未启用 Passkey 登录",
@@ -118,7 +119,7 @@ func PasskeyRegisterBegin(c *gin.Context) {
}
func PasskeyRegisterFinish(c *gin.Context) {
if !system_setting.GetPasskeySettings().Enabled {
if !system_setting.PasskeySettingsSnapshot().Enabled {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "管理员未启用 Passkey 登录",
@@ -142,12 +143,6 @@ func PasskeyRegisterFinish(c *gin.Context) {
return
}
wa, err := passkeysvc.BuildWebAuthn(c.Request)
if err != nil {
writeSecurityOperationError(c, err)
return
}
credentialRecord, err := model.GetPasskeyByUserID(user.Id)
if err != nil && !errors.Is(err, model.ErrPasskeyNotFound) {
writeSecurityOperationError(c, err)
@@ -171,6 +166,12 @@ func PasskeyRegisterFinish(c *gin.Context) {
writeSecurityOperationError(c, err)
return
}
wa, err := passkeysvc.BuildWebAuthnForRPID(c.Request, sessionData.RelyingPartyID)
if err != nil {
writeSecurityOperationError(c, err)
return
}
c.Set("passkey_rp_id", sessionData.RelyingPartyID)
if sessionData.UserVerification != protocol.VerificationRequired {
writeSecurityOperationError(c, model.ErrAuthFlowInvalid)
return
@@ -193,6 +194,7 @@ func PasskeyRegisterFinish(c *gin.Context) {
return
}
passkeyCredential.RPID = &sessionData.RelyingPartyID
if err := model.RegisterPasskeyForSession(identity, passkeyCredential); err != nil {
writeSecurityOperationError(c, err)
return
@@ -281,7 +283,7 @@ func PasskeyStatus(c *gin.Context) {
}
func PasskeyLoginBegin(c *gin.Context) {
if !system_setting.GetPasskeySettings().Enabled {
if !system_setting.PasskeySettingsSnapshot().Enabled {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "管理员未启用 Passkey 登录",
@@ -289,7 +291,16 @@ func PasskeyLoginBegin(c *gin.Context) {
return
}
wa, err := passkeysvc.BuildWebAuthn(c.Request)
var request struct {
RPID string `json:"rp_id"`
}
if c.Request.Body != nil && c.Request.ContentLength != 0 {
if common.DecodeJson(c.Request.Body, &request) != nil {
common.ApiErrorMsg(c, "无效的 Passkey 验证请求")
return
}
}
wa, rpIDs, err := passkeysvc.BuildLoginWebAuthn(c.Request, request.RPID, "")
if err != nil {
writeSecurityOperationError(c, err)
return
@@ -316,6 +327,7 @@ func PasskeyLoginBegin(c *gin.Context) {
"message": "",
"data": gin.H{
"options": assertion,
"rp_ids": rpIDs,
"flow_token": flowToken,
"expires_at": expiresAt,
},
@@ -323,7 +335,7 @@ func PasskeyLoginBegin(c *gin.Context) {
}
func PasskeyLoginFinish(c *gin.Context) {
if !system_setting.GetPasskeySettings().Enabled {
if !system_setting.PasskeySettingsSnapshot().Enabled {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "管理员未启用 Passkey 登录",
@@ -342,12 +354,6 @@ func PasskeyLoginFinish(c *gin.Context) {
return
}
wa, err := passkeysvc.BuildWebAuthn(c.Request)
if err != nil {
writeSecurityOperationError(c, err)
return
}
sessionData, _, err := passkeysvc.PopSessionDataFlow(
request.FlowToken,
model.AuthFlowPurposePasskeyLogin,
@@ -357,6 +363,12 @@ func PasskeyLoginFinish(c *gin.Context) {
writeSecurityOperationError(c, err)
return
}
wa, err := passkeysvc.BuildWebAuthnForRPID(c.Request, sessionData.RelyingPartyID)
if err != nil {
writeSecurityOperationError(c, err)
return
}
c.Set("passkey_rp_id", sessionData.RelyingPartyID)
if sessionData.UserVerification != protocol.VerificationRequired {
writeSecurityOperationError(c, model.ErrAuthFlowInvalid)
return
@@ -369,6 +381,10 @@ func PasskeyLoginFinish(c *gin.Context) {
return nil, fmt.Errorf("未找到 Passkey 凭证: %w", err)
}
if credential.RPID != nil && *credential.RPID != "" && *credential.RPID != sessionData.RelyingPartyID {
return nil, service.ErrVerificationFailed
}
// 通过凭证获取用户
user := &model.User{Id: credential.UserID}
if err := user.FillUserById(); err != nil {
@@ -415,7 +431,7 @@ func PasskeyLoginFinish(c *gin.Context) {
return
}
if err := model.UpdatePasskeyAssertionState(modelUser.Id, credential, time.Now()); err != nil {
if err := model.UpdatePasskeyAssertionState(modelUser.Id, credential, time.Now(), sessionData.RelyingPartyID); err != nil {
writeSecurityOperationError(c, err)
return
}
@@ -474,7 +490,7 @@ func AdminResetPasskey(c *gin.Context) {
}
func PasskeyVerifyBegin(c *gin.Context) {
if !system_setting.GetPasskeySettings().Enabled {
if !system_setting.PasskeySettingsSnapshot().Enabled {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "管理员未启用 Passkey 登录",
@@ -516,7 +532,11 @@ func PasskeyVerifyBegin(c *gin.Context) {
return
}
wa, err := passkeysvc.BuildWebAuthn(c.Request)
credentialRPID := ""
if credential.RPID != nil {
credentialRPID = *credential.RPID
}
wa, rpIDs, err := passkeysvc.BuildLoginWebAuthn(c.Request, request.RPID, credentialRPID)
if err != nil {
writeSecurityOperationError(c, err)
return
@@ -544,6 +564,7 @@ func PasskeyVerifyBegin(c *gin.Context) {
"message": "",
"data": gin.H{
"options": assertion,
"rp_ids": rpIDs,
"flow_token": flowToken,
"expires_at": expiresAt,
},
@@ -551,7 +572,7 @@ func PasskeyVerifyBegin(c *gin.Context) {
}
func PasskeyVerifyFinish(c *gin.Context) {
if !system_setting.GetPasskeySettings().Enabled {
if !system_setting.PasskeySettingsSnapshot().Enabled {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "管理员未启用 Passkey 登录",
@@ -576,12 +597,6 @@ func PasskeyVerifyFinish(c *gin.Context) {
return
}
wa, err := passkeysvc.BuildWebAuthn(c.Request)
if err != nil {
writeSecurityOperationError(c, err)
return
}
credential, err := model.GetPasskeyByUserID(user.Id)
if err != nil {
c.JSON(http.StatusOK, gin.H{
@@ -605,11 +620,21 @@ func PasskeyVerifyFinish(c *gin.Context) {
writeSecurityOperationError(c, err)
return
}
wa, err := passkeysvc.BuildWebAuthnForRPID(c.Request, sessionData.RelyingPartyID)
if err != nil {
writeSecurityOperationError(c, err)
return
}
c.Set("passkey_rp_id", sessionData.RelyingPartyID)
if sessionData.UserVerification != protocol.VerificationRequired {
writeSecurityOperationError(c, model.ErrAuthFlowInvalid)
return
}
if credential.RPID != nil && *credential.RPID != "" && *credential.RPID != sessionData.RelyingPartyID {
writeSecurityOperationError(c, service.ErrVerificationFailed)
return
}
waUser := passkeysvc.NewWebAuthnUser(user, credential)
validatedCredential, err := wa.ValidateLogin(waUser, *sessionData, parsedCredential)
if err != nil {
@@ -617,7 +642,7 @@ func PasskeyVerifyFinish(c *gin.Context) {
return
}
if err := model.UpdatePasskeyAssertionState(user.Id, validatedCredential, time.Now()); err != nil {
if err := model.UpdatePasskeyAssertionState(user.Id, validatedCredential, time.Now(), sessionData.RelyingPartyID); err != nil {
writeSecurityOperationError(c, err)
return
}
+754
View File
@@ -1,9 +1,15 @@
package controller
import (
"bytes"
"crypto/ecdsa"
"crypto/elliptic"
"crypto/rand"
"crypto/sha256"
"encoding/base64"
"errors"
"fmt"
"maps"
"net/http"
"net/http/httptest"
"strings"
@@ -11,18 +17,743 @@ import (
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/middleware"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/service"
passkeysvc "github.com/QuantumNous/new-api/service/passkey"
"github.com/QuantumNous/new-api/setting/system_setting"
"github.com/gin-gonic/gin"
"github.com/go-webauthn/webauthn/protocol"
"github.com/go-webauthn/webauthn/webauthn"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)
type passkeyTestBody struct {
*strings.Reader
}
type passkeyDomainBegin struct {
FlowToken string `json:"flow_token"`
RPIDs []string `json:"rp_ids"`
Options struct {
PublicKey struct {
Challenge string `json:"challenge"`
RPID string `json:"rpId"`
RP struct {
ID string `json:"id"`
} `json:"rp"`
} `json:"publicKey"`
} `json:"options"`
}
func TestGetStatusDoesNotExposePasskeyOrigins(t *testing.T) {
settings := system_setting.GetPasskeySettings()
originalSettings := *settings
originalOptionMap := common.OptionMap
originalServerAddress := system_setting.ServerAddress
t.Cleanup(func() {
*settings = originalSettings
common.OptionMap = originalOptionMap
system_setting.ServerAddress = originalServerAddress
})
common.OptionMap = map[string]string{}
system_setting.ServerAddress = "https://www.example.com"
*settings = system_setting.PasskeySettings{Enabled: true, RPID: "example.com"}
for _, origins := range []string{"https://www.example.com,https://private.example.com", "", "[]"} {
t.Run("origins="+origins, func(t *testing.T) {
settings.Origins = origins
response := httptest.NewRecorder()
context, _ := gin.CreateTestContext(response)
context.Request = httptest.NewRequest(http.MethodGet, "/api/status", nil)
GetStatus(context)
require.Equal(t, http.StatusOK, response.Code)
var payload struct {
Success bool `json:"success"`
Data map[string]any `json:"data"`
}
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &payload))
require.True(t, payload.Success)
assert.NotContains(t, payload.Data, "passkey_origins")
assert.NotContains(t, response.Body.String(), "private.example.com")
assert.Equal(t, true, payload.Data["passkey_login"])
assert.Equal(t, origins, settings.Origins)
})
}
}
func passkeyDomainRequest(t *testing.T, path string, payload any, identity service.AuthIdentity, origin string, handler gin.HandlerFunc) *httptest.ResponseRecorder {
t.Helper()
body, err := common.Marshal(payload)
require.NoError(t, err)
return securityEnrollmentRequest(http.MethodPost, path, string(body), "", identity, func(c *gin.Context) {
c.Request.Header.Set("Origin", origin)
handler(c)
})
}
func decodePasskeyDomainBegin(t *testing.T, response *httptest.ResponseRecorder) passkeyDomainBegin {
t.Helper()
var result securityEnrollmentResponse
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
require.True(t, result.Success, response.Body.String())
var begin passkeyDomainBegin
require.NoError(t, common.Unmarshal(result.Data, &begin))
return begin
}
// Reuse the existing software authenticator, replacing the domain-bound bytes
// and recomputing the signature exactly as a browser/authenticator would.
func passkeyDomainAssertion(t *testing.T, key *ecdsa.PrivateKey, challenge, rpID, origin string, userID int, verified bool) map[string]any {
t.Helper()
var assertion map[string]any
require.NoError(t, common.Unmarshal(securityPasskeyResponse(t, key, challenge, false, 1, verified), &assertion))
response := assertion["response"].(map[string]any)
clientData, err := common.Marshal(map[string]string{"type": "webauthn.get", "challenge": challenge, "origin": origin})
require.NoError(t, err)
authData, err := base64.RawURLEncoding.DecodeString(response["authenticatorData"].(string))
require.NoError(t, err)
rpHash := sha256.Sum256([]byte(rpID))
copy(authData, rpHash[:])
clientHash := sha256.Sum256(clientData)
signedHash := sha256.Sum256(append(append([]byte{}, authData...), clientHash[:]...))
signature, err := ecdsa.SignASN1(rand.Reader, key, signedHash[:])
require.NoError(t, err)
response["clientDataJSON"] = base64.RawURLEncoding.EncodeToString(clientData)
response["authenticatorData"] = base64.RawURLEncoding.EncodeToString(authData)
response["signature"] = base64.RawURLEncoding.EncodeToString(signature)
response["userHandle"] = base64.RawURLEncoding.EncodeToString([]byte(fmt.Sprint(userID)))
return assertion
}
func TestPasskeyDomainsPreserveCredentialsAcrossVerificationFlows(t *testing.T) {
for _, kind := range []string{"direct", "login factor", "sensitive action", "historical mixed case"} {
t.Run(kind, func(t *testing.T) {
user, identity := setupSecurityEnrollmentTest(t)
key := newSecurityLoginPasskey(t, user.Id)
settings := system_setting.GetPasskeySettings()
legacyRPID := "www.example.com"
if kind == "historical mixed case" {
legacyRPID = "WWW.example.com"
}
settings.LegacyRPIDs = legacyRPID
settings.Origins = "https://example.com,https://www.example.com"
beginPath, finishPath := "/api/user/passkey/login/begin", "/api/user/passkey/login/finish"
beginHandler, finishHandler := PasskeyLoginBegin, PasskeyLoginFinish
request := map[string]any{"rp_id": legacyRPID}
if kind == "login factor" {
pending, err := service.StartLoginVerification(user, "password", nil)
require.NoError(t, err)
request["flow_token"] = pending.FlowToken
beginPath, finishPath = "/api/user/login/passkey/begin", "/api/user/login/passkey/finish"
beginHandler, finishHandler = LoginPasskeyBegin, LoginPasskeyFinish
} else if kind == "sensitive action" {
request["scope"] = service.VerificationScopeAccessTokenGenerate
beginPath, finishPath = "/api/user/passkey/verify/begin", "/api/user/passkey/verify/finish"
beginHandler, finishHandler = PasskeyVerifyBegin, PasskeyVerifyFinish
}
begin := decodePasskeyDomainBegin(t, passkeyDomainRequest(t, beginPath, request, identity, "https://www.example.com", beginHandler))
assert.Equal(t, legacyRPID, begin.Options.PublicKey.RPID)
assert.Equal(t, []string{"example.com", legacyRPID}, begin.RPIDs)
before, err := model.GetPasskeyByUserID(user.Id)
require.NoError(t, err)
require.Nil(t, before.RPID)
finish := map[string]any{"flow_token": begin.FlowToken, "credential": passkeyDomainAssertion(t, key, begin.Options.PublicKey.Challenge, legacyRPID, "https://www.example.com", user.Id, true)}
if kind == "login factor" {
finish["flow_token"], finish["passkey_flow_token"] = request["flow_token"], begin.FlowToken
}
response := passkeyDomainRequest(t, finishPath, finish, identity, "https://www.example.com", finishHandler)
var result securityEnrollmentResponse
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
require.True(t, result.Success, response.Body.String())
after, err := model.GetPasskeyByUserID(user.Id)
require.NoError(t, err)
require.NotNil(t, after.RPID)
assert.Equal(t, legacyRPID, *after.RPID)
assert.Equal(t, before.CredentialID, after.CredentialID)
assert.Equal(t, before.PublicKey, after.PublicKey)
if kind == "sensitive action" {
var proof service.SecurityProof
require.NoError(t, common.Unmarshal(result.Data, &proof))
assert.Equal(t, service.VerificationScopeAccessTokenGenerate, proof.Scope)
}
count, err := model.CountActiveUserSessions(user.Id, time.Now().Unix())
require.NoError(t, err)
replay := passkeyDomainRequest(t, finishPath, finish, identity, "https://www.example.com", finishHandler)
require.NoError(t, common.Unmarshal(replay.Body.Bytes(), &result))
assert.False(t, result.Success)
afterCount, err := model.CountActiveUserSessions(user.Id, time.Now().Unix())
require.NoError(t, err)
assert.Equal(t, count, afterCount)
// Once identified, a browser hint cannot override the credential's binding.
wa, ids, err := passkeysvc.BuildLoginWebAuthn(httptest.NewRequest(http.MethodPost, beginPath, nil), "example.com", *after.RPID)
require.NoError(t, err)
assert.Equal(t, legacyRPID, wa.Config.RPID)
assert.Equal(t, []string{legacyRPID}, ids)
})
}
}
func TestPasskeyDomainsRejectInvalidAssertions(t *testing.T) {
for _, failure := range []string{"wrong RP ID", "unlisted origin", "origin outside RP scope", "wrong user", "bad signature", "missing user verification", "expired challenge", "removed domain", "known different domain"} {
t.Run(failure, func(t *testing.T) {
user, identity := setupSecurityEnrollmentTest(t)
key := newSecurityLoginPasskey(t, user.Id)
settings := system_setting.GetPasskeySettings()
settings.LegacyRPIDs = "www.example.com"
settings.Origins = "https://example.com,https://www.example.com"
begin := decodePasskeyDomainBegin(t, passkeyDomainRequest(t, "/api/user/passkey/login/begin", map[string]string{"rp_id": "www.example.com"}, identity, "https://www.example.com", PasskeyLoginBegin))
rpID, origin, handle := "www.example.com", "https://www.example.com", user.Id
switch failure {
case "wrong RP ID":
rpID = "example.com"
case "unlisted origin":
origin = "https://evil.example.com"
case "origin outside RP scope":
origin = "https://example.com"
case "wrong user":
handle++
case "expired challenge":
flow, err := model.GetAuthFlow(begin.FlowToken, model.AuthFlowMatch{Purpose: model.AuthFlowPurposePasskeyLogin})
require.NoError(t, err)
var payload struct {
SessionData webauthn.SessionData `json:"session_data"`
Security passkeysvc.FlowSecurity `json:"security"`
}
require.NoError(t, common.UnmarshalJsonStr(flow.Payload, &payload))
payload.SessionData.Expires = time.Now().Add(-time.Second)
encoded, err := common.Marshal(payload)
require.NoError(t, err)
require.NoError(t, model.DB.Model(flow).Update("payload", string(encoded)).Error)
case "removed domain":
settings.LegacyRPIDs = ""
case "known different domain":
require.NoError(t, model.DB.Model(&model.PasskeyCredential{}).Where("user_id = ?", user.Id).Update("rp_id", "example.com").Error)
}
assertion := passkeyDomainAssertion(t, key, begin.Options.PublicKey.Challenge, rpID, origin, handle, failure != "missing user verification")
if failure == "bad signature" {
assertion["response"].(map[string]any)["signature"] = "AA"
}
response := passkeyDomainRequest(t, "/api/user/passkey/login/finish", map[string]any{"flow_token": begin.FlowToken, "credential": assertion}, identity, "https://www.example.com", PasskeyLoginFinish)
var result securityEnrollmentResponse
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
assert.False(t, result.Success, response.Body.String())
assert.Empty(t, response.Header().Values("Set-Cookie"))
count, err := model.CountActiveUserSessions(user.Id, time.Now().Unix())
require.NoError(t, err)
assert.EqualValues(t, 1, count)
stored, err := model.GetPasskeyByUserID(user.Id)
require.NoError(t, err)
if failure == "known different domain" {
assert.Equal(t, "example.com", *stored.RPID)
} else {
assert.Nil(t, stored.RPID)
}
})
}
}
func TestPasskeyDomainChoicesRespectOriginAndConfiguration(t *testing.T) {
_, identity := setupSecurityEnrollmentTest(t)
settings := system_setting.GetPasskeySettings()
settings.LegacyRPIDs = "www.example.com"
settings.Origins = "https://example.com,https://www.example.com"
for _, test := range []struct {
origin, hint string
success bool
ids []string
}{
{"https://www.example.com", "", true, []string{"example.com", "www.example.com"}},
{"https://example.com", "", true, []string{"example.com"}},
{"https://example.com", "www.example.com", false, nil},
{"https://www.example.com", "unconfigured.example.com", false, nil},
{"https://attacker.test", "", false, nil},
} {
response := passkeyDomainRequest(t, "/api/user/passkey/login/begin", map[string]string{"rp_id": test.hint}, identity, test.origin, PasskeyLoginBegin)
var result securityEnrollmentResponse
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
assert.Equal(t, test.success, result.Success, "%+v", test)
if test.success {
assert.Equal(t, test.ids, decodePasskeyDomainBegin(t, response).RPIDs)
} else {
assert.Equal(t, "PASSKEY_RP_ID_UNAVAILABLE", result.Code)
}
}
settings.RPID, settings.LegacyRPIDs, settings.Origins, settings.AllowInsecureOrigin = "localhost", "", "http://localhost:3000,http://localhost:3001", true
for _, origin := range []string{"http://localhost:3000", "http://localhost:3001"} {
begin := decodePasskeyDomainBegin(t, passkeyDomainRequest(t, "/api/user/passkey/login/begin", nil, identity, origin, PasskeyLoginBegin))
assert.Equal(t, "localhost", begin.Options.PublicKey.RPID)
}
}
func setupPasskeyDomainOptions(t *testing.T) {
t.Helper()
require.NoError(t, model.DB.AutoMigrate(&model.Option{}))
common.OptionMapRWMutex.RLock()
previousOptions := maps.Clone(common.OptionMap)
previousAddress := system_setting.ServerAddress
common.OptionMapRWMutex.RUnlock()
t.Cleanup(func() {
common.OptionMapRWMutex.Lock()
common.OptionMap = previousOptions
system_setting.ServerAddress = previousAddress
common.OptionMapRWMutex.Unlock()
})
}
func TestPasskeyDomainRemovalProtectsExistingAndUnknownCredentials(t *testing.T) {
for _, binding := range []string{"www.example.com", ""} {
t.Run("binding="+binding, func(t *testing.T) {
user, _ := setupSecurityEnrollmentTest(t)
setupPasskeyDomainOptions(t)
settings := system_setting.GetPasskeySettings()
settings.LegacyRPIDs = "www.example.com"
settings.Origins = "https://example.com,https://www.example.com"
credential := &model.PasskeyCredential{UserID: user.Id, CredentialID: "existing", PublicKey: "key"}
if binding != "" {
credential.RPID = &binding
}
require.NoError(t, model.DB.Create(credential).Error)
assert.Error(t, model.UpdateOption("passkey.legacy_rp_ids", ""))
assert.Equal(t, "www.example.com", system_setting.PasskeySettingsSnapshot().LegacyRPIDs)
})
}
}
func TestPasskeyDomainRemovalWithoutAffectedCredentialsNeedsNoOverride(t *testing.T) {
user, _ := setupSecurityEnrollmentTest(t)
setupPasskeyDomainOptions(t)
settings := system_setting.GetPasskeySettings()
settings.LegacyRPIDs = "www.example.com,WWW.example.com"
upper := "WWW.example.com"
require.NoError(t, model.DB.Create(&model.PasskeyCredential{UserID: user.Id, RPID: &upper, CredentialID: "uppercase", PublicKey: "key"}).Error)
change, err := model.UpdatePasskeyDomainOptions(map[string]string{"passkey.legacy_rp_ids": upper}, false, "")
require.NoError(t, err)
assert.Equal(t, []string{"www.example.com"}, change.RemovedRPIDs)
assert.Zero(t, change.AffectedCredentials)
assert.Zero(t, change.UnknownCredentials)
assert.False(t, change.ConfirmationRequired)
assert.Equal(t, upper, settings.LegacyRPIDs)
}
func TestPasskeySnapshotDoesNotMaterializeDefaults(t *testing.T) {
setupSecurityEnrollmentTest(t)
setupPasskeyDomainOptions(t)
settings := system_setting.GetPasskeySettings()
settings.RPID, settings.Origins = "", ""
system_setting.ServerAddress = "https://www.example.com"
assert.Equal(t, "www.example.com", system_setting.PasskeySettingsSnapshot().EffectiveRPID())
assert.Empty(t, settings.RPID)
assert.Empty(t, settings.Origins)
}
func TestPasskeyDomainPreviewConfirmationAndAudit(t *testing.T) {
user, identity := setupSecurityEnrollmentTest(t)
setupPasskeyDomainOptions(t)
settings := system_setting.GetPasskeySettings()
settings.LegacyRPIDs = "www.example.com,WWW.example.com"
settings.Origins = "https://example.com,https://www.example.com"
lower, upper, empty := "www.example.com", "WWW.example.com", ""
for i, binding := range []*string{&lower, &upper, nil, &empty, &lower} {
credential := &model.PasskeyCredential{UserID: user.Id + i, RPID: binding, CredentialID: fmt.Sprint(i), PublicKey: "key"}
require.NoError(t, model.DB.Create(credential).Error)
if i == 4 {
require.NoError(t, model.DB.Delete(credential).Error)
}
}
request := map[string]any{"rp_id": "example.com", "legacy_rp_ids": upper, "origins": settings.Origins, "preview": true}
response := passkeyDomainRequest(t, "/api/option/passkey/domains", request, identity, "https://example.com", UpdatePasskeyDomains)
var result securityEnrollmentResponse
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
require.True(t, result.Success, response.Body.String())
var preview model.PasskeyDomainChange
require.NoError(t, common.Unmarshal(result.Data, &preview))
assert.Equal(t, []string{lower}, preview.RemovedRPIDs)
assert.EqualValues(t, 1, preview.AffectedCredentials)
assert.EqualValues(t, 2, preview.UnknownCredentials)
assert.True(t, preview.ConfirmationRequired)
require.NotEmpty(t, preview.RemovalConfirmation)
options, err := model.AllOption()
require.NoError(t, err)
assert.Empty(t, options, "preview must not materialize default rows")
assert.Equal(t, "www.example.com,WWW.example.com", settings.LegacyRPIDs)
assert.ErrorIs(t, model.UpdateOptionsBulk(map[string]string{"passkey.legacy_rp_ids": upper, "Notice": "must roll back"}), model.ErrPasskeyDomainRemovalConfirmation)
blocked := passkeyDomainRequest(t, "/api/option/", map[string]string{"key": "passkey.legacy_rp_ids", "value": upper}, identity, "https://example.com", UpdateOption)
assert.Equal(t, http.StatusConflict, blocked.Code)
assert.Contains(t, blocked.Body.String(), "PASSKEY_RP_ID_REMOVAL_CONFIRMATION_REQUIRED")
request["preview"] = false
request["removal_confirmation"] = preview.RemovalConfirmation
// A new unknown credential makes the reviewed impact stale.
require.NoError(t, model.DB.Create(&model.PasskeyCredential{UserID: user.Id + 10, CredentialID: "new", PublicKey: "key"}).Error)
response = passkeyDomainRequest(t, "/api/option/passkey/domains", request, identity, "https://example.com", UpdatePasskeyDomains)
assert.Equal(t, http.StatusConflict, response.Code)
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
require.NoError(t, common.Unmarshal(result.Data, &preview))
assert.EqualValues(t, 3, preview.UnknownCredentials)
options, err = model.AllOption()
require.NoError(t, err)
assert.Empty(t, options)
request["removal_confirmation"] = preview.RemovalConfirmation
response = passkeyDomainRequest(t, "/api/option/passkey/domains", request, identity, "https://example.com", UpdatePasskeyDomains)
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
require.True(t, result.Success, response.Body.String())
assert.Equal(t, upper, settings.LegacyRPIDs)
var audits []model.AuditLog
require.NoError(t, model.LOG_DB.Find(&audits).Error)
for _, audit := range audits {
if audit.Action == "option.passkey_domains_blocked" {
assert.False(t, audit.Success)
assert.Equal(t, http.StatusConflict, audit.Status)
}
if audit.Action == "option.passkey_domains_confirmed" {
assert.True(t, audit.Success)
require.NotNil(t, audit.Other.Op)
params, err := common.Marshal(audit.Other.Op.Params)
require.NoError(t, err)
var impact struct {
Known, Unknown int64
Confirmed bool
RemovedRPIDs []string `json:"removed_rp_ids"`
}
require.NoError(t, common.Unmarshal(params, &impact))
assert.EqualValues(t, 1, impact.Known)
assert.EqualValues(t, 3, impact.Unknown)
assert.True(t, impact.Confirmed)
assert.Equal(t, []string{lower}, impact.RemovedRPIDs)
}
}
encoded, err := common.Marshal(audits)
require.NoError(t, err)
assert.Contains(t, string(encoded), "option.passkey_domains_confirmed")
assert.Contains(t, string(encoded), "option.passkey_domains_blocked")
assert.NotContains(t, string(encoded), preview.RemovalConfirmation)
// Simulate an old node's cache after removal: persistence must use DB trust.
settings.LegacyRPIDs = lower + "," + upper
assert.ErrorIs(t, model.UpdatePasskeyAssertionState(user.Id, &webauthn.Credential{ID: []byte("0")}, time.Now(), lower), system_setting.ErrPasskeyRPIDUnavailable)
assert.ErrorIs(t, model.RegisterPasskeyForSession(identity, &model.PasskeyCredential{UserID: user.Id, RPID: &lower, CredentialID: "replacement", PublicKey: "key"}), system_setting.ErrPasskeyRPIDUnavailable)
}
func TestPasskeyDomainSettingsAcceptExactHTTPSInternalHosts(t *testing.T) {
for _, test := range []struct {
rpID, origins string
valid bool
}{
{"intranet", "https://intranet:8443", true},
{"intranet", "https://child.intranet", false},
{"intranet", "http://intranet", false},
{"com", "https://com", false},
{"127.0.0.1", "https://127.0.0.1", false},
{"intranet:8443", "https://intranet:8443", false},
{"localhost", "http://localhost:3000", true},
} {
t.Run(test.rpID+"/"+test.origins, func(t *testing.T) {
setupSecurityEnrollmentTest(t)
setupPasskeyDomainOptions(t)
change, err := model.UpdatePasskeyDomainOptions(map[string]string{"passkey.rp_id": test.rpID, "passkey.legacy_rp_ids": "", "passkey.origins": test.origins}, false, "")
if !test.valid {
assert.ErrorIs(t, err, system_setting.ErrPasskeyRPIDInvalid)
return
}
require.NoError(t, err)
assert.Equal(t, test.rpID, change.RPID)
assert.Equal(t, "example.com", change.LegacyRPIDs)
})
}
}
func TestPasskeyImplicitDomainChangesRetainPreviousTrust(t *testing.T) {
setupSecurityEnrollmentTest(t)
setupPasskeyDomainOptions(t)
settings := system_setting.GetPasskeySettings()
settings.RPID, settings.Origins = "", ""
system_setting.ServerAddress = "https://www.example.com"
require.NoError(t, model.UpdateOption("ServerAddress", "https://example.com"))
assert.Equal(t, []string{"example.com", "www.example.com"}, system_setting.PasskeySettingsSnapshot().RelyingPartyIDs())
assert.Equal(t, "https://www.example.com,https://example.com", settings.Origins)
// Moving the retained RP into the primary is not a removal.
change, err := model.UpdatePasskeyDomainOptions(map[string]string{"passkey.rp_id": "www.example.com", "passkey.legacy_rp_ids": ""}, false, "")
require.NoError(t, err)
assert.Empty(t, change.RemovedRPIDs)
assert.Equal(t, "example.com", change.LegacyRPIDs)
}
func TestPasskeyRegistrationAndDomainRemovalSerialize(t *testing.T) {
user, identity := setupSecurityEnrollmentTest(t)
setupPasskeyDomainOptions(t)
system_setting.GetPasskeySettings().Origins = "https://example.com,https://www.example.com"
require.NoError(t, model.UpdateOption("passkey.legacy_rp_ids", "www.example.com"))
registering, release := make(chan struct{}), make(chan struct{})
require.NoError(t, model.DB.Callback().Create().Before("gorm:create").Register("passkey_registration_barrier", func(tx *gorm.DB) {
if _, ok := tx.Statement.Dest.(*model.PasskeyCredential); ok {
close(registering)
<-release
}
}))
t.Cleanup(func() { _ = model.DB.Callback().Create().Remove("passkey_registration_barrier") })
rpID := "www.example.com"
registered, removed := make(chan error, 1), make(chan error, 1)
go func() {
registered <- model.RegisterPasskeyForSession(identity, &model.PasskeyCredential{UserID: user.Id, RPID: &rpID, CredentialID: "concurrent", PublicKey: "key"})
}()
<-registering
removing := make(chan struct{})
go func() {
close(removing)
removed <- model.UpdateOption("passkey.legacy_rp_ids", "")
}()
<-removing
// A concurrent snapshot must remain read-only while the write is pending.
assert.Equal(t, []string{"example.com", "www.example.com"}, system_setting.PasskeySettingsSnapshot().RelyingPartyIDs())
close(release)
require.NoError(t, <-registered)
assert.ErrorIs(t, <-removed, model.ErrPasskeyDomainRemovalConfirmation)
stored, err := model.GetPasskeyByUserID(user.Id)
require.NoError(t, err)
require.NotNil(t, stored.RPID)
assert.Equal(t, rpID, *stored.RPID)
}
func TestPasskeyDomainEndpointRequiresRoot(t *testing.T) {
user, _ := setupSecurityEnrollmentTest(t)
setupPasskeyDomainOptions(t)
bundle, err := service.CreateLoginSession(user.Id, "password", "127.0.0.1", "domain-settings-test")
require.NoError(t, err)
router := gin.New()
router.PUT("/api/option/passkey/domains", middleware.RootAuth(), UpdatePasskeyDomains)
for _, bearer := range []string{"", bundle.AccessToken} {
request := httptest.NewRequest(http.MethodPut, "/api/option/passkey/domains", strings.NewReader(`{"rp_id":"example.com","legacy_rp_ids":"","origins":"https://example.com","preview":true}`))
if bearer != "" {
request.Header.Set("Authorization", "Bearer "+bearer)
}
response := httptest.NewRecorder()
router.ServeHTTP(response, request)
assert.Contains(t, []int{http.StatusUnauthorized, http.StatusForbidden}, response.Code)
}
options, err := model.AllOption()
require.NoError(t, err)
assert.Empty(t, options)
}
func TestPasskeyDefaultsDoNotRestrictUnrelatedServerAddressSetup(t *testing.T) {
setupSecurityEnrollmentTest(t)
setupPasskeyDomainOptions(t)
settings := system_setting.GetPasskeySettings()
settings.Enabled, settings.RPID, settings.LegacyRPIDs, settings.Origins = false, "", "", ""
system_setting.ServerAddress = ""
require.NoError(t, model.UpdateOption("ServerAddress", "http://127.0.0.1:3000"))
assert.Equal(t, "http://127.0.0.1:3000", system_setting.ServerAddress)
assert.Empty(t, settings.RPID)
assert.Empty(t, settings.LegacyRPIDs)
}
func TestPasskeyRegistrationRejectsUnavailableRequestOrigin(t *testing.T) {
_, identity := setupSecurityEnrollmentTest(t)
settings := system_setting.GetPasskeySettings()
settings.RPID = "www.example.com"
settings.LegacyRPIDs = "example.com"
settings.Origins = "https://example.com,https://www.example.com"
proof := issueSecurityEnrollmentProof(t, identity, service.VerificationOperation{Scope: "passkey.register"}, "password")
response := securityEnrollmentRequest(http.MethodPost, "/api/user/passkey/register/begin", "", proof, identity, func(c *gin.Context) {
c.Request.Header.Set("Origin", "https://example.com")
PasskeyRegisterBegin(c)
})
var result securityEnrollmentResponse
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
assert.False(t, result.Success)
assert.Equal(t, "PASSKEY_RP_ID_UNAVAILABLE", result.Code)
var count int64
require.NoError(t, model.DB.Model(&model.AuthFlow{}).Where("purpose = ?", model.AuthFlowPurposePasskeyRegister).Count(&count).Error)
assert.Zero(t, count)
}
func TestPasskeyRPIDRotationRetainsDefaultsAndRollsBackOnFailure(t *testing.T) {
setupSecurityEnrollmentTest(t)
setupPasskeyDomainOptions(t)
settings := system_setting.GetPasskeySettings()
settings.RPID, settings.LegacyRPIDs = "", ""
settings.Origins = "https://www.example.com,https://example.com,https://login.example.com"
system_setting.ServerAddress = "https://www.example.com"
require.NoError(t, model.DB.Create(&model.Option{Key: "passkey.rp_id", Value: ""}).Error)
require.NoError(t, model.UpdateOption("passkey.rp_id", "example.com"))
assert.Equal(t, []string{"example.com", "www.example.com"}, system_setting.PasskeySettingsSnapshot().RelyingPartyIDs())
// Simulate a second node with an older local cache. Persisted domains win.
system_setting.GetPasskeySettings().RPID = "www.example.com"
require.NoError(t, model.UpdateOption("passkey.rp_id", "login.example.com"))
assert.Equal(t, []string{"login.example.com", "www.example.com", "example.com"}, system_setting.PasskeySettingsSnapshot().RelyingPartyIDs())
before, err := model.AllOption()
require.NoError(t, err)
snapshot := system_setting.PasskeySettingsSnapshot()
failure := errors.New("synthetic option write failure")
require.NoError(t, model.DB.Callback().Update().Before("gorm:update").Register("passkey_option_failure", func(tx *gorm.DB) {
if option, ok := tx.Statement.Dest.(*model.Option); ok && option.Key == "passkey.legacy_rp_ids" {
tx.AddError(failure)
}
}))
assert.ErrorIs(t, model.UpdateOptionsBulk(map[string]string{"passkey.rp_id": "next.example.com", "passkey.legacy_rp_ids": "www.example.com"}), failure)
require.NoError(t, model.DB.Callback().Update().Remove("passkey_option_failure"))
after, err := model.AllOption()
require.NoError(t, err)
assert.ElementsMatch(t, before, after)
assert.Equal(t, snapshot, system_setting.PasskeySettingsSnapshot())
for _, invalid := range []string{"*.example.com", "https://example.com", "example.com:443", "example.com/path", "com"} {
assert.ErrorIs(t, model.UpdateOption("passkey.rp_id", invalid), system_setting.ErrPasskeyRPIDInvalid)
assert.ErrorIs(t, model.UpdateOption("passkey.legacy_rp_ids", invalid), system_setting.ErrPasskeyRPIDInvalid)
}
assert.Equal(t, snapshot, system_setting.PasskeySettingsSnapshot())
// Older settings accepted mixed case. Rotation must preserve its hash input.
require.NoError(t, model.DB.Model(&model.Option{}).Where(&model.Option{Key: "passkey.rp_id"}).Update("value", "WWW.example.com").Error)
require.NoError(t, model.UpdateOption("passkey.rp_id", "next.example.com"))
assert.Contains(t, system_setting.PasskeySettingsSnapshot().RelyingPartyIDs(), "WWW.example.com")
}
func TestPasskeyRPIDRotationKeepsInFlightRegistration(t *testing.T) {
user, identity := setupSecurityEnrollmentTest(t)
setupPasskeyDomainOptions(t)
system_setting.GetPasskeySettings().Origins = "https://example.com,https://www.example.com"
proof := issueSecurityEnrollmentProof(t, identity, service.VerificationOperation{Scope: "passkey.register"}, "password")
begin := decodePasskeyDomainBegin(t, securityEnrollmentRequest(http.MethodPost, "/api/user/passkey/register/begin", "", proof, identity, PasskeyRegisterBegin))
assert.Equal(t, "example.com", begin.Options.PublicKey.RP.ID)
require.NoError(t, model.UpdateOption("passkey.rp_id", "www.example.com"))
key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
require.NoError(t, err)
response := passkeyDomainRequest(t, "/api/user/passkey/register/finish", passkeyFinishRequest{FlowToken: begin.FlowToken, Credential: securityPasskeyResponse(t, key, begin.Options.PublicKey.Challenge, true, 0)}, identity, "https://example.com", PasskeyRegisterFinish)
var result securityEnrollmentResponse
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
require.True(t, result.Success, response.Body.String())
credential, err := model.GetPasskeyByUserID(user.Id)
require.NoError(t, err)
require.NotNil(t, credential.RPID)
assert.Equal(t, "example.com", *credential.RPID)
// A later registration uses the current primary, not the retained domain.
wa, err := passkeysvc.BuildWebAuthn(httptest.NewRequest(http.MethodPost, "https://www.example.com/api/user/passkey/register/begin", nil))
require.NoError(t, err)
options, _, err := wa.BeginRegistration(passkeysvc.NewWebAuthnUser(user, credential))
require.NoError(t, err)
assert.Equal(t, "www.example.com", options.Response.RelyingParty.ID)
}
func TestPasskeyDomainFailuresDoNotLogCredentials(t *testing.T) {
gin.SetMode(gin.TestMode)
var output bytes.Buffer
common.LogWriterMu.Lock()
previousWriter := gin.DefaultErrorWriter
gin.DefaultErrorWriter = &output
common.LogWriterMu.Unlock()
t.Cleanup(func() {
common.LogWriterMu.Lock()
gin.DefaultErrorWriter = previousWriter
common.LogWriterMu.Unlock()
})
response := securityEnrollmentRequest(http.MethodPost, "/api/user/passkey/login/finish", "", "", service.AuthIdentity{}, func(c *gin.Context) {
c.Set("passkey_rp_id", "www.example.com")
writeSecurityOperationError(c, protocol.ErrVerification.WithDetails("private-flow-token").WithInfo("private-challenge"))
})
assert.Contains(t, response.Body.String(), "SECURITY_VERIFICATION_FAILED")
assert.Contains(t, output.String(), "verification_error")
assert.Contains(t, output.String(), "www.example.com")
for _, secret := range []string{"private-flow-token", "private-challenge"} {
assert.NotContains(t, response.Body.String(), secret)
assert.NotContains(t, output.String(), secret)
}
}
func TestPasskeyDomainErrorsRespectRequestLanguage(t *testing.T) {
user, _ := setupSecurityEnrollmentTest(t)
setupPasskeyDomainOptions(t)
system_setting.GetPasskeySettings().LegacyRPIDs = "www.example.com"
require.NoError(t, model.DB.Create(&model.PasskeyCredential{UserID: user.Id, CredentialID: "unknown-domain", PublicKey: "key"}).Error)
for _, locale := range []struct {
language, invalid, unavailable, removal string
}{
{"zh-CN", "通行密钥域名无效。请填写域名,不包含协议、端口、路径或通配符。", "此通行密钥域名无法在当前网站使用。请前往原网站或选择其他验证方式。", "请核对受影响的通行密钥并确认删除域名。配置或影响范围发生变化后,需要重新确认。"},
{"zh-TW", "通行金鑰網域無效。請填寫網域,不包含通訊協定、連接埠、路徑或萬用字元。", "此通行金鑰網域無法在目前網站使用。請前往原網站或選擇其他驗證方式。", "請核對受影響的通行金鑰並確認刪除網域。設定或影響範圍變更後,需要重新確認。"},
{"en", "Invalid Passkey domain. Enter a domain without a scheme, port, path or wildcard.", "This Passkey domain is not available on this website. Use its original website or another verification method.", "Review the affected Passkeys and confirm the domain removal. If the settings or impact have changed, confirmation is required again."},
} {
t.Run(locale.language, func(t *testing.T) {
for _, request := range []struct {
path, body, code, message string
handler gin.HandlerFunc
}{
{"/api/option/", `{"key":"passkey.rp_id","value":"localhost:3000"}`, "PASSKEY_RP_ID_INVALID", locale.invalid, UpdateOption},
{"/api/option/", `{"key":"passkey.legacy_rp_ids","value":"localhost:3001"}`, "PASSKEY_RP_ID_INVALID", locale.invalid, UpdateOption},
{"/api/option/", `{"key":"passkey.legacy_rp_ids","value":""}`, "PASSKEY_RP_ID_REMOVAL_CONFIRMATION_REQUIRED", locale.removal, UpdateOption},
{"/api/user/passkey/login/begin", `{"rp_id":"unconfigured.example.com"}`, "PASSKEY_RP_ID_UNAVAILABLE", locale.unavailable, PasskeyLoginBegin},
} {
response := httptest.NewRecorder()
c, _ := gin.CreateTestContext(response)
method := http.MethodPost
if request.path == "/api/option/" {
method = http.MethodPut
}
c.Request = httptest.NewRequest(method, request.path, strings.NewReader(request.body))
c.Request.Header.Set("Accept-Language", locale.language)
c.Request.Header.Set("Origin", "https://example.com")
c.Set("id", user.Id)
c.Set("role", common.RoleRootUser)
request.handler(c)
var result securityEnrollmentResponse
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
expectedStatus := http.StatusOK
if request.code == "PASSKEY_RP_ID_REMOVAL_CONFIRMATION_REQUIRED" {
expectedStatus = http.StatusConflict
}
assert.Equal(t, expectedStatus, response.Code)
assert.False(t, result.Success)
assert.Equal(t, request.code, result.Code)
assert.Equal(t, request.message, result.Message)
}
})
}
}
func TestPasskeyRPIDMigrationPreservesExistingCredentials(t *testing.T) {
for _, upgrade := range []bool{false, true} {
t.Run(fmt.Sprintf("upgrade=%t", upgrade), func(t *testing.T) {
user, identity := setupSecurityEnrollmentTest(t)
key := newSecurityLoginPasskey(t, user.Id)
before, err := model.GetPasskeyByUserID(user.Id)
require.NoError(t, err)
if upgrade {
// This schema is copied from the latest released model at
// v1.0.0-rc.36, with the same GORM and database driver versions.
require.NoError(t, model.DB.Migrator().DropTable(&model.PasskeyCredential{}))
require.NoError(t, model.DB.AutoMigrate(&passkeyCredentialBeforeRPID{}))
encoded, err := common.Marshal(before)
require.NoError(t, err)
var legacy passkeyCredentialBeforeRPID
require.NoError(t, common.Unmarshal(encoded, &legacy))
require.NoError(t, model.DB.Create(&legacy).Error)
}
pool, err := model.DB.DB()
require.NoError(t, err)
for range 2 {
// Each migration represents a new application startup. Discard
// PostgreSQL's old SELECT * prepared plans along with its old pool.
pool.SetMaxIdleConns(0)
pool.SetMaxIdleConns(2)
require.NoError(t, model.DB.AutoMigrate(&model.PasskeyCredential{}))
}
after, err := model.GetPasskeyByUserID(user.Id)
require.NoError(t, err)
assert.Equal(t, before, after)
assert.True(t, model.DB.Migrator().HasIndex(&model.PasskeyCredential{}, "idx_passkey_credentials_deleted_at"))
assert.Error(t, model.DB.Create(&model.PasskeyCredential{UserID: user.Id, CredentialID: "another-key", PublicKey: "key"}).Error)
assert.Error(t, model.DB.Create(&model.PasskeyCredential{UserID: user.Id + 1, CredentialID: before.CredentialID, PublicKey: "key"}).Error)
system_setting.GetPasskeySettings().LegacyRPIDs = "www.example.com"
system_setting.GetPasskeySettings().Origins = "https://example.com,https://www.example.com"
begin := decodePasskeyDomainBegin(t, passkeyDomainRequest(t, "/api/user/passkey/login/begin", map[string]string{"rp_id": "www.example.com"}, identity, "https://www.example.com", PasskeyLoginBegin))
response := passkeyDomainRequest(t, "/api/user/passkey/login/finish", map[string]any{"flow_token": begin.FlowToken, "credential": passkeyDomainAssertion(t, key, begin.Options.PublicKey.Challenge, "www.example.com", "https://www.example.com", user.Id, true)}, identity, "https://www.example.com", PasskeyLoginFinish)
var result securityEnrollmentResponse
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
assert.True(t, result.Success, response.Body.String())
})
}
}
func (*passkeyTestBody) Close() error { return nil }
func TestParsePasskeyFinishRequestDoesNotRewriteRequestBody(t *testing.T) {
@@ -70,3 +801,26 @@ func TestPasskeyRegisterFinishRejectsUnapprovedFlowWithoutConsumingIt(t *testing
require.NoError(t, err)
assert.Nil(t, flow.ConsumedAt)
}
type passkeyCredentialBeforeRPID struct {
ID int `json:"id" gorm:"primaryKey"`
UserID int `json:"user_id" gorm:"uniqueIndex;not null"`
CredentialID string `json:"credential_id" gorm:"type:varchar(512);uniqueIndex;not null"` // base64 encoded
PublicKey string `json:"public_key" gorm:"type:text;not null"` // base64 encoded
AttestationType string `json:"attestation_type" gorm:"type:varchar(255)"`
AAGUID string `json:"aaguid" gorm:"type:varchar(512)"` // base64 encoded
SignCount uint32 `json:"sign_count" gorm:"default:0"`
CloneWarning bool `json:"clone_warning"`
UserPresent bool `json:"user_present"`
UserVerified bool `json:"user_verified"`
BackupEligible bool `json:"backup_eligible"`
BackupState bool `json:"backup_state"`
Transports string `json:"transports" gorm:"type:text"`
Attachment string `json:"attachment" gorm:"type:varchar(32)"`
LastUsedAt *time.Time `json:"last_used_at"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
DeletedAt gorm.DeletedAt `json:"-" gorm:"index"`
}
func (passkeyCredentialBeforeRPID) TableName() string { return "passkey_credentials" }
+96
View File
@@ -0,0 +1,96 @@
package controller
import (
"bytes"
"context"
"errors"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/gin-gonic/gin"
"github.com/tidwall/gjson"
"net/http"
"net/http/httptest"
"strings"
)
var errDetectionServiceUnavailable = errors.New("detection service unreachable")
const localPelicanID = "local-pelican"
const pelicanPrompt = "Create a complete self-contained HTML page showing a pelican riding a bicycle. Draw the pelican and bicycle with inline SVG and animate pedaling and rotating wheels using CSS. Make the scene visually polished. Do not load external resources or use JavaScript. Return only the complete HTML document."
func detectionHTML(raw []byte, contentType string) string {
result := gjson.ParseBytes(raw)
if strings.Contains(contentType, "text/event-stream") {
result = gjson.Result{}
for line := range strings.SplitSeq(string(raw), "\n") {
if strings.HasPrefix(line, "data:") {
event := gjson.Parse(strings.TrimSpace(strings.TrimPrefix(line, "data:")))
if event.Get("type").String() == "response.completed" {
result = event.Get("response")
}
}
}
}
if result.Get("status").String() != "completed" {
return ""
}
var output strings.Builder
for _, item := range result.Get("output").Array() {
if item.Get("type").String() != "message" {
continue
}
for _, part := range item.Get("content").Array() {
if part.Get("type").String() == "output_text" {
output.WriteString(part.Get("text").String())
}
}
}
html := strings.TrimSpace(output.String())
lower := strings.ToLower(html)
start := strings.Index(lower, "<!doctype html")
if start < 0 {
start = strings.Index(lower, "<html")
}
end := strings.LastIndex(lower, "</html>")
if start < 0 || end < start || len(html) > 2<<20 {
return ""
}
return html[start : end+7]
}
func processLocalPelican(ctx context.Context, run model.DetectionRun) {
if ctx.Err() != nil {
return
}
if model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Updates(map[string]any{"remote_id": localPelicanID, "quality": "unknown", "result": ""}).Error != nil {
return
}
var call model.DetectionCall
if model.DB.First(&call, "id = ?", run.ID).Error != nil {
return
}
if call.State == "running" {
return
}
if call.State == "ready" {
payload, _ := common.Marshal(map[string]any{"model": "gpt-6-astra", "input": pelicanPrompt})
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequestWithContext(ctx, http.MethodPost, "/api/detection-proxy/v1/responses", bytes.NewReader(payload))
c.Request.Header.Set("Authorization", "Bearer "+pelicanCredential(run.ID))
DetectionProxy(c)
if model.DB.First(&call, "id = ?", run.ID).Error != nil {
return
}
}
if call.State != "finished" {
return
}
if model.DB.First(&run, "id = ?", run.ID).Error != nil {
return
}
if run.HTML != "" {
finishDetection(run, "succeeded", "")
} else {
finishDetection(run, "failed", "local_generation_failed")
}
}
+14
View File
@@ -0,0 +1,14 @@
package controller
import (
"github.com/stretchr/testify/assert"
"testing"
)
func TestDetectionHTMLRequiresCompletedDocument(t *testing.T) {
raw := []byte(`{"status":"completed","output":[{"type":"message","content":[{"type":"output_text","text":"<html><body>pelican</body></html>"}]}]}`)
assert.Equal(t, "<html><body>pelican</body></html>", detectionHTML(raw, "application/json"))
assert.Empty(t, detectionHTML([]byte(`{"status":"incomplete","output":[]}`), "application/json"))
sse := []byte("data: {\"type\":\"response.completed\",\"response\":" + string(raw) + "}\n\n")
assert.Equal(t, detectionHTML(raw, "application/json"), detectionHTML(sse, "text/event-stream"))
}
@@ -0,0 +1,50 @@
package controller
import (
"context"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/billingexpr"
"github.com/QuantumNous/new-api/service"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
"net/http"
"net/http/httptest"
"testing"
"time"
)
func TestLocalPelicanGeneratesOnceWithUnknownVerdict(t *testing.T) {
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
conn, err := db.DB()
require.NoError(t, err)
conn.SetMaxOpenConns(1)
old := model.DB
model.DB = db
t.Cleanup(func() { model.DB = old; _ = conn.Close() })
require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.DetectionRun{}, &model.DetectionCall{}, &model.DetectionSchedule{}))
calls := 0
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
calls++
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"status":"completed","usage":{"input_tokens":10,"output_tokens":5},"output":[{"type":"message","content":[{"type":"output_text","text":"<html><body>pelican</body></html>"}]}]}`))
}))
defer upstream.Close()
url := upstream.URL
channel := model.Channel{Type: 1, Key: "test-only", BaseURL: &url, Status: common.ChannelStatusEnabled, Models: "gpt-6-astra"}
require.NoError(t, db.Create(&channel).Error)
price, _ := common.Marshal(service.DetectionPrice{Snapshot: billingexpr.BillingSnapshot{ExprString: "p+c", ExprHash: billingexpr.ExprHashString("p+c"), GroupRatio: 1, QuotaPerUnit: 500000, ExprVersion: 1}})
run := model.DetectionRun{ID: "local-test", RequestKey: "local-test", Kind: "pelican", ChannelID: channel.Id, Status: "running", Deadline: time.Now().Unix() + 60, Reserved: 100, Price: string(price)}
require.NoError(t, db.Create(&run).Error)
require.NoError(t, db.Create(&model.DetectionCall{ID: run.ID, RunID: run.ID, SecretHash: model.DetectionSecretHash(pelicanCredential(run.ID)), State: "ready", ExpiresAt: run.Deadline}).Error)
processLocalPelican(context.Background(), run)
require.NoError(t, db.First(&run, "id = ?", run.ID).Error)
require.Equal(t, "succeeded", run.Status)
require.Equal(t, "unknown", run.Quality)
require.Contains(t, run.HTML, "pelican")
require.True(t, run.Settled)
processLocalPelican(context.Background(), run)
require.Equal(t, 1, calls)
}
+4 -13
View File
@@ -53,9 +53,10 @@ func GetPerfMetrics(c *gin.Context) {
}
result, err := perfmetrics.Query(perfmetrics.QueryParams{
Model: modelName,
Group: c.Query("group"),
Hours: hours,
Model: modelName,
Group: c.Query("group"),
Hours: hours,
AllowedGroups: append(lo.Keys(ratio_setting.GetGroupRatioCopy()), "auto"),
})
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{
@@ -65,18 +66,8 @@ func GetPerfMetrics(c *gin.Context) {
return
}
result.Groups = filterActiveGroups(result.Groups)
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": result,
})
}
func filterActiveGroups(groups []perfmetrics.GroupResult) []perfmetrics.GroupResult {
activeRatios := ratio_setting.GetGroupRatioCopy()
return lo.Filter(groups, func(g perfmetrics.GroupResult, _ int) bool {
_, ok := activeRatios[g.Group]
return ok || g.Group == "auto"
})
}
+190 -198
View File
@@ -25,12 +25,8 @@ import (
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting"
"github.com/QuantumNous/new-api/setting/operation_setting"
"github.com/bytedance/gopkg/util/gopool"
"github.com/samber/lo"
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
)
@@ -77,6 +73,13 @@ func publicUpstreamErrorContent(statusCode int) string {
}
func clientFacingOpenAIError(err *types.NewAPIError, upstreamFailure bool, requestId string) types.OpenAIError {
if err.GetErrorCode() == relaycommon.ResponseGuardErrorCode {
return types.OpenAIError{
Message: common.MessageWithRequestId(err.Error(), requestId),
Type: "api_error",
Code: relaycommon.ResponseGuardErrorCode,
}
}
if upstreamFailure {
return types.OpenAIError{
Message: common.MessageWithRequestId(publicUpstreamErrorMessage, requestId),
@@ -122,7 +125,11 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
defer func() {
if newAPIError != nil {
service.RecordRequestPolicyTermination(c, newAPIError)
logger.LogError(c, fmt.Sprintf("relay error: %s", common.LocalLogPreview(newAPIError.Error())))
if c.GetBool("response_guard_error_sent") {
return
}
switch relayFormat {
case types.RelayFormatOpenAIRealtime:
helper.WssError(c, ws, clientFacingOpenAIError(newAPIError, upstreamFailure, requestId))
@@ -156,59 +163,25 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
return
}
needSensitiveCheck := setting.ShouldCheckPromptSensitive()
needCountToken := constant.CountToken
// Avoid building huge CombineText (strings.Join) when token counting and sensitive check are both disabled.
var meta *types.TokenCountMeta
if needSensitiveCheck || needCountToken {
meta = request.GetTokenCountMeta()
} else {
meta = fastTokenCountMetaForPricing(request)
}
if needSensitiveCheck && meta != nil {
contains, words := service.CheckSensitiveText(meta.CombineText)
if contains {
logger.LogWarn(c, fmt.Sprintf("user sensitive words detected: %s", strings.Join(words, ", ")))
newAPIError = types.NewError(err, types.ErrorCodeSensitiveWordsDetected)
return
}
}
tokens, err := service.EstimateRequestToken(c, meta, relayInfo)
if err != nil {
newAPIError = types.NewError(err, types.ErrorCodeCountTokenFailed)
return
}
relayInfo.SetEstimatePromptTokens(tokens)
priceData, err := helper.ModelPriceHelper(c, relayInfo, tokens, meta)
if err != nil {
newAPIError = types.NewError(err, types.ErrorCodeModelPriceError, types.ErrOptionWithStatusCode(http.StatusBadRequest))
return
}
// common.SetContextKey(c, constant.ContextKeyTokenCountMeta, meta)
if priceData.FreeModel {
logger.LogInfo(c, fmt.Sprintf("模型 %s 免费,跳过预扣费", relayInfo.OriginModelName))
} else {
newAPIError = service.PreConsumeBilling(c, priceData.QuotaToPreConsume, relayInfo)
if newAPIError != nil {
return
}
}
defer func() {
// Only return quota if downstream failed and quota was actually pre-consumed
if newAPIError != nil {
newAPIError = service.NormalizeViolationFeeError(newAPIError)
if relayInfo.Billing != nil {
relayInfo.Billing.Refund(c)
}
service.ChargeViolationFeeIfNeeded(c, relayInfo, newAPIError)
recovered := recover()
resultErr := newAPIError
if recovered != nil {
resultErr = types.NewError(fmt.Errorf("relay panic: %v", recovered), types.ErrorCodeBadResponse)
}
if relayFormat != types.RelayFormatOpenAIRealtime {
perfmetrics.RecordRelayResult(c.Request.Context(), relayInfo, resultErr)
}
if recovered != nil {
panic(recovered)
}
}()
if newAPIError = relay.PrepareRequestBilling(c, relayInfo); newAPIError != nil {
return
}
defer func() {
newAPIError = relay.RefundFailedRequestBilling(c, relayInfo, newAPIError)
}()
retryParam := &service.RetryParam{
@@ -220,8 +193,30 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
}
relayInfo.RetryIndex = 0
relayInfo.LastError = nil
maxRetries := common.RetryTimes
ordinaryRetries := 0
if relayInfo.DurableAttemptBilling {
maxRetries = 2
timeout := constant.StreamingTimeout
if timeout <= 0 {
timeout = 300
}
relayInfo.GuardDeadline = time.Now().Add(time.Duration(timeout) * time.Second)
defer func() { service.EndFailedRelayAttempt(c, relayInfo) }()
}
for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() {
for ; retryParam.GetRetry() <= maxRetries; retryParam.IncreaseRetry() {
if relayInfo.DurableAttemptBilling && time.Now().After(relayInfo.GuardDeadline) {
newAPIError = types.NewErrorWithStatusCode(errors.New("response model confirmation budget exhausted"), relaycommon.ResponseGuardErrorCode, http.StatusGatewayTimeout, types.ErrOptionWithSkipRetry())
break
}
if relayInfo.ResponseGuard != nil && relayInfo.UpstreamRequests >= relaycommon.ResponseGuardMaxAttempts {
newAPIError = types.NewErrorWithStatusCode(errors.New("response model attempt budget exhausted"), relaycommon.ResponseGuardErrorCode, http.StatusBadGateway, types.ErrOptionWithSkipRetry())
break
}
relayInfo.StreamStatus = nil
relayInfo.PerformanceBusinessRejection = false
relayInfo.PerformanceOutputTokens = 0
relayInfo.RetryIndex = retryParam.GetRetry()
channel, channelErr := getChannel(c, relayInfo, retryParam)
if channelErr != nil {
@@ -229,11 +224,15 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
newAPIError = channelErr
break
}
addUsedChannel(c, channel.Id)
service.AppendUsedChannel(c, channel.Id)
if billingErr := service.PrepareTieredBillingForSelectedGroup(c, relayInfo); billingErr != nil {
newAPIError = billingErr
break
}
if billingErr := service.BeginRelayAttempt(c, relayInfo); billingErr != nil {
newAPIError = billingErr
break
}
bodyStorage, bodyErr := common.GetBodyStorage(c)
if bodyErr != nil {
@@ -259,19 +258,40 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
}
if newAPIError == nil {
service.MarkRequestPolicySuccess(c, relayInfo.StreamStatus)
relayInfo.LastError = nil
service.RecordGroupOutcome(currentRequestGroup(c, channel), true)
return
}
upstreamFailure = true
if isGroupAvailabilityFailure(newAPIError) {
service.RecordGroupOutcome(currentRequestGroup(c, channel), false)
}
newAPIError = service.NormalizeViolationFeeError(newAPIError)
relayInfo.LastError = newAPIError
service.EndFailedRelayAttempt(c, relayInfo)
remaining := common.RetryTimes - retryParam.GetRetry()
if relayInfo.DurableAttemptBilling {
remaining = common.RetryTimes - ordinaryRetries
}
decision := service.DecideRelayRetry(c, newAPIError, remaining)
if canSwitchRejectedAstra(c, relayInfo, newAPIError) {
decision = service.PolicyDecision{Action: "retry", Reason: "model_mismatch_switch", Source: "model_guard"}
service.ClearPoolAffinity(c)
} else if decision.Action == "retry" {
ordinaryRetries++
}
service.RecordPolicyFailure(c, channel.Id, newAPIError, decision)
processChannelError(c, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()), newAPIError, relayInfo)
if !shouldRetry(c, newAPIError, common.RetryTimes-retryParam.GetRetry()) {
if decision.Action != "retry" {
break
}
if service.RequestRoutingPolicy(c) == nil {
service.MarkRequestGroupFailed(c, currentRequestGroup(c, channel))
}
}
useChannel := c.GetStringSlice("use_channel")
@@ -279,11 +299,24 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
retryLogStr := fmt.Sprintf("重试:%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]"))
logger.LogInfo(c, retryLogStr)
}
if newAPIError != nil {
gopool.Go(func() {
perfmetrics.RecordRelaySample(relayInfo, false, 0)
})
}
func canSwitchRejectedAstra(c *gin.Context, info *relaycommon.RelayInfo, err *types.NewAPIError) bool {
if !info.AstraAutoSwitch || info.ResponseGuard == nil {
return false
}
if !info.DurableAttemptBilling || info.UpstreamRequests >= 3 || err.GetErrorCode() != relaycommon.ResponseGuardErrorCode || c.Writer.Written() || c.Request.Context().Err() != nil || service.RequestRoutingPolicy(c) == nil {
return false
}
if _, pinned, _ := service.GetChannelConstraints(c).ResolvedPin(); pinned {
return false
}
attempt := info.ResponseGuard.Last()
if attempt == nil || attempt.Reason != "model_mismatch" {
return false
}
request, ok := info.Request.(*dto.OpenAIResponsesRequest)
return ok && request.PreviousResponseID == "" && (len(request.Conversation) == 0 || string(request.Conversation) == "null")
}
// CountClaudeTokens implements Anthropic's token-counting utility endpoint.
@@ -319,47 +352,12 @@ func CountClaudeTokens(c *gin.Context) {
}
var upgrader = websocket.Upgrader{
Subprotocols: []string{"realtime"}, // WS 握手支持的协议,如果有使用 Sec-WebSocket-Protocol,则必须在此声明对应的 Protocol TODO add other protocol
Subprotocols: []string{"realtime", "responses"}, // WS 握手支持的协议,如果有使用 Sec-WebSocket-Protocol,则必须在此声明对应的 Protocol
CheckOrigin: func(r *http.Request) bool {
return true // 允许跨域
},
}
func addUsedChannel(c *gin.Context, channelId int) {
useChannel := c.GetStringSlice("use_channel")
useChannel = append(useChannel, fmt.Sprintf("%d", channelId))
c.Set("use_channel", useChannel)
}
func fastTokenCountMetaForPricing(request dto.Request) *types.TokenCountMeta {
if request == nil {
return &types.TokenCountMeta{}
}
meta := &types.TokenCountMeta{
TokenType: types.TokenTypeTokenizer,
}
switch r := request.(type) {
case *dto.GeneralOpenAIRequest:
maxCompletionTokens := lo.FromPtrOr(r.MaxCompletionTokens, uint(0))
maxTokens := lo.FromPtrOr(r.MaxTokens, uint(0))
if maxCompletionTokens > maxTokens {
meta.MaxTokens = int(maxCompletionTokens)
} else {
meta.MaxTokens = int(maxTokens)
}
case *dto.OpenAIResponsesRequest:
meta.MaxTokens = int(lo.FromPtrOr(r.MaxOutputTokens, uint(0)))
case *dto.ClaudeRequest:
meta.MaxTokens = int(lo.FromPtr(r.MaxTokens))
case *dto.ImageRequest:
// Pricing for image requests depends on ImagePriceRatio; safe to compute even when CountToken is disabled.
return r.GetTokenCountMeta()
default:
// Best-effort: leave CombineText empty to avoid large allocations.
}
return meta
}
func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service.RetryParam) (*model.Channel, *types.NewAPIError) {
if info.ChannelMeta == nil {
autoBan := c.GetBool("auto_ban")
@@ -367,12 +365,14 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service
if !autoBan {
autoBanInt = 0
}
return &model.Channel{
channel := &model.Channel{
Id: c.GetInt("channel_id"),
Type: c.GetInt("channel_type"),
Name: c.GetString("channel_name"),
AutoBan: &autoBanInt,
}, nil
}
service.RequestPolicy(c).BeginAttempt(channel, info.UsingGroup)
return channel, nil
}
channel, selectGroup, err := service.CacheGetRandomSatisfiedChannel(retryParam)
if err != nil {
@@ -384,6 +384,7 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service
info.PriceData.GroupRatioInfo = helper.HandleGroupRatio(c, info)
service.RequestPolicy(c).BeginAttempt(channel, selectGroup)
newAPIError := middleware.SetupContextForSelectedChannel(c, channel, info.OriginModelName)
if newAPIError != nil {
return nil, newAPIError
@@ -391,76 +392,45 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service
return channel, nil
}
func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) bool {
if openaiErr == nil {
// currentRequestGroup returns the group that handled the current attempt,
// falling back to the channel's own group when the selection did not record one.
func currentRequestGroup(c *gin.Context, channel *model.Channel) string {
if group := common.GetContextKeyString(c, constant.ContextKeyAutoGroup); group != "" {
return group
}
if channel != nil {
return channel.Group
}
return ""
}
// isGroupAvailabilityFailure reports whether an error says something about the
// selected group's health rather than about the client request itself.
func isGroupAvailabilityFailure(err *types.NewAPIError) bool {
if err == nil {
return false
}
if service.ShouldSkipRetryAfterChannelAffinityFailure(c) {
return false
}
if types.IsChannelError(openaiErr) {
if types.IsChannelError(err) {
return true
}
if types.IsSkipRetryError(openaiErr) {
return false
}
if retryTimes <= 0 {
return false
}
if service.GetChannelConstraints(c).SuppressesRetry() {
return false
}
code := openaiErr.StatusCode
if code >= 200 && code < 300 {
return false
}
if code < 100 || code > 599 {
return true
}
if operation_setting.IsAlwaysSkipRetryCode(openaiErr.GetErrorCode()) {
return false
}
return operation_setting.ShouldRetryByStatusCode(code)
return err.StatusCode >= 500
}
func shouldRetry(c *gin.Context, err *types.NewAPIError, retryTimes int) bool {
return service.ShouldRetryRelayError(c, err, retryTimes)
}
func processChannelError(c *gin.Context, channelError types.ChannelError, err *types.NewAPIError, relayInfo *relaycommon.RelayInfo) {
logger.LogError(c, fmt.Sprintf("channel error (channel #%d, status code: %d): %s", channelError.ChannelId, err.StatusCode, common.LocalLogPreview(err.Error())))
// 不要使用context获取渠道信息,异步处理时可能会出现渠道信息不一致的情况
// do not use context to get channel info, there may be inconsistent channel info when processing asynchronously
if service.ShouldDisableChannel(err) && channelError.AutoBan {
gopool.Go(func() {
service.DisableChannel(channelError, err.ErrorWithStatusCode())
})
}
if constant.ErrorLogEnabled && types.IsRecordErrorLog(err) {
// 保存错误日志到mysql中
userId := c.GetInt("id")
tokenName := c.GetString("token_name")
modelName := c.GetString("original_model")
tokenId := c.GetInt("token_id")
userGroup := c.GetString("group")
other := model.NewLogOther()
if c.Request != nil && c.Request.URL != nil {
other.SetPublic("request_path", c.Request.URL.Path)
}
other.SetPublic("error_type", err.GetErrorType())
other.SetPublic("error_code", err.GetErrorCode())
other.SetPublic("status_code", err.StatusCode)
service.AppendRelayLogAdminInfo(c, relayInfo, other)
service.AppendTaskPluginContextAuditInfo(c, other)
startTime := common.GetContextKeyTime(c, constant.ContextKeyRequestStartTime)
if startTime.IsZero() {
startTime = time.Now()
}
useTimeSeconds := int(time.Since(startTime).Seconds())
other.SetRoot("upstream_error", err.MaskSensitiveErrorWithStatusCode())
model.RecordErrorLog(c, userId, channelError.ChannelId, modelName, tokenName, publicUpstreamErrorContent(err.StatusCode), tokenId, useTimeSeconds, common.GetContextKeyBool(c, constant.ContextKeyIsStream), userGroup, other)
}
service.ProcessChannelError(c, channelError, err, relayInfo)
}
func RelayMidjourney(c *gin.Context) {
policy := service.RequestPolicy(c)
defer func() {
if policy.Attempts > 0 && !policy.Successful {
service.RecordRequestPolicyTermination(c, types.NewErrorWithStatusCode(errors.New("Midjourney submission failed"), types.ErrorCodeBadResponseStatusCode, http.StatusBadGateway, types.ErrOptionWithSkipRetry()))
}
}()
relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatMjProxy, nil, nil)
if err != nil {
@@ -488,6 +458,7 @@ func RelayMidjourney(c *gin.Context) {
//err = relayMidjourneySubmit(c, relayMode)
log.Println(mjErr)
if mjErr != nil {
policy.Successful = false
statusCode := http.StatusBadRequest
if mjErr.Code == 30 {
mjErr.Result = "当前分组负载已饱和,请稍后再试,或升级账户以提升服务质量。"
@@ -627,6 +598,7 @@ func executeTaskSubmissionWith(
relayInfo *relaycommon.RelayInfo,
submit taskSubmitAttempt,
) (*taskSubmissionOutcome, *taskdto.TaskError) {
policy := service.RequestPolicy(c)
diagnostics := newTaskPluginSubmitDiagnostics(c)
diagnostics.start(relayInfo)
var result *relay.TaskSubmitResult
@@ -664,6 +636,7 @@ func executeTaskSubmissionWith(
if lockedCh, ok := relayInfo.LockedChannel.(*model.Channel); ok && lockedCh != nil {
channel = lockedCh
policy.BeginAttempt(channel, relayInfo.UsingGroup)
if retryParam.GetRetry() > 0 {
if setupErr := middleware.SetupContextForSelectedChannel(c, channel, relayInfo.OriginModelName); setupErr != nil {
taskErr = service.TaskErrorWrapperLocal(setupErr.Err, "setup_locked_channel_failed", http.StatusInternalServerError)
@@ -675,13 +648,13 @@ func executeTaskSubmissionWith(
channel, channelErr = getChannel(c, relayInfo, retryParam)
if channelErr != nil {
logger.LogError(c, channelErr.Error())
taskErr = service.TaskErrorWrapperLocal(channelErr.Err, "get_channel_failed", http.StatusInternalServerError)
taskErr = service.TaskErrorWrapperLocal(channelErr.Err, "get_channel_failed", channelErr.StatusCode)
break
}
}
diagnostics.attempt(retryParam.GetRetry()+1, channel, relayInfo.LockedChannel != nil)
addUsedChannel(c, channel.Id)
service.AppendUsedChannel(c, channel.Id)
bodyStorage, bodyErr := common.GetBodyStorage(c)
if bodyErr != nil {
stage = "read_body"
@@ -706,15 +679,19 @@ func executeTaskSubmissionWith(
break
}
taskAPIError := taskSubmissionAPIError(taskErr)
relayInfo.LastError = taskAPIError
decision := decideTaskRetry(c, taskErr, common.RetryTimes-retryParam.GetRetry())
service.RecordPolicyFailure(c, channel.Id, taskAPIError, decision)
if !taskErr.LocalError {
processChannelError(c,
*types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey,
common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()),
types.NewOpenAIError(taskErr.Error, types.ErrorCodeBadResponseStatusCode, taskErr.StatusCode),
taskAPIError,
relayInfo)
}
willRetry := shouldRetryTaskRelay(c, channel.Id, taskErr, common.RetryTimes-retryParam.GetRetry())
willRetry := decision.Action == "retry"
diagnostics.attemptFailed(retryParam.GetRetry()+1, channel, taskErr, willRetry)
if !willRetry {
break
@@ -816,6 +793,11 @@ func executeTaskSubmissionWith(
diagnostics.failed("settle", "billing_error", taskErr, true)
return nil, taskErr
}
if task.Status != model.TaskStatusFailure {
service.MarkRequestPolicySuccess(c, nil)
} else {
policy.AddEvent(service.PolicyEvent{Decision: service.PolicyDecision{Action: "stop", Reason: "task_failed", Source: "upstream"}})
}
service.LogTaskConsumption(c, relayInfo, task)
diagnostics.complete(task, result.Quota)
@@ -876,6 +858,7 @@ func presentTaskSubmission(c *gin.Context, outcome *taskSubmissionOutcome) {
}
func respondTaskSubmissionError(c *gin.Context, taskErr *taskdto.TaskError) {
service.RecordRequestPolicyTermination(c, taskSubmissionAPIError(taskErr))
newTaskPluginSubmitDiagnostics(c).presentError(taskErr)
if middleware.RespondTaskPluginError(c, taskErr) {
return
@@ -891,44 +874,53 @@ func respondTaskError(c *gin.Context, taskErr *taskdto.TaskError) {
c.JSON(taskErr.StatusCode, taskErr)
}
func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *taskdto.TaskError, retryTimes int) bool {
if taskErr == nil {
return false
// taskSubmissionAPIError adapts a task error for the shared relay error paths.
// TaskError.Error is nil for many local rejections, so fall back to the message.
func taskSubmissionAPIError(taskErr *taskdto.TaskError) *types.NewAPIError {
err := taskErr.Error
if err == nil {
err = errors.New(taskErr.Message)
}
if service.ShouldSkipRetryAfterChannelAffinityFailure(c) {
return false
}
if retryTimes <= 0 {
return false
}
if service.GetChannelConstraints(c).SuppressesRetry() {
return false
}
if taskErr.StatusCode == http.StatusTooManyRequests {
return true
}
if taskErr.StatusCode == 307 {
return true
}
if taskErr.StatusCode/100 == 5 {
return types.NewOpenAIError(err, types.ErrorCodeBadResponseStatusCode, taskErr.StatusCode)
}
// decideTaskRetry is the single retry decision for task submissions. The
// reason is recorded in the request policy decision events of the log details.
func decideTaskRetry(c *gin.Context, taskErr *taskdto.TaskError, retryTimes int) service.PolicyDecision {
stop := service.PolicyDecision{Action: "stop", Source: "system"}
retry := service.PolicyDecision{Action: "retry", Reason: "retry_status_matched", Source: "system"}
switch {
case taskErr == nil:
stop.Reason = "request_completed"
case taskErr.NoRetry:
stop.Reason = "task_accepted"
case service.ShouldSkipRetryAfterChannelAffinityFailure(c):
stop.Reason, stop.Source = "strict_session", "session_rule"
if source := service.RequestPolicy(c).SessionModeSource; source != "" {
stop.Source = source
}
case retryTimes <= 0:
stop.Reason, stop.Source = "attempt_budget_exhausted", "global"
case service.GetChannelConstraints(c).SuppressesRetry():
stop.Reason, stop.Source = "pinned_channel", "channel_constraint"
case taskErr.StatusCode == http.StatusTooManyRequests, taskErr.StatusCode == 307:
return retry
case taskErr.StatusCode/100 == 5:
// 超时不重试
if operation_setting.IsAlwaysSkipRetryStatusCode(taskErr.StatusCode) {
return false
stop.Reason = "system_retry_exclusion"
break
}
return true
}
if taskErr.StatusCode == http.StatusBadRequest {
return false
}
if taskErr.StatusCode == 408 {
return retry
case taskErr.StatusCode == http.StatusBadRequest, taskErr.StatusCode == 408:
// azure处理超时不重试
return false
stop.Reason = "status_not_retryable"
case taskErr.LocalError:
stop.Reason = "local_rejection"
case taskErr.StatusCode/100 == 2:
stop.Reason = "system_retry_exclusion"
default:
return retry
}
if taskErr.LocalError {
return false
}
if taskErr.StatusCode/100 == 2 {
return false
}
return true
return stop
}
+69
View File
@@ -0,0 +1,69 @@
package controller
import (
"encoding/json"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/gin-gonic/gin"
)
func relayAttemptGuardResult(a model.RelayAttempt) string {
if a.GuardResult == "model_mismatch" {
return "mismatch_rejected"
}
if a.GuardResult == "unchecked" && a.ReturnedModel != "" && a.ReturnedModel == a.RequestedModel {
return "accepted"
}
if a.GuardResult == "accepted" && a.Number > 1 {
return "switched_success"
}
return a.GuardResult
}
// Attempt history is owner-scoped; channel identifiers are administrator-only.
func GetRelayAttempts(c *gin.Context) {
requestID := c.Param("request_id")
if requestID == "" || len(requestID) > 64 {
c.AbortWithStatus(400)
return
}
query := model.DB.Where("request_id = ?", requestID)
admin := c.GetInt("role") >= common.RoleAdminUser
if !admin {
query = query.Where("user_id = ?", c.GetInt("id"))
}
var attempts []model.RelayAttempt
if err := query.Order("number asc, id asc").Limit(100).Find(&attempts).Error; err != nil {
common.ApiError(c, err)
return
}
rows := make([]gin.H, 0, len(attempts))
total := 0
pending := false
for _, a := range attempts {
var usage dto.Usage
_ = common.UnmarshalJsonStr(a.UsageJSON, &usage)
row := gin.H{"number": a.Number, "group": a.GroupName, "requested_model": a.RequestedModel, "returned_model": a.ReturnedModel, "guard_result": relayAttemptGuardResult(a), "usage_source": a.UsageSource, "state": a.State, "quota": a.Quota, "billing_source": a.Funding, "usage": usage}
row["charged"] = a.State == "settled" && a.Quota > 0
row["charge_explanation"] = "Each consumed attempt is charged independently, including rejected results."
var log model.Log
var other map[string]json.RawMessage
price := map[string]json.RawMessage{}
if common.UnmarshalJsonStr(a.LogJSON, &log) == nil && common.UnmarshalJsonStr(log.Other, &other) == nil {
for _, key := range []string{"group_ratio", "model_ratio", "completion_ratio", "cache_ratio", "model_price", "billing_mode", "expr_b64", "matched_tier", "billing_tokens", "billing_unit", "fixed_price"} {
if value, ok := other[key]; ok {
price[key] = value
}
}
}
row["price"] = price
if admin {
row["channel_id"] = a.ChannelID
}
rows = append(rows, row)
total += a.Quota
pending = pending || a.State != "settled"
}
common.ApiSuccess(c, gin.H{"attempts": rows, "total_quota": total, "pending": pending})
}
+66
View File
@@ -0,0 +1,66 @@
package controller
import (
"context"
"errors"
hostdto "github.com/QuantumNous/new-api/dto"
"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/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
"net/http/httptest"
"testing"
)
func TestRelayAttemptSwitchSafety(t *testing.T) {
for _, tc := range []struct {
name, reason, previous, conversation string
count int
pin, cancel, written, disabled, want bool
}{
{name: "explicit mismatch", reason: "model_mismatch", count: 1, want: true},
{name: "token disabled", reason: "model_mismatch", count: 1, disabled: true},
{name: "third channel allowed", reason: "model_mismatch", count: 2, want: true},
{name: "third channel final", reason: "model_mismatch", count: 3},
{name: "missing model", reason: "missing_model", count: 1},
{name: "confirmation timeout", reason: "confirmation_timeout", count: 1},
{name: "pinned", reason: "model_mismatch", count: 1, pin: true},
{name: "cancelled", reason: "model_mismatch", count: 1, cancel: true},
{name: "already delivered", reason: "model_mismatch", count: 1, written: true},
{name: "upstream response id", reason: "model_mismatch", count: 1, previous: "resp_private"},
{name: "upstream conversation", reason: "model_mismatch", count: 1, conversation: `"conv_private"`},
} {
t.Run(tc.name, func(t *testing.T) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
c.Request = httptest.NewRequest("POST", "/v1/responses", nil).WithContext(ctx)
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "a"}}})
if tc.pin {
service.GetChannelConstraints(c).AddPin(hostdto.ChannelPin{ChannelId: 1, Source: hostdto.PinSourceToken, Rank: hostdto.PinRankToken, RetryMode: hostdto.PinRetrySingleAttempt})
}
if tc.cancel {
cancel()
}
if tc.written {
_, _ = c.Writer.WriteString("already sent")
}
info := &relaycommon.RelayInfo{DurableAttemptBilling: true, UpstreamRequests: tc.count, Request: &dto.OpenAIResponsesRequest{PreviousResponseID: tc.previous, Conversation: []byte(tc.conversation)}, ResponseGuard: &relaycommon.ResponseModelGuard{Attempts: []*relaycommon.ResponseGuardAttempt{{Reason: tc.reason}}}}
err := types.NewError(errors.New("guard rejected"), relaycommon.ResponseGuardErrorCode)
info.AstraAutoSwitch = !tc.disabled
require.Equal(t, tc.want, canSwitchRejectedAstra(c, info, err))
})
}
}
func TestRelayAttemptGuardResult(t *testing.T) {
require.Equal(t, "mismatch_rejected", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "model_mismatch", Number: 1}))
require.Equal(t, "switched_success", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "accepted", Number: 2}))
require.Equal(t, "accepted", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "accepted", Number: 1}))
require.Equal(t, "accepted", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "unchecked", Number: 1, RequestedModel: "gpt-6-astra", ReturnedModel: "gpt-6-astra"}))
require.Equal(t, "unchecked", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "unchecked", Number: 1, RequestedModel: "gpt-6-astra"}))
require.Equal(t, "transport_error", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "transport_error", Number: 1}))
}
+18
View File
@@ -2,6 +2,7 @@ package controller
import (
"errors"
"fmt"
"net/http"
"strings"
"testing"
@@ -43,3 +44,20 @@ func TestClientFacingErrorKeepsLocalValidationMessage(t *testing.T) {
t.Fatalf("local validation message was hidden: %q", got.Message)
}
}
func TestClientFacingResponseGuardErrorNamesDeclaredModel(t *testing.T) {
err := types.NewErrorWithStatusCode(
errors.New("Requested model gpt-6-astra, but the upstream response declared gpt-5.6-luna."),
"response_model_rejected",
http.StatusConflict,
)
got := clientFacingOpenAIError(err, true, "local-request-id")
if fmt.Sprint(got.Code) != "response_model_rejected" {
t.Fatalf("unexpected public code: %v", got.Code)
}
for _, expected := range []string{"gpt-6-astra", "gpt-5.6-luna", "local-request-id"} {
if !strings.Contains(got.Message, expected) {
t.Fatalf("public mismatch message %q does not contain %q", got.Message, expected)
}
}
}
+193 -5
View File
@@ -3,8 +3,10 @@ package controller
import (
"context"
"errors"
"fmt"
"net/http"
"net/http/httptest"
"os"
"strings"
"testing"
"time"
@@ -22,14 +24,18 @@ import (
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/driver/mysql"
"gorm.io/driver/postgres"
"gorm.io/gorm"
"gorm.io/gorm/schema"
)
type taskSubmissionTestBilling struct {
events *[]string
settleErr error
onSettle func()
refunds int
events *[]string
settleErr error
reserveErr error
onSettle func()
refunds int
}
func (b *taskSubmissionTestBilling) Settle(int) error {
@@ -49,7 +55,7 @@ func (b *taskSubmissionTestBilling) NeedsRefund() bool { return b.refunds
func (b *taskSubmissionTestBilling) GetPreConsumedQuota() int { return 0 }
func (b *taskSubmissionTestBilling) Reserve(int) error {
*b.events = append(*b.events, "reserve")
return nil
return b.reserveErr
}
func TestPresentTaskSubmissionUsesNativePresenterAfterPersistence(t *testing.T) {
@@ -387,3 +393,185 @@ func taskSubmissionRelayInfo(billing relaycommon.BillingSettler) *relaycommon.Re
ChannelMeta: &relaycommon.ChannelMeta{ChannelId: 1, ChannelType: constant.ChannelTypeTaskPlugin},
}
}
// Uses the real submit adaptor, expression evaluator, BillingSession, task row
// and consume log. Set TEST_TASK_DB_DIALECT plus TEST_MYSQL_DSN or
// TEST_POSTGRES_DSN to exercise the same contract on an external test database.
// Unique table prefixes keep the fixture isolated from all existing tables.
func TestImmediateTaskSettlementDatabase(t *testing.T) {
dialect := common.DatabaseType(os.Getenv("TEST_TASK_DB_DIALECT"))
var driver gorm.Dialector
switch dialect {
case "", common.DatabaseTypeSQLite:
dialect = common.DatabaseTypeSQLite
driver = sqlite.Open(":memory:")
case common.DatabaseTypeMySQL:
require.NotEmpty(t, os.Getenv("TEST_MYSQL_DSN"))
driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
case common.DatabaseTypePostgreSQL:
require.NotEmpty(t, os.Getenv("TEST_POSTGRES_DSN"))
driver = postgres.New(postgres.Config{DSN: os.Getenv("TEST_POSTGRES_DSN"), PreferSimpleProtocol: true})
default:
t.Fatalf("unsupported test dialect %q", dialect)
}
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("tsubmit_%d_", time.Now().UnixNano())}})
require.NoError(t, err)
sqlDB, err := db.DB()
require.NoError(t, err)
sqlDB.SetMaxOpenConns(1)
t.Cleanup(func() { _ = sqlDB.Close() })
models := []any{&model.User{}, &model.Channel{}, &model.Task{}, &model.Log{}}
require.NoError(t, db.AutoMigrate(models...))
t.Cleanup(func() { require.NoError(t, db.Migrator().DropTable(models...)) })
var version string
if dialect == common.DatabaseTypeSQLite {
require.NoError(t, db.Raw("SELECT sqlite_version()").Scan(&version).Error)
} else {
require.NoError(t, db.Raw("SELECT version()").Scan(&version).Error)
}
t.Logf("database: %s %s", dialect, version)
oldDB, oldLogDB := model.DB, model.LOG_DB
oldMain, oldLog := common.MainDatabaseType(), common.LogDatabaseType()
oldRedis, oldMemory, oldBatch, oldConsume, oldExport := common.RedisEnabled, common.MemoryCacheEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled, common.DataExportEnabled
model.DB, model.LOG_DB = db, db
common.SetDatabaseTypes(dialect, dialect)
common.RedisEnabled, common.MemoryCacheEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled, common.DataExportEnabled = false, false, false, true, false
t.Cleanup(func() {
model.DB, model.LOG_DB = oldDB, oldLogDB
common.SetDatabaseTypes(oldMain, oldLog)
common.RedisEnabled, common.MemoryCacheEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled, common.DataExportEnabled = oldRedis, oldMemory, oldBatch, oldConsume, oldExport
})
const expression = `u("units") == 7 ? tier("missing", u("missing") * 1.0) : u("units") == 8 ? tier("invalid", -1.0) : u("units") > 3 ? tier("bulk", u("units") * 0.01) : tier("small", u("units") * 0.01)`
withTieredBillingConfig(t, map[string]string{"document-model": "tiered_expr"}, map[string]string{"document-model": expression})
const source = `
export const meta={apiVersion:1,key:"generic-settlement",name:"Generic settlement",version:"1.0.0",author:{name:"Test"},models:["document-model"],fetchMode:"per_task",usageSchema:{units:{type:"number",unit:"count"}}};
export function buildSubmitRequest(ctx){return {url:ctx.baseUrl+"/compile",body:ctx.requestBody};}
export function parseSubmitResponse(ctx,resp){return {taskId:"vendor-job",taskData:resp.body,immediate:{status:resp.body.status,reason:"provider rejected job"}};}
export function extractUsage(){return {units:4};}
export function extractUsageOnComplete(ctx,result,body){return body.usage;}
export function parseTaskResult(){throw new Error("completed submissions must not poll");}
export function buildQueryRequest(){throw new Error("completed submissions must not poll");}
`
plugin, err := pluginruntime.CompilePlugin(source, pluginruntime.Options{})
require.NoError(t, err)
for index, tc := range []struct {
name, status string
actual any
count float64
}{
{"partial", "SUCCESS", 2, 2}, {"zero", "SUCCESS", 0, 0}, {"larger", "SUCCESS", 6, 6},
{"invalid usage", "SUCCESS", -1, 4}, {"expression failure", "SUCCESS", 7, 4}, {"negative result", "SUCCESS", 8, 4}, {"missing usage", "SUCCESS", nil, 4}, {"failed", "FAILURE", 9, 0},
} {
t.Run(tc.name, func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
body := map[string]any{"status": tc.status}
if tc.actual != nil {
body["usage"] = map[string]any{"units": tc.actual}
}
encoded, err := common.Marshal(body)
if err != nil {
panic(err)
}
_, _ = w.Write(encoded)
}))
defer server.Close()
initial := int(20 * common.QuotaPerUnit)
user := model.User{Username: fmt.Sprintf("task_user_%d", index), AffCode: fmt.Sprintf("task_aff_%d", index), Quota: initial}
require.NoError(t, db.Create(&user).Error)
ch := model.Channel{Name: "test provider", Type: constant.ChannelTypeTaskPlugin}
require.NoError(t, db.Create(&ch).Error)
c := taskSubmissionTestContext()
c.Set("group", "default")
c.Set("username", user.Username)
c.Set("task_request", map[string]any{"model": "document-model"})
c.Set(pluginruntime.ContextKeyPinnedPlugin, pluginruntime.PinnedPlugin{Plugin: plugin})
common.SetContextKey(c, constant.ContextKeyOriginalModel, "document-model")
common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, server.URL)
common.SetContextKey(c, constant.ContextKeyChannelId, ch.Id)
common.SetContextKey(c, constant.ContextKeyChannelType, ch.Type)
info := taskSubmissionRelayInfo(nil)
info.UserId = user.Id
info.OriginModelName = "document-model"
info.UserGroup = "default"
info.UsingGroup = "default"
info.IsPlayground = true
info.UserSetting.BillingPreference = "wallet_only"
info.PublicTaskID = model.GenerateTaskID()
info.LockedChannel = &ch
outcome, taskErr := executeTaskSubmission(c, info)
require.Nil(t, taskErr)
require.NotNil(t, outcome)
want := common.QuotaRound(tc.count * 0.01 * common.QuotaPerUnit)
assert.Equal(t, want, outcome.Result.Quota)
assert.Equal(t, want, info.PriceData.Quota)
var stored model.Task
require.NoError(t, db.Where("task_id = ?", info.PublicTaskID).First(&stored).Error)
assert.Equal(t, want, stored.Quota)
assert.Equal(t, model.TaskStatus(tc.status), stored.Status)
assert.Positive(t, stored.FinishTime)
assert.Equal(t, float64(4), info.TieredBillingSnapshot.EstimatedQuotaBeforeGroup/(0.01*common.QuotaPerUnit))
if tc.status == "SUCCESS" {
assert.Equal(t, tc.count, stored.PrivateData.BillingContext.TieredSnapshot.UsageFacts["units"])
}
var updated model.User
require.NoError(t, db.First(&updated, user.Id).Error)
assert.Equal(t, initial-want, updated.Quota)
assert.Equal(t, want, updated.UsedQuota)
var logs []model.Log
require.NoError(t, db.Where("user_id = ?", user.Id).Find(&logs).Error)
require.Len(t, logs, 1)
assert.Equal(t, want, logs[0].Quota)
var other map[string]any
require.NoError(t, common.UnmarshalJsonStr(logs[0].Other, &other))
if tc.status == "SUCCESS" {
assert.Equal(t, tc.count, other["usage_facts"].(map[string]any)["units"])
}
assert.False(t, c.Writer.Written(), "presentation must follow persistence and settlement")
require.NoError(t, info.Billing.Settle(want))
info.Billing.Refund(c)
require.NoError(t, db.First(&updated, user.Id).Error)
assert.Equal(t, initial-want, updated.Quota, "terminal settlement is idempotent")
})
}
}
func TestAcceptedSubmitStreamNeverRetries(t *testing.T) {
c := taskSubmissionTestContext()
assert.Equal(t, service.PolicyDecision{Action: "stop", Reason: "task_accepted", Source: "system"}, decideTaskRetry(c, &dto.TaskError{StatusCode: 502, LocalError: true, NoRetry: true}, 3))
}
// Local task rejections carry a message but no cause; the response and the
// decision record must still be produced.
func TestRespondTaskSubmissionErrorWithoutCause(t *testing.T) {
previousErrorLog := constant.ErrorLogEnabled
constant.ErrorLogEnabled = false
t.Cleanup(func() { constant.ErrorLogEnabled = previousErrorLog })
c := taskSubmissionTestContext()
taskErr := &dto.TaskError{Code: "get_channel_failed", Message: "no channel", StatusCode: http.StatusServiceUnavailable, LocalError: true}
require.NotPanics(t, func() { respondTaskSubmissionError(c, taskErr) })
assert.Equal(t, http.StatusServiceUnavailable, c.Writer.Status())
events := service.RequestPolicy(c).Events()
require.Len(t, events, 1)
assert.Equal(t, service.PolicyDecision{Action: "stop", Reason: "request_failed", Source: "system"}, events[0].Decision)
assert.Equal(t, http.StatusServiceUnavailable, events[0].Status)
assert.Equal(t, service.PolicyDecision{Action: "stop", Reason: "local_rejection", Source: "system"}, decideTaskRetry(c, &dto.TaskError{StatusCode: http.StatusForbidden, LocalError: true, Message: "billing"}, 2), "local errors stop once the status rules do not force a retry")
}
func TestExecuteTaskSubmissionRefundsWhenFinalReserveFails(t *testing.T) {
events := []string{}
setupTaskSubmissionDatabase(t, true, &events)
billing := &taskSubmissionTestBilling{events: &events, reserveErr: errors.New("insufficient funds")}
c := taskSubmissionTestContext()
info := taskSubmissionRelayInfo(billing)
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
return &relay.TaskSubmitResult{Platform: "plugin", Quota: 600, Immediate: &relaycommon.TaskInfo{Status: "SUCCESS"}}, nil
})
require.Nil(t, outcome)
require.NotNil(t, taskErr)
assert.Equal(t, http.StatusForbidden, taskErr.StatusCode)
assert.Equal(t, []string{"reserve", "refund"}, events)
assert.Equal(t, 1, billing.refunds)
assert.False(t, c.Writer.Written())
}
+27
View File
@@ -0,0 +1,27 @@
package controller
import (
"net/http"
"github.com/QuantumNous/new-api/model"
"github.com/gin-gonic/gin"
)
func GetRequestPolicy(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"success": true, "data": gin.H{"options": model.CurrentRequestPolicy().Options}})
}
func UpdateRequestPolicy(c *gin.Context) {
var request struct {
Options map[string]string `json:"options"`
}
if err := c.ShouldBindJSON(&request); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()})
return
}
if err := model.UpdateRequestPolicyOptions(request.Options); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()})
return
}
GetRequestPolicy(c)
}
+129
View File
@@ -0,0 +1,129 @@
package controller
import (
"bytes"
"context"
"errors"
"fmt"
"net"
"net/http"
"sync"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/middleware"
"github.com/QuantumNous/new-api/relay"
"github.com/QuantumNous/new-api/relay/helper"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/gin-gonic/gin"
)
type responsesWSRequestContextKey struct{}
type responsesWSRequestState struct {
requestID string
handle func(*gin.Context) *types.NewAPIError
apiError *types.NewAPIError
}
// Each response.create runs the ordinary request middleware to completion,
// without routing another HTTP request or retaining a pooled Gin context.
var responsesWSRequestEngine = sync.OnceValue(func() *gin.Engine {
engine := gin.New()
engine.ForwardedByClientIP = false
_ = engine.SetTrustedProxies(nil)
engine.POST("/v1/responses", func(c *gin.Context) {
state := c.Request.Context().Value(responsesWSRequestContextKey{}).(*responsesWSRequestState)
c.Set(common.RequestIdKey, state.requestID)
common.SetContextKey(c, constant.ContextKeyRequestStartTime, time.Now())
c.Next()
}, middleware.BodyStorageCleanup(), middleware.TokenAuth(), middleware.ModelRequestRateLimit(), func(c *gin.Context) {
state := c.Request.Context().Value(responsesWSRequestContextKey{}).(*responsesWSRequestState)
state.apiError = state.handle(c)
if state.apiError != nil {
status := state.apiError.StatusCode
if status < http.StatusBadRequest {
status = http.StatusInternalServerError
}
c.Status(status)
}
})
return engine
})
type responsesWSResponseWriter struct {
header http.Header
status int
body bytes.Buffer
}
func (w *responsesWSResponseWriter) Header() http.Header {
return w.header
}
func (w *responsesWSResponseWriter) WriteHeader(status int) {
if w.status == 0 {
w.status = status
}
}
func (w *responsesWSResponseWriter) Write(data []byte) (int, error) {
w.WriteHeader(http.StatusOK)
return w.body.Write(data)
}
func newResponsesWSRequestRunner(c *gin.Context) relay.ResponsesWSRequestRunner {
// Capture credentials before channel selection or header overrides. Resolve
// the peer once with the public router's trusted-proxy configuration.
headers := c.Request.Header.Clone()
for _, name := range []string{"Connection", "Upgrade", "Sec-WebSocket-Key", "Sec-WebSocket-Version", "Sec-WebSocket-Extensions", "Sec-WebSocket-Protocol", "Content-Length", "Content-Encoding"} {
headers.Del(name)
}
remoteAddr := net.JoinHostPort(c.ClientIP(), "0")
return func(request *http.Request, requestID string, handle func(*gin.Context) *types.NewAPIError) *types.NewAPIError {
state := &responsesWSRequestState{requestID: requestID, handle: handle}
ctx := context.WithValue(request.Context(), responsesWSRequestContextKey{}, state)
ctx = context.WithValue(ctx, common.RequestIdKey, requestID)
request = request.Clone(ctx)
request.Method = http.MethodPost
request.URL.Path = "/v1/responses"
request.URL.RawPath = ""
request.Header = headers.Clone()
request.Header.Set("Content-Type", "application/json")
request.RemoteAddr = remoteAddr
response := &responsesWSResponseWriter{header: make(http.Header)}
responsesWSRequestEngine().ServeHTTP(response, request)
if state.apiError != nil {
return state.apiError
}
if response.status < http.StatusBadRequest {
return nil
}
var body struct {
Error *types.OpenAIError `json:"error"`
}
if common.Unmarshal(response.body.Bytes(), &body) == nil && body.Error != nil {
return types.WithOpenAIError(*body.Error, response.status, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
}
// The existing in-memory rate limiter returns a bare 429 response.
return types.NewErrorWithStatusCode(errors.New(http.StatusText(response.status)), types.ErrorCodeInvalidRequest, response.status, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
}
}
func ResponsesWebSocket(c *gin.Context) {
requestID := c.GetString(common.RequestIdKey)
runner := newResponsesWSRequestRunner(c)
ws, err := upgrader.Upgrade(c.Writer, c.Request, nil)
if err != nil {
return
}
defer ws.Close()
if apiError := relay.ResponsesWebSocketHelper(c, ws, runner); apiError != nil {
logger.LogError(c, fmt.Sprintf("responses websocket relay error: %s", common.LocalLogPreview(apiError.Error())))
apiError.SetMessage(common.MessageWithRequestId(apiError.Error(), requestID))
helper.WssError(c, ws, apiError.ToOpenAIError())
}
}
File diff suppressed because it is too large. Load diff
+19
View File
@@ -3,12 +3,18 @@ package controller
import (
"errors"
"net/http"
"slices"
"strings"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/middleware"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/oauth"
"github.com/QuantumNous/new-api/service"
passkeysvc "github.com/QuantumNous/new-api/service/passkey"
"github.com/QuantumNous/new-api/setting/system_setting"
"github.com/gin-gonic/gin"
"github.com/go-webauthn/webauthn/protocol"
)
@@ -34,6 +40,10 @@ func writeSecurityOperationError(c *gin.Context, err error) {
var code, message string
var protocolError *protocol.Error
switch {
case errors.Is(err, passkeysvc.ErrRPIDUnavailable):
code, message = "PASSKEY_RP_ID_UNAVAILABLE", i18n.T(c, i18n.MsgPasskeyRPIDUnavailable)
case errors.Is(err, system_setting.ErrPasskeyRPIDInvalid):
code, message = "PASSKEY_RP_ID_INVALID", i18n.T(c, i18n.MsgPasskeyRPIDInvalid)
case errors.Is(err, service.ErrAccountEmailInvalid), errors.Is(err, service.ErrAccountEmailRestricted):
code, message = "EMAIL_ADDRESS_REJECTED", err.Error()
case errors.Is(err, model.ErrEmailAlreadyTaken):
@@ -112,6 +122,15 @@ func writeSecurityOperationError(c *gin.Context, err error) {
return
}
c.Set("security_error_code", code)
if strings.Contains(c.Request.URL.Path, "/passkey/") {
reason := "verification_failed"
if errors.As(err, &protocolError) && slices.Contains([]string{"invalid_request", "challenge_mismatch", "parse_error", "auth_data", "verification_error", "invalid_signature", "invalid_key_type", "unsupported_key_algorithm"}, protocolError.Type) {
reason = protocolError.Type
}
// Protocol details can contain challenges and client-controlled data.
// Only fixed categories and the server-selected public RP ID are logged.
logger.LogWarn(c.Request.Context(), "passkey verification rejected: code=%s reason=%s rp_id=%q", code, reason, c.GetString("passkey_rp_id"))
}
c.JSON(status, gin.H{"success": false, "code": code, "message": message})
}
+1 -1
View File
@@ -64,7 +64,7 @@ func setupSecurityEnrollmentTest(t *testing.T) (*model.User, service.AuthIdentit
var version string
require.NoError(t, db.Raw(versionQuery).Scan(&version).Error)
t.Logf("database: %s %s", dialect, version)
require.NoError(t, db.AutoMigrate(&model.User{}, &model.UserSession{}, &model.TwoFA{}, &model.TwoFABackupCode{}, &model.PasskeyCredential{}, &model.AuthFlow{}, &model.UserOAuthBinding{}))
require.NoError(t, db.AutoMigrate(&model.User{}, &model.UserSession{}, &model.TwoFA{}, &model.TwoFABackupCode{}, &model.PasskeyCredential{}, &model.AuthFlow{}, &model.UserOAuthBinding{}, &model.Option{}))
require.NoError(t, logDB.AutoMigrate(&model.AuditLog{}))
model.DB, model.LOG_DB = db, logDB
dbType := common.DatabaseTypeSQLite
+84
View File
@@ -0,0 +1,84 @@
package controller
import (
"net/http"
"strings"
"unicode/utf8"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/gin-gonic/gin"
)
func GetStatusPageAdmin(c *gin.Context) {
results, err := currentGroupMonitors(c.Request.Context(), true)
if err != nil {
common.ApiError(c, err)
return
}
c.JSON(http.StatusOK, gin.H{"success": true, "data": results})
}
func UpdateStatusPageGroup(c *gin.Context) {
var input struct {
Group string `json:"group"`
Alias string `json:"alias"`
Hidden *bool `json:"hidden"`
}
if c.ShouldBindJSON(&input) != nil || strings.TrimSpace(input.Group) == "" || utf8.RuneCountInString(input.Group) > 64 {
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "Invalid group"})
return
}
reset := strings.HasSuffix(c.FullPath(), "/reset")
input.Alias = strings.TrimSpace(input.Alias)
if !reset && (input.Hidden == nil || utf8.RuneCountInString(input.Alias) > 80 || strings.ContainsAny(input.Alias, "\r\n\t")) {
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "Invalid display settings"})
return
}
channels, err := model.GetStatusChannels(c.Request.Context())
if err != nil {
common.ApiError(c, err)
return
}
found := false
for _, channel := range channels {
for group := range strings.SplitSeq(channel.Group, ",") {
if strings.TrimSpace(group) == input.Group {
found = true
}
}
}
if !found {
c.JSON(http.StatusNotFound, gin.H{"success": false, "message": "Group not found"})
return
}
var display *model.StatusPageGroup
if !reset {
display = &model.StatusPageGroup{Alias: input.Alias, Hidden: *input.Hidden}
}
if err := model.SaveStatusPageGroup(c.Request.Context(), input.Group, display); err != nil {
common.ApiError(c, err)
return
}
c.JSON(http.StatusOK, gin.H{"success": true})
}
// Main-site group status is independent of the public uptime display settings.
func GetSiteGroupStatus(c *gin.Context) {
results, err := currentGroupMonitors(c.Request.Context(), true)
if err != nil {
common.ApiError(c, err)
return
}
for i := range results {
monitors := make([]Monitor, 0, len(results[i].Monitors))
for _, monitor := range results[i].Monitors {
monitor.Name = monitor.Group
monitor.Alias = ""
monitor.Hidden = false
monitors = append(monitors, monitor)
}
results[i].Monitors = monitors
}
c.JSON(http.StatusOK, gin.H{"success": true, "data": results})
}
+29 -2
View File
@@ -66,9 +66,17 @@ func GetCurrentSystemTask(c *gin.Context) {
}
func ListSystemTasks(c *gin.Context) {
limit, _ := strconv.Atoi(c.Query("limit"))
var query struct {
model.SystemTaskFilter
Offset int `form:"offset" binding:"gte=0"`
Limit int `form:"limit"`
}
if err := c.ShouldBindQuery(&query); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "invalid system task filters"})
return
}
tasks, err := model.ListSystemTasks(limit)
tasks, total, err := model.ListSystemTasks(query.SystemTaskFilter, query.Offset, query.Limit)
if err != nil {
common.ApiError(c, err)
return
@@ -83,6 +91,25 @@ func ListSystemTasks(c *gin.Context) {
"success": true,
"message": "",
"data": responses,
"total": total,
})
}
func DeleteSystemTaskHistory(c *gin.Context) {
var filter model.SystemTaskFilter
if err := c.ShouldBindQuery(&filter); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "invalid system task filters"})
return
}
deleted, err := model.DeleteSystemTaskHistory(filter)
if err != nil {
common.ApiError(c, err)
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"message": "",
"data": gin.H{"deleted_count": deleted},
})
}
+72
View File
@@ -3,6 +3,7 @@ package controller
import (
"context"
"fmt"
"sync"
"time"
"github.com/QuantumNous/new-api/common"
@@ -18,6 +19,9 @@ import (
// instances and each run is recorded as one task row. Call this before
// service.StartSystemTaskRunner.
func RegisterScheduledSystemTasks() {
service.RegisterSystemTaskHandler(detectionHandler{})
service.RegisterSystemTaskHandler(astraSamplingHandler{})
service.RegisterSystemTaskHandler(channelAvailabilityHandler{})
service.RegisterSystemTaskHandler(channelTestHandler{})
service.RegisterSystemTaskHandler(modelUpdateHandler{})
service.RegisterSystemTaskHandler(midjourneyPollHandler{})
@@ -29,6 +33,74 @@ func RegisterScheduledSystemTasks() {
// into the system task runner.
type channelTestHandler struct{}
type channelAvailabilityHandler struct{}
func (channelAvailabilityHandler) Type() string { return model.SystemTaskTypeChannelAvailability }
func (channelAvailabilityHandler) Enabled() bool { return service.AstraAvailabilityEnabled() }
func (channelAvailabilityHandler) Interval() time.Duration {
return time.Duration(service.AstraAvailabilityIntervalMinutes()) * time.Minute
}
func (channelAvailabilityHandler) ScheduleFromCreation() bool { return true }
func (channelAvailabilityHandler) NewPayload() any { return nil }
func (channelAvailabilityHandler) Run(ctx context.Context, task *model.SystemTask, runnerID string) {
testUserID, err := resolveChannelTestUserID(nil)
if err != nil {
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, err)
return
}
channels, err := model.GetAllChannels(0, 0, true, false)
if err != nil {
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, err)
return
}
selected := make([]*model.Channel, 0, len(channels))
for _, channel := range channels {
if channel.Status == common.ChannelStatusEnabled {
selected = append(selected, channel)
}
}
settings := operation_setting.GetMonitorSetting()
state := service.ChannelAvailabilityState{Rounds: []service.ChannelAvailabilityRound{}}
var mu sync.Mutex
var persistErr error
ctx, cancel := context.WithCancel(ctx)
defer cancel()
disableThreshold := int64(common.ChannelDisableThreshold * 1000)
if disableThreshold == 0 {
disableThreshold = 10000000
}
summary := runChannelTestWorkers(ctx, selected, settings.ChannelTestConcurrency, func(ctx context.Context, channel *model.Channel) channelTestSummary {
result := testChannelForHealthCheck(ctx, channel, testUserID, false, disableThreshold)
if result.Tested == 0 {
return result
}
mu.Lock()
defer mu.Unlock()
if persistErr == nil {
state.Rounds = append(state.Rounds, service.ChannelAvailabilityRound{ChannelID: channel.Id, At: time.Now().Unix(), Success: result.Succeeded == 1})
persistErr = model.UpdateSystemTaskState(task.TaskID, runnerID, state)
if persistErr != nil {
cancel()
}
}
return result
}, nil)
if persistErr != nil {
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, summary, persistErr)
return
}
if ctx.Err() != nil {
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, summary, ctx.Err())
return
}
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusSucceeded, summary, nil)
}
func (channelTestHandler) Type() string { return model.SystemTaskTypeChannelTest }
func (channelTestHandler) Enabled() bool {
+73
View File
@@ -0,0 +1,73 @@
package controller
import (
"net/http"
"net/http/httptest"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)
func TestSystemTaskListFiltersAndPaginationResponse(t *testing.T) {
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
sqlDB, err := db.DB()
require.NoError(t, err)
previousDB := model.DB
model.DB = db
t.Cleanup(func() {
model.DB = previousDB
require.NoError(t, sqlDB.Close())
})
require.NoError(t, db.AutoMigrate(&model.SystemTask{}))
require.NoError(t, db.Create(&[]model.SystemTask{
{TaskID: "older", Type: model.SystemTaskTypeModelUpdate, Status: model.SystemTaskStatusFailed},
{TaskID: "newer", Type: model.SystemTaskTypeModelUpdate, Status: model.SystemTaskStatusFailed},
{TaskID: "other-status", Type: model.SystemTaskTypeModelUpdate, Status: model.SystemTaskStatusSucceeded},
{TaskID: "other-type", Type: model.SystemTaskTypeChannelTest, Status: model.SystemTaskStatusFailed},
}).Error)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodGet, "/api/system-task/list?scope=history&type=model_update&status=failed&offset=1&limit=1", nil)
ListSystemTasks(c)
require.Equal(t, http.StatusOK, recorder.Code)
var response struct {
Success bool `json:"success"`
Total int64 `json:"total"`
Data []model.SystemTaskResponse `json:"data"`
}
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
assert.True(t, response.Success)
assert.EqualValues(t, 2, response.Total)
require.Len(t, response.Data, 1)
assert.Equal(t, "older", response.Data[0].TaskID)
}
func TestSystemTaskInvalidFiltersAreRejected(t *testing.T) {
for _, tc := range []struct {
query string
handler gin.HandlerFunc
}{
{"?scope=invalid", ListSystemTasks},
{"?status=unknown", ListSystemTasks},
{"?offset=-1", ListSystemTasks},
{"?limit=invalid", ListSystemTasks},
{"?status=unknown", DeleteSystemTaskHistory},
{"?scope=invalid", DeleteSystemTaskHistory},
} {
t.Run(tc.query, func(t *testing.T) {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodGet, "/"+tc.query, nil)
tc.handler(c)
assert.Equal(t, http.StatusBadRequest, recorder.Code)
assert.JSONEq(t, `{"success":false,"message":"invalid system task filters"}`, recorder.Body.String())
})
}
}
+27 -14
View File
@@ -15,6 +15,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/jsplugin"
@@ -26,6 +27,15 @@ import (
const maxTaskPluginSourceBytes = 1024 * 1024
func taskPluginCompileError(c *gin.Context, err error) {
var unknownField *jsplugin.UnknownMetaFieldError
if errors.As(err, &unknownField) {
common.ApiErrorI18n(c, i18n.MsgTaskPluginUnknownMetaField, map[string]any{"Field": unknownField.Field})
return
}
common.ApiError(c, err)
}
type taskPluginUploadRequest struct {
Source string `json:"source" binding:"required"`
Enabled *bool `json:"enabled"`
@@ -57,7 +67,7 @@ func UploadTaskPlugin(c *gin.Context) {
temporary := jsplugin.NewRegistry()
loaded, err := temporary.Register(request.Source, jsplugin.Options{})
if err != nil {
common.ApiErrorMsg(c, err.Error())
taskPluginCompileError(c, err)
return
}
if err = jsplugin.ValidateV1Meta(loaded.Meta); err != nil {
@@ -335,7 +345,7 @@ func GetTaskPlugin(c *gin.Context) {
if err == nil {
loaded, compileErr := jsplugin.NewRegistry().Register(plugin.Source, jsplugin.Options{Key: plugin.Key, Version: plugin.Version})
if compileErr != nil {
common.ApiErrorMsg(c, compileErr.Error())
taskPluginCompileError(c, compileErr)
return
}
common.ApiSuccess(c, taskPluginDetail{Plugin: plugin, Meta: loaded.Meta, Source: plugin.Source, Layer: "override", HasIcon: plugin.HasIcon()})
@@ -352,7 +362,7 @@ func GetTaskPlugin(c *gin.Context) {
}
loaded, err := jsplugin.NewRegistry().RegisterFactory(source, jsplugin.Options{Key: key})
if err != nil {
common.ApiError(c, err)
taskPluginCompileError(c, err)
return
}
_, _, hasIcon := plugins.Icon(key)
@@ -384,7 +394,7 @@ func DryRunTaskPlugin(c *gin.Context) {
}
loaded, err := jsplugin.NewRegistry().Register(detailSource, jsplugin.Options{Key: c.Param("key")})
if err != nil {
common.ApiErrorMsg(c, err.Error())
taskPluginCompileError(c, err)
return
}
args := make([]any, len(request.Args))
@@ -473,7 +483,7 @@ func ActivateTaskPlugin(c *gin.Context) {
return
}
if _, err = jsplugin.NewRegistry().Register(target.Source, jsplugin.Options{Key: target.Key, Version: target.Version}); err != nil {
common.ApiErrorMsg(c, err.Error())
taskPluginCompileError(c, err)
return
}
if err = model.ActivateTaskPlugin(target.Key, target.Version); err != nil {
@@ -649,15 +659,18 @@ func GetTaskPluginOptions(c *gin.Context) {
_, _, hasIcon = plugins.Icon(meta.Key)
}
options = append(options, gin.H{
"key": meta.Key,
"name": meta.Name,
"icon": meta.Icon,
"hasIcon": hasIcon,
"baseUrl": meta.BaseURL,
"sortPriority": meta.SortPriority,
"website": meta.Website,
"models": meta.Models,
"usageSchema": meta.UsageSchema,
"key": meta.Key,
"name": meta.Name,
"description": meta.Description,
"icon": meta.Icon,
"hasIcon": hasIcon,
"baseUrl": meta.BaseURL,
"sortPriority": meta.SortPriority,
"website": meta.Website,
"models": meta.Models,
"channelTypes": meta.ChannelTypes,
"usageSchema": meta.UsageSchema,
"usageProfiles": meta.UsageProfiles,
})
}
}
+52 -7
View File
@@ -13,6 +13,7 @@ import (
"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/pkg/jsplugin"
"github.com/QuantumNous/new-api/plugins"
@@ -358,15 +359,18 @@ func TestMasterSwitchEmptiesOptionsAndKeepsList(t *testing.T) {
assert.Equal(t, "kling", item.Meta.Key)
}
func TestGetTaskPluginOptionsIncludesUsageSchemaIconAndBaseURL(t *testing.T) {
func TestGetTaskPluginOptionsIncludesDescriptionUsageSchemaIconBaseURLAndChannelTypes(t *testing.T) {
setupTaskPluginControllerTest(t)
const key = "usage-options-probe"
source := `
export const meta = {
apiVersion: 1, key: "usage-options-probe", name: "Usage Options", version: "1.0.0", author: {name: "Test"},
description: {en: "Video generation via the vendor API", zh: "通过厂商接口生成视频"},
icon: "text:UO", baseUrl: "http://localhost:9000/",
channelTypes: [1990, 1991],
models: ["usage-options-model"], fetchMode: "per_task",
usageSchema: {seconds: {type: "number", unit: "second", description: "Generated media duration."}}
usageSchema: {seconds: {type: "number", unit: "second", description: "Video generation unit price"}},
usageProfiles: [{models:["usage-options-model"],schema:{image_count:{type:"number",unit:"count"}}}]
};
export function buildSubmitRequest() { return {}; }
export function parseSubmitResponse() { return {}; }
@@ -386,10 +390,13 @@ export function parseTaskResult() { return {}; }
var response struct {
Success bool `json:"success"`
Data []struct {
Key string `json:"key"`
Icon string `json:"icon"`
BaseURL string `json:"baseUrl"`
UsageSchema map[string]jsplugin.UsageFieldSchema `json:"usageSchema"`
Key string `json:"key"`
Description jsplugin.LocalizedText `json:"description"`
Icon string `json:"icon"`
BaseURL string `json:"baseUrl"`
ChannelTypes []int `json:"channelTypes"`
UsageSchema map[string]jsplugin.UsageFieldSchema `json:"usageSchema"`
UsageProfiles []jsplugin.UsageProfile `json:"usageProfiles"`
} `json:"data"`
}
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
@@ -398,9 +405,14 @@ export function parseTaskResult() { return {}; }
if option.Key != key {
continue
}
assert.Equal(t, jsplugin.LocalizedText{"en": "Video generation via the vendor API", "zh": "通过厂商接口生成视频"}, option.Description)
assert.Equal(t, "second", option.UsageSchema["seconds"].Unit)
assert.Equal(t, "Generated media duration.", option.UsageSchema["seconds"].Description["en"])
assert.Equal(t, "Video generation unit price", option.UsageSchema["seconds"].Description["en"])
require.Len(t, option.UsageProfiles, 1)
assert.Equal(t, []string{"usage-options-model"}, option.UsageProfiles[0].Models)
assert.Equal(t, "count", option.UsageProfiles[0].Schema["image_count"].Unit)
assert.Equal(t, "text:UO", option.Icon)
assert.Equal(t, []int{1990, 1991}, option.ChannelTypes)
assert.Equal(t, "http://localhost:9000", option.BaseURL, "the drawer prefills the normalized plugin default")
return
}
@@ -888,6 +900,39 @@ func TestUploadTaskPluginPreflightConflict(t *testing.T) {
}
}
func TestUploadTaskPluginLocalizesUnknownMetaField(t *testing.T) {
require.NoError(t, i18n.Init())
body, err := common.Marshal(map[string]any{
"source": `export const meta = { futureField: true };`,
})
require.NoError(t, err)
for _, tc := range []struct {
language string
message string
}{
{"en", `Plugin metadata contains an unknown field "futureField". If this plugin was downloaded from the official marketplace, it may require a newer version of new-api. Try updating new-api and installing the plugin again.`},
{"zh-CN", "插件元数据包含未知字段“futureField”。如果插件来自官方市场,可能需要更高版本的 new-api。请尝试更新 new-api 后重新安装插件。"},
{"zh-TW", "外掛中繼資料包含未知欄位「futureField」。如果外掛來自官方市集,可能需要較新版本的 new-api。請嘗試更新 new-api 後重新安裝外掛。"},
} {
t.Run(tc.language, func(t *testing.T) {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodPost, "/api/plugin/task", bytes.NewReader(body))
c.Request.Header.Set("Content-Type", "application/json")
c.Request.Header.Set("Accept-Language", tc.language)
UploadTaskPlugin(c)
require.Equal(t, http.StatusOK, recorder.Code)
var response struct {
Success bool `json:"success"`
Message string `json:"message"`
}
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
assert.False(t, response.Success)
assert.Equal(t, tc.message, response.Message)
})
}
}
func TestUploadTaskPluginRejectsMetaViolatingV1Schema(t *testing.T) {
setupTaskPluginControllerTest(t)
cases := []struct {
+152 -33
View File
@@ -32,14 +32,34 @@ func (input *tokenAutoGroupsInput) UnmarshalJSON(data []byte) error {
return common.Unmarshal(data, &input.Groups)
}
type tokenGroupPoolInput struct {
Set bool
Entries []model.TokenGroupEntry
}
func (input *tokenGroupPoolInput) UnmarshalJSON(data []byte) error {
input.Set = true
if strings.TrimSpace(string(data)) == "null" {
input.Entries = nil
return nil
}
return common.Unmarshal(data, &input.Entries)
}
type tokenRequest struct {
model.Token
AutoGroups tokenAutoGroupsInput `json:"auto_groups"`
AstraDetectionEnabled *bool `json:"astra_detection_enabled"`
AstraAutoSwitch *bool `json:"astra_auto_switch"`
RoutingPolicy *model.TokenRoutingPolicy `json:"routing_policy"`
AutoGroups tokenAutoGroupsInput `json:"auto_groups"`
GroupPool tokenGroupPoolInput `json:"group_pool"`
}
type tokenResponse struct {
*model.Token
AutoGroups []string `json:"auto_groups"`
RoutingPolicy *model.TokenRoutingPolicy `json:"routing_policy,omitempty"`
AutoGroups []string `json:"auto_groups"`
GroupPool []model.TokenGroupEntry `json:"group_pool,omitempty"`
}
func maxTokenQuota() int {
@@ -58,15 +78,20 @@ func buildMaskedTokenResponse(token *model.Token) *tokenResponse {
}
maskedToken := *token
maskedToken.Key = token.GetMaskedKey()
autoGroups, err := token.GetAutoGroups()
groupPool, err := token.GetGroupPool()
if err != nil {
common.SysError(fmt.Sprintf("failed to parse auto groups for token %d: %v", token.Id, err))
autoGroups = nil
groupPool = nil
}
autoGroups := make([]string, 0, len(groupPool))
for _, entry := range groupPool {
autoGroups = append(autoGroups, entry.Group)
}
if len(autoGroups) == 0 {
autoGroups = nil
}
return &tokenResponse{Token: &maskedToken, AutoGroups: autoGroups}
policy, _ := token.GetRoutingPolicy()
return &tokenResponse{Token: &maskedToken, AutoGroups: autoGroups, GroupPool: groupPool, RoutingPolicy: policy}
}
func buildMaskedTokenResponses(tokens []*model.Token) []*tokenResponse {
@@ -128,18 +153,56 @@ func setTokenAutoGroups(c *gin.Context, token *model.Token, groups []string) boo
}
func setTokenBackupGroup(c *gin.Context, token *model.Token, groups []string) bool {
if len(groups) > 1 {
common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsTooMany, map[string]any{"Max": 1})
maxCount := setting.GetMaxTokenAutoGroups()
if len(groups) > maxCount {
common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsTooMany, map[string]any{"Max": maxCount})
return false
}
if len(groups) == 1 && groups[0] == token.Group {
common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsDuplicate, map[string]any{"Group": groups[0]})
return false
for _, group := range groups {
if group == token.Group {
common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsDuplicate, map[string]any{"Group": group})
return false
}
}
if !setTokenAutoGroups(c, token, groups) {
return false
}
token.CrossGroupRetry = len(groups) == 1
token.CrossGroupRetry = len(groups) > 0
return true
}
func setTokenGroupPool(c *gin.Context, token *model.Token, entries []model.TokenGroupEntry) bool {
maxCount := setting.GetMaxTokenAutoGroups()
if len(entries) > maxCount {
common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsTooMany, map[string]any{"Max": maxCount})
return false
}
userGroup, err := getTokenRequestUserGroup(c)
if err != nil {
common.ApiError(c, err)
return false
}
seen := make(map[string]struct{}, len(entries))
for _, entry := range entries {
if entry.Group == token.Group {
common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsDuplicate, map[string]any{"Group": entry.Group})
return false
}
if _, ok := seen[entry.Group]; ok {
common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsDuplicate, map[string]any{"Group": entry.Group})
return false
}
seen[entry.Group] = struct{}{}
if !service.IsUserSelectableGroup(userGroup, entry.Group) {
common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsInvalid, map[string]any{"Group": entry.Group})
return false
}
}
if err := token.SetGroupPool(entries); err != nil {
common.ApiError(c, err)
return false
}
token.CrossGroupRetry = len(entries) > 0
return true
}
@@ -299,6 +362,16 @@ func AddToken(c *gin.Context) {
return
}
token := request.Token
token.AstraDetectionEnabled = true
if request.AstraDetectionEnabled != nil {
token.AstraDetectionEnabled = *request.AstraDetectionEnabled
}
if request.AstraAutoSwitch != nil {
token.AstraAutoSwitch = *request.AstraAutoSwitch
}
if !token.AstraDetectionEnabled {
token.AstraAutoSwitch = false
}
if len(token.Name) > 50 {
common.ApiErrorI18n(c, i18n.MsgTokenNameTooLong)
return
@@ -331,15 +404,36 @@ func AddToken(c *gin.Context) {
})
return
}
if token.Group == "auto" {
if request.RoutingPolicy != nil {
if !setTokenRoutingPolicy(c, &token, request.RoutingPolicy) {
return
}
} else if token.Group == "auto" {
if !setTokenAutoGroups(c, &token, request.AutoGroups.Groups) {
return
}
} else if request.GroupPool.Set {
if !setTokenGroupPool(c, &token, request.GroupPool.Entries) {
return
}
} else {
if !setTokenBackupGroup(c, &token, request.AutoGroups.Groups) {
return
}
}
if token.RoutingPolicy == "" {
policy, err := token.LegacyRoutingPolicy()
if err != nil {
common.ApiError(c, err)
return
}
body, err := common.Marshal(policy)
if err != nil {
common.ApiError(c, err)
return
}
token.RoutingPolicy = string(body)
}
key, err := common.GenerateKey()
if err != nil {
common.ApiErrorI18n(c, i18n.MsgTokenGenerateFailed)
@@ -347,20 +441,23 @@ func AddToken(c *gin.Context) {
return
}
cleanToken := model.Token{
UserId: c.GetInt("id"),
Name: token.Name,
Key: key,
CreatedTime: common.GetTimestamp(),
AccessedTime: common.GetTimestamp(),
ExpiredTime: token.ExpiredTime,
RemainQuota: token.RemainQuota,
UnlimitedQuota: token.UnlimitedQuota,
ModelLimitsEnabled: token.ModelLimitsEnabled,
ModelLimits: token.ModelLimits,
AllowIps: token.AllowIps,
Group: token.Group,
CrossGroupRetry: token.CrossGroupRetry,
AutoGroups: token.AutoGroups,
UserId: c.GetInt("id"),
Name: token.Name,
Key: key,
CreatedTime: common.GetTimestamp(),
AccessedTime: common.GetTimestamp(),
ExpiredTime: token.ExpiredTime,
RemainQuota: token.RemainQuota,
UnlimitedQuota: token.UnlimitedQuota,
ModelLimitsEnabled: token.ModelLimitsEnabled,
ModelLimits: token.ModelLimits,
AllowIps: token.AllowIps,
Group: token.Group,
CrossGroupRetry: token.CrossGroupRetry,
AstraAutoSwitch: token.AstraAutoSwitch,
AstraDetectionEnabled: token.AstraDetectionEnabled,
AutoGroups: token.AutoGroups,
RoutingPolicy: token.RoutingPolicy,
}
err = cleanToken.Insert()
if err != nil {
@@ -454,15 +551,34 @@ func UpdateToken(c *gin.Context) {
cleanToken.ModelLimitsEnabled = token.ModelLimitsEnabled
cleanToken.ModelLimits = token.ModelLimits
cleanToken.AllowIps = token.AllowIps
cleanToken.Group = token.Group
cleanToken.CrossGroupRetry = token.CrossGroupRetry
if token.Group != "auto" {
if !setTokenBackupGroup(c, cleanToken, request.AutoGroups.Groups) {
if request.AstraDetectionEnabled != nil {
cleanToken.AstraDetectionEnabled = *request.AstraDetectionEnabled
}
if request.AstraAutoSwitch != nil {
cleanToken.AstraAutoSwitch = *request.AstraAutoSwitch
}
if !cleanToken.AstraDetectionEnabled {
cleanToken.AstraAutoSwitch = false
}
if request.RoutingPolicy != nil {
if !setTokenRoutingPolicy(c, cleanToken, request.RoutingPolicy) {
return
}
} else if request.AutoGroups.Set {
if !setTokenAutoGroups(c, cleanToken, request.AutoGroups.Groups) {
return
} else if cleanToken.RoutingPolicy == "" {
cleanToken.Group = token.Group
cleanToken.CrossGroupRetry = token.CrossGroupRetry
if token.Group != "auto" {
if request.GroupPool.Set {
if !setTokenGroupPool(c, cleanToken, request.GroupPool.Entries) {
return
}
} else if !setTokenBackupGroup(c, cleanToken, request.AutoGroups.Groups) {
return
}
} else if request.AutoGroups.Set {
if !setTokenAutoGroups(c, cleanToken, request.AutoGroups.Groups) {
return
}
}
}
}
@@ -490,7 +606,10 @@ func UpdateToken(c *gin.Context) {
(previous.AllowIps != nil && cleanToken.AllowIps != nil && *previous.AllowIps != *cleanToken.AllowIps)},
{"group", previous.Group != cleanToken.Group},
{"cross_group_retry", previous.CrossGroupRetry != cleanToken.CrossGroupRetry},
{"astra_auto_switch", previous.AstraAutoSwitch != cleanToken.AstraAutoSwitch},
{"astra_detection_enabled", previous.AstraDetectionEnabled != cleanToken.AstraDetectionEnabled},
{"auto_groups", previous.AutoGroups != cleanToken.AutoGroups},
{"routing_policy", previous.RoutingPolicy != cleanToken.RoutingPolicy},
} {
if field.changed {
changedFields = append(changedFields, field.name)
+66
View File
@@ -0,0 +1,66 @@
package controller
import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting"
"github.com/gin-gonic/gin"
)
func setTokenRoutingPolicy(c *gin.Context, token *model.Token, policy *model.TokenRoutingPolicy) bool {
body, err := common.Marshal(policy)
if err != nil {
common.ApiError(c, err)
return false
}
if token.RoutingPolicy != "" && token.RoutingPolicy == string(body) {
return true
}
userGroup, err := getTokenRequestUserGroup(c)
if err != nil {
common.ApiError(c, err)
return false
}
if err := policy.Validate(setting.GetMaxTokenAutoGroups()+1, func(group string) bool { return service.IsUserSelectableGroup(userGroup, group) }); err != nil {
common.ApiError(c, err)
return false
}
token.RoutingPolicy = string(body)
token.CrossGroupRetry = true
token.AutoGroups = ""
if policy.Mode == "inherit" {
token.Group = "auto"
} else if policy.Mode == "user" {
token.Group = ""
} else {
token.Group = policy.Resolve("", nil, func(string) bool { return true })[0].Group
}
return true
}
func PreviewTokenRouting(c *gin.Context) {
var request struct {
Policy model.TokenRoutingPolicy `json:"routing_policy"`
Model string `json:"model"`
}
if err := c.ShouldBindJSON(&request); err != nil {
common.ApiError(c, err)
return
}
if err := request.Policy.Validate(setting.GetMaxTokenAutoGroups()+1, func(string) bool { return true }); err != nil {
common.ApiError(c, err)
return
}
group, err := getTokenRequestUserGroup(c)
if err != nil {
common.ApiError(c, err)
return
}
inherited := service.GetUserAutoGroup(group)
if request.Policy.Mode == "user" {
inherited = []string{group}
}
entries := request.Policy.Resolve(request.Model, inherited, func(candidate string) bool { return service.IsUserSelectableGroup(group, candidate) })
common.ApiSuccess(c, entries)
}
+104
View File
@@ -118,6 +118,91 @@ func setupTokenControllerTestDB(t *testing.T) *gorm.DB {
return db
}
func TestTokenAstraMigration(t *testing.T) {
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
for _, upgrade := range []bool{false, true} {
t.Run(fmt.Sprintf("%s/upgrade=%t", dialect, upgrade), func(t *testing.T) {
var db *gorm.DB
if dialect == "sqlite" {
db = openTokenControllerTestDB(t)
} else {
dsn := os.Getenv("TEST_" + strings.ToUpper(dialect) + "_DSN")
if dsn == "" {
t.Skip("external test database is not configured")
}
var managed *bool
db, managed = openTokenControllerExternalDB(t, dialect, dsn)
*managed = true
if db.Migrator().HasTable("options") {
t.Fatal("requires an empty test database")
}
}
t.Cleanup(func() { require.NoError(t, db.Migrator().DropTable(&model.Option{})) })
if upgrade {
require.NoError(t, db.AutoMigrate(&legacyToken{}))
require.NoError(t, db.Create(&legacyToken{UserId: 7, Key: "legacy-astra", Name: "preserved", RemainQuota: 123}).Error)
}
require.NoError(t, db.AutoMigrate(&model.Token{}, &model.Option{}))
require.NoError(t, model.MigrateTokenAstraAutoSwitch())
require.NoError(t, model.MigrateTokenAstraDetection())
if upgrade {
var legacy model.Token
require.NoError(t, db.First(&legacy, 1).Error)
require.True(t, legacy.AstraAutoSwitch)
require.True(t, legacy.AstraDetectionEnabled)
require.Equal(t, 123, legacy.RemainQuota)
legacy.AstraAutoSwitch = false
legacy.AstraDetectionEnabled = false
require.NoError(t, legacy.Update())
}
fresh := model.Token{UserId: 7, Key: "new-astra"}
require.NoError(t, db.Create(&fresh).Error)
for range 2 {
require.NoError(t, db.AutoMigrate(&model.Token{}, &model.Option{}))
require.NoError(t, model.MigrateTokenAstraAutoSwitch())
require.NoError(t, model.MigrateTokenAstraDetection())
}
var tokens []model.Token
require.NoError(t, db.Find(&tokens).Error)
for _, token := range tokens {
require.False(t, token.AstraAutoSwitch, "restart must preserve opt-outs and new-token defaults")
require.False(t, token.AstraDetectionEnabled)
}
})
}
}
}
func TestTokenAstraSwitchAPIPreservesOmittedAndPersistsFalse(t *testing.T) {
db := setupTokenControllerTestDB(t)
token := seedToken(t, db, 7, "astra", "astra-api-key")
token.AstraAutoSwitch = true
token.AstraDetectionEnabled = true
require.NoError(t, token.Update())
for _, enabled := range []any{nil, false, true} {
body := map[string]any{"id": token.Id, "name": "astra", "unlimited_quota": true, "group": "default", "expired_time": -1}
if enabled != nil {
body["astra_auto_switch"] = enabled
}
c, writer := newAuthenticatedContext(t, http.MethodPut, "/api/token/", body, 7)
UpdateToken(c)
require.True(t, decodeAPIResponse(t, writer).Success, writer.Body.String())
require.NoError(t, db.First(token, token.Id).Error)
want := enabled == nil || enabled == true
require.Equal(t, want, token.AstraAutoSwitch)
var response struct {
AstraAutoSwitch bool `json:"astra_auto_switch"`
}
require.NoError(t, common.Unmarshal(decodeAPIResponse(t, writer).Data, &response))
require.Equal(t, want, response.AstraAutoSwitch)
}
c, writer := newAuthenticatedContext(t, http.MethodPut, "/api/token/", map[string]any{"id": token.Id, "name": "foreign", "astra_auto_switch": false}, 8)
UpdateToken(c)
require.False(t, decodeAPIResponse(t, writer).Success)
require.NoError(t, db.First(token, token.Id).Error)
require.True(t, token.AstraAutoSwitch)
}
func openTokenControllerExternalDB(t *testing.T, dialect string, dsn string) (*gorm.DB, *bool) {
t.Helper()
@@ -892,3 +977,22 @@ func verifyAPITokenAudit(t *testing.T) {
assert.EqualValues(t, 1, count)
})
}
func TestTokenDetectionOffDisablesSwitchAndSurvivesOmittedUpdate(t *testing.T) {
db := setupTokenControllerTestDB(t)
token := seedToken(t, db, 7, "detection", "detection-api-key")
token.AstraDetectionEnabled, token.AstraAutoSwitch = true, true
require.NoError(t, token.Update())
for _, detection := range []any{false, nil, true} {
body := map[string]any{"id": token.Id, "name": "detection", "unlimited_quota": true, "group": "default", "expired_time": -1}
if detection != nil {
body["astra_detection_enabled"] = detection
}
c, writer := newAuthenticatedContext(t, http.MethodPut, "/api/token/", body, 7)
UpdateToken(c)
require.True(t, decodeAPIResponse(t, writer).Success, writer.Body.String())
require.NoError(t, db.First(token, token.Id).Error)
require.Equal(t, detection == true, token.AstraDetectionEnabled)
require.False(t, token.AstraAutoSwitch)
}
}
+83 -146
View File
@@ -2,9 +2,7 @@ package controller
import (
"context"
"errors"
"net/http"
"strconv"
"strings"
"sync"
"time"
@@ -15,29 +13,31 @@ import (
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
"golang.org/x/sync/errgroup"
)
const (
requestTimeout = 30 * time.Second
httpTimeout = 10 * time.Second
uptimeKeySuffix = "_24"
apiStatusPath = "/api/status-page/"
apiHeartbeatPath = "/api/status-page/heartbeat/"
)
const requestTimeout = 30 * time.Second
type Monitor struct {
Metrics *perfmetrics.GroupResult `json:"metrics,omitempty"`
Name string `json:"name"`
Uptime float64 `json:"uptime"`
Status int `json:"status"`
Group string `json:"group,omitempty"`
Ping int `json:"ping,omitempty"`
LastChecked string `json:"last_checked,omitempty"`
History []MonitorSample `json:"history,omitempty"`
CacheStats *model.GroupCacheStats `json:"cache_stats,omitempty"`
ProbeStatus int `json:"probe_status"`
ChannelDisabled bool `json:"channel_disabled"`
ProbeHistory []service.ProbeHistoryPoint `json:"probe_history"`
ID string `json:"id"`
Models []string `json:"models"`
Alias string `json:"alias,omitempty"`
Hidden bool `json:"hidden,omitempty"`
Pelican *service.PelicanSummary `json:"pelican,omitempty"`
StatusSource string `json:"status_source"`
StatusCheckedAt int64 `json:"status_checked_at,omitempty"`
Fingerprint *service.AstraSamplingSummary `json:"fingerprint,omitempty"`
Metrics *perfmetrics.GroupResult `json:"metrics,omitempty"`
Name string `json:"name"`
Uptime float64 `json:"uptime"`
Status int `json:"status"`
Group string `json:"group,omitempty"`
Ping int `json:"ping,omitempty"`
LastChecked string `json:"last_checked,omitempty"`
History []MonitorSample `json:"history,omitempty"`
CacheStats *model.GroupCacheStats `json:"cache_stats,omitempty"`
ProbeStatus int `json:"probe_status"`
ChannelDisabled bool `json:"channel_disabled"`
}
type MonitorSample struct {
@@ -51,113 +51,6 @@ type UptimeGroupResult struct {
Monitors []Monitor `json:"monitors"`
}
func getAndDecode(ctx context.Context, client *http.Client, url string, dest any) error {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
if err != nil {
return err
}
resp, err := client.Do(req)
if err != nil {
return err
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return errors.New("non-200 status")
}
return common.DecodeJson(resp.Body, dest)
}
func fetchGroupData(ctx context.Context, client *http.Client, groupConfig map[string]any) UptimeGroupResult {
url, _ := groupConfig["url"].(string)
slug, _ := groupConfig["slug"].(string)
categoryName, _ := groupConfig["categoryName"].(string)
result := UptimeGroupResult{
CategoryName: categoryName,
Monitors: []Monitor{},
}
if url == "" || slug == "" {
return result
}
baseURL := strings.TrimSuffix(url, "/")
var statusData struct {
PublicGroupList []struct {
ID int `json:"id"`
Name string `json:"name"`
MonitorList []struct {
ID int `json:"id"`
Name string `json:"name"`
} `json:"monitorList"`
} `json:"publicGroupList"`
}
var heartbeatData struct {
HeartbeatList map[string][]struct {
Status int `json:"status"`
Ping int `json:"ping"`
Time string `json:"time"`
} `json:"heartbeatList"`
UptimeList map[string]float64 `json:"uptimeList"`
}
g, gCtx := errgroup.WithContext(ctx)
g.Go(func() error {
return getAndDecode(gCtx, client, baseURL+apiStatusPath+slug, &statusData)
})
g.Go(func() error {
return getAndDecode(gCtx, client, baseURL+apiHeartbeatPath+slug, &heartbeatData)
})
if g.Wait() != nil {
return result
}
for _, pg := range statusData.PublicGroupList {
if len(pg.MonitorList) == 0 {
continue
}
for _, m := range pg.MonitorList {
monitor := Monitor{
Name: m.Name,
Group: pg.Name,
Status: -1,
}
monitorID := strconv.Itoa(m.ID)
if uptime, exists := heartbeatData.UptimeList[monitorID+uptimeKeySuffix]; exists {
monitor.Uptime = uptime
}
if heartbeats, exists := heartbeatData.HeartbeatList[monitorID]; exists && len(heartbeats) > 0 {
latest := heartbeats[len(heartbeats)-1]
monitor.Status = latest.Status
monitor.Ping = latest.Ping
monitor.LastChecked = latest.Time
monitor.History = make([]MonitorSample, 0, len(heartbeats))
for _, heartbeat := range heartbeats {
monitor.History = append(monitor.History, MonitorSample{
Status: heartbeat.Status,
Ping: heartbeat.Ping,
Time: heartbeat.Time,
})
}
}
result.Monitors = append(result.Monitors, monitor)
}
}
return result
}
var groupCacheSnapshot struct {
sync.Mutex
updated time.Time
@@ -168,18 +61,35 @@ var groupCacheSnapshot struct {
const requestStatusFreshness = 2 * time.Hour
func currentGroupMonitors(ctx context.Context, results []UptimeGroupResult) ([]UptimeGroupResult, error) {
func currentGroupMonitors(ctx context.Context, admin bool) ([]UptimeGroupResult, error) {
channels, err := model.GetStatusChannels(ctx)
if err != nil {
return nil, err
}
settings, err := model.GetStatusPageGroups(ctx)
if err != nil {
return nil, err
}
models, err := model.GetStatusPageModels(ctx, channels)
if err != nil {
return nil, err
}
usable := service.GetUserUsableGroups("")
now := time.Now()
fingerprints, fingerprintErr := service.AstraSamplingSummaries(ctx, now, settings)
if fingerprintErr != nil {
return nil, fingerprintErr
}
detections, detectionErr := service.GroupDetectionSummaries(ctx, now, fingerprints, settings)
if detectionErr != nil {
return nil, detectionErr
}
active := make(map[string]bool)
order := make([]string, 0)
for _, channel := range channels {
for _, group := range strings.Split(channel.Group, ",") {
group = strings.TrimSpace(group)
if _, visible := usable[group]; !visible {
if _, visible := usable[group]; (!visible || settings[group].Hidden) && !admin {
continue
}
if _, exists := active[group]; !exists {
@@ -202,8 +112,8 @@ func currentGroupMonitors(ctx context.Context, results []UptimeGroupResult) ([]U
}
groupCacheSnapshot.stats = stats
names := []string{}
for _, item := range model.GetPricing() {
names = append(names, item.ModelName)
if err := model.DB.WithContext(ctx).Model(&model.Ability{}).Where("enabled = ?", true).Distinct().Pluck("model", &names).Error; err != nil {
return nil, err
}
performance, metricsErr := perfmetrics.QueryGroups(names, 24)
if metricsErr != nil {
@@ -216,21 +126,26 @@ func currentGroupMonitors(ctx context.Context, results []UptimeGroupResult) ([]U
groupCacheSnapshot.updated = now
groupCacheSnapshot.channels = string(encoded)
}
byName := make(map[string]Monitor)
for _, result := range results {
for _, monitor := range result.Monitors {
if previous, exists := byName[monitor.Name]; !exists || (monitor.Status == 1 && previous.Status != 1) {
byName[monitor.Name] = monitor
}
}
}
successes := perfmetrics.LatestUserSuccess(ctx, order)
monitors := make([]Monitor, 0, len(order))
for _, group := range order {
monitor, exists := byName[group]
if !exists {
monitor = Monitor{Name: group, Status: -1}
monitor := Monitor{ID: model.DetectionSecretHash(group)[:16], Name: group, Status: -1, Models: models[group]}
if monitor.Models == nil {
monitor.Models = []string{}
}
if settings[group].Alias != "" {
monitor.Name = settings[group].Alias
}
if admin {
monitor.Alias, monitor.Hidden = settings[group].Alias, settings[group].Hidden
}
if admin {
monitor.Group = group
}
monitor.Fingerprint = fingerprints[group]
if monitor.Fingerprint == nil {
monitor.Fingerprint = &service.AstraSamplingSummary{Enabled: service.AstraSamplingEnabled(), Available: fingerprintErr == nil, History: []service.AstraSamplingHistory{}}
}
monitor.Group = group
monitor.ProbeStatus = monitor.Status
monitor.ChannelDisabled = !active[group]
if !active[group] {
@@ -239,16 +154,30 @@ func currentGroupMonitors(ctx context.Context, results []UptimeGroupResult) ([]U
stats := groupCacheSnapshot.stats[group]
monitor.CacheStats = &stats
applyRequestMetrics(&monitor, groupCacheSnapshot.metrics[group])
if !admin && monitor.Metrics != nil {
monitor.Metrics.Group = monitor.Name
}
monitor.Pelican = detections[group].Pelican
monitor.StatusSource = "requests"
if monitor.ChannelDisabled {
monitor.StatusSource = "disabled"
} else if monitor.Status < 0 {
monitor.Status = detections[group].Status
monitor.StatusSource = "detection"
monitor.StatusCheckedAt = detections[group].At
}
monitor.ProbeHistory = detections[group].ProbeHistory
applyUserSuccess(&monitor, successes[group], detections[group].At, now.Unix())
monitors = append(monitors, monitor)
}
return []UptimeGroupResult{{CategoryName: "Groups", Monitors: monitors}}, nil
}
func GetUptimeKumaStatus(c *gin.Context) {
func GetGroupStatus(c *gin.Context) {
ctx, cancel := context.WithTimeout(c.Request.Context(), requestTimeout)
defer cancel()
results, err := currentGroupMonitors(ctx, nil)
results, err := currentGroupMonitors(ctx, false)
if err != nil {
common.SysError("group status statistics: " + err.Error())
c.JSON(http.StatusServiceUnavailable, gin.H{"success": false, "message": "Status data is temporarily unavailable"})
@@ -267,6 +196,14 @@ func requestStatus(rate float64) int {
return 0
}
func applyUserSuccess(monitor *Monitor, success, probe, now int64) {
if !monitor.ChannelDisabled && success > 0 && success >= now-900 && success >= probe {
monitor.Status = 1
monitor.StatusSource = "requests"
monitor.StatusCheckedAt = success
}
}
func applyRequestMetrics(monitor *Monitor, metrics perfmetrics.GroupResult) {
monitor.Status = -1
monitor.Uptime = 0
+191 -31
View File
@@ -3,39 +3,199 @@ package controller
import (
"context"
"fmt"
"net/http"
"net/http/httptest"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/middleware"
"github.com/QuantumNous/new-api/model"
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/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"
"gorm.io/gorm"
"gorm.io/gorm/schema"
"net/http/httptest"
"os"
"strings"
"testing"
"time"
)
func TestFetchGroupDataUsesLatestHeartbeat(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/api/status-page/status":
_, _ = fmt.Fprint(w, `{"publicGroupList":[{"id":1,"name":"hidden","monitorList":[{"id":7,"name":"GPT Mix"}]}]}`)
case "/api/status-page/heartbeat/status":
_, _ = fmt.Fprint(w, `{"heartbeatList":{"7":[{"status":0,"time":"2026-09-05 10:00:00","ping":0},{"status":1,"time":"2026-09-05 10:03:00","ping":321}]},"uptimeList":{"7_24":0.95}}`)
default:
http.NotFound(w, r)
}
}))
t.Cleanup(server.Close)
result := fetchGroupData(context.Background(), server.Client(), map[string]interface{}{
"url": server.URL,
"slug": "status",
"categoryName": "API",
})
require.Len(t, result.Monitors, 1)
monitor := result.Monitors[0]
assert.Equal(t, "GPT Mix", monitor.Name)
assert.Equal(t, 1, monitor.Status)
assert.Equal(t, 321, monitor.Ping)
assert.Equal(t, "2026-09-05 10:03:00", monitor.LastChecked)
assert.Equal(t, 0.95, monitor.Uptime)
assert.Len(t, monitor.History, 2)
func TestStatusPageDatabaseMatrix(t *testing.T) {
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
t.Run(dialect, func(t *testing.T) {
var driver gorm.Dialector
switch dialect {
case "sqlite":
driver = sqlite.Open(":memory:")
case "mysql":
if os.Getenv("TEST_MYSQL_DSN") == "" {
t.Skip("TEST_MYSQL_DSN unset")
}
driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
case "postgres":
if os.Getenv("TEST_POSTGRES_DSN") == "" {
t.Skip("TEST_POSTGRES_DSN unset")
}
driver = postgres.Open(os.Getenv("TEST_POSTGRES_DSN"))
}
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("status_%d_", time.Now().UnixNano())}})
require.NoError(t, err)
conn, err := db.DB()
require.NoError(t, err)
conn.SetMaxOpenConns(1)
oldDB, oldLog := model.DB, model.LOG_DB
oldType, oldLogType := common.MainDatabaseType(), common.LogDatabaseType()
oldRedis, oldSecret, oldSecure := common.RedisEnabled, common.SessionSecret, common.SessionCookieSecure
oldGroups := setting.UserUsableGroups2JSONString()
model.DB, model.LOG_DB = db, db
common.SetDatabaseTypes(common.DatabaseType(dialect), common.DatabaseType(dialect))
common.RedisEnabled, common.SessionSecret, common.SessionCookieSecure = false, "status-test-only", true
rows := []any{&model.Option{}, &model.Channel{}, &model.Ability{}, &model.SystemTask{}, &model.DetectionRun{}, &model.DetectionCall{}, &model.DetectionSchedule{}, &model.Log{}, &model.PerfMetric{}, &model.User{}, &model.UserSession{}}
require.NoError(t, db.AutoMigrate(rows...))
require.NoError(t, db.Table("audit_logs").AutoMigrate(&model.AuditLog{}))
t.Cleanup(func() {
require.NoError(t, db.Migrator().DropTable("audit_logs"))
for _, row := range rows {
require.NoError(t, db.Migrator().DropTable(row))
}
model.DB, model.LOG_DB = oldDB, oldLog
common.SetDatabaseTypes(oldType, oldLogType)
common.RedisEnabled, common.SessionSecret, common.SessionCookieSecure = oldRedis, oldSecret, oldSecure
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(oldGroups))
groupCacheSnapshot.updated = time.Time{}
require.NoError(t, conn.Close())
})
var version string
q := "select version()"
if dialect == "sqlite" {
q = "select sqlite_version()"
}
require.NoError(t, db.Raw(q).Scan(&version).Error)
t.Log("database version:", version)
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"internal":"","shared":"","hidden":""}`))
now := time.Now().Unix()
channels := []model.Channel{{Id: 1, Name: "private-channel", Key: "private-key", Group: "internal,shared", Models: fingerprint.Model, Status: 1}, {Id: 2, Group: "hidden", Models: fingerprint.Model, Status: 1}, {Id: 3, Group: "internal", Models: "disabled-model", Status: 2}}
require.NoError(t, db.Create(&channels).Error)
abilities := []model.Ability{{Group: "internal", Model: fingerprint.Model, ChannelId: 1, Enabled: true}, {Group: "shared", Model: fingerprint.Model, ChannelId: 1, Enabled: true}, {Group: "hidden", Model: fingerprint.Model, ChannelId: 2, Enabled: true}, {Group: "internal", Model: "disabled-model", ChannelId: 3, Enabled: true}, {Group: "internal", Model: "disabled-ability", ChannelId: 1}}
require.NoError(t, db.Create(&abilities).Error)
state, err := common.Marshal(service.AstraSamplingState{Rounds: []service.AstraSamplingRound{{ID: "old-round", ChannelID: 1, StartedAt: now - 20, FinishedAt: now - 10, Result: fingerprint.Result{Status: "match", Closest: fingerprint.Model, ValidSamples: 3}}}})
require.NoError(t, err)
require.NoError(t, db.Create(&model.SystemTask{TaskID: "old-task", Type: service.AstraSamplingTaskType, Status: model.SystemTaskStatusSucceeded, CreatedAt: now - 20, State: string(state)}).Error)
runs := []model.DetectionRun{{ID: "old-work", RequestKey: "old-work", ChannelID: 1, GroupName: "internal", Kind: "pelican", Status: "succeeded", Quality: "normal", CreatedAt: now - 20, Settled: true, HTML: "<h1>art</h1>", Charged: 17}, {ID: "in-flight", RequestKey: "in-flight", ChannelID: 1, Kind: "pelican", Status: "running", CreatedAt: now - 10}, {ID: "hidden-art", RequestKey: "hidden-art", ChannelID: 2, Kind: "pelican", Status: "succeeded", CreatedAt: now - 20, HTML: "hidden"}}
require.NoError(t, db.Create(&runs).Error)
tokens := map[int]string{}
for _, role := range []int{1, 10} {
require.NoError(t, db.Create(&model.User{Id: role, Username: fmt.Sprintf("role-%d", role), Role: role, Status: 1, AuthVersion: 1, AffCode: fmt.Sprintf("role-%d", role)}).Error)
sid := fmt.Sprintf("session-%d", role)
require.NoError(t, db.Create(&model.UserSession{SID: sid, UserID: role, Version: 1, UserAuthVersion: 1, Status: model.UserSessionStatusActive, RefreshHash: sid, ExpiresAt: now + 3600}).Error)
tokens[role], _, err = service.IssueAccessToken(service.AuthIdentity{UserID: role, SessionID: sid, UserAuthVersion: 1, SessionVersion: 1})
require.NoError(t, err)
}
engine := gin.New()
engine.GET("/api/uptime/status", GetGroupStatus)
engine.GET("/api/group-status", GetSiteGroupStatus)
engine.PUT("/api/uptime/group", middleware.AdminAuth(), middleware.SessionCookieOriginGuard(), UpdateStatusPageGroup)
engine.POST("/api/uptime/group/reset", middleware.AdminAuth(), middleware.SessionCookieOriginGuard(), UpdateStatusPageGroup)
engine.GET("/api/detections/artwork/:id", GetDetectionArtwork)
engine.GET("/api/group-status/artwork/:id", GetSiteDetectionArtwork)
request := func(method, path, body, token, origin string) *httptest.ResponseRecorder {
req := httptest.NewRequest(method, "https://status.test"+path, strings.NewReader(body))
req.Header.Set("Content-Type", "application/json")
if token != "" {
req.Header.Set("Authorization", "Bearer "+token)
}
if origin != "" {
req.Header.Set("Origin", origin)
}
out := httptest.NewRecorder()
engine.ServeHTTP(out, req)
return out
}
for _, token := range []string{"", tokens[1]} {
assert.Contains(t, []int{401, 403}, request("POST", "/api/uptime/group/reset", `{"group":"internal"}`, token, "https://status.test").Code)
}
assert.Equal(t, 403, request("POST", "/api/uptime/group/reset", `{"group":"internal"}`, tokens[10], "https://evil.test").Code)
assert.Equal(t, 400, request("PUT", "/api/uptime/group", `{"group":"internal"}`, tokens[10], "https://status.test").Code)
assert.Equal(t, 404, request("POST", "/api/uptime/group/reset", `{"group":"missing"}`, tokens[10], "https://status.test").Code)
for _, body := range []string{`{"group":"internal","alias":"#688","hidden":false}`, `{"group":"hidden","alias":"secret-alias","hidden":true}`} {
res := request("PUT", "/api/uptime/group", body, tokens[10], "https://status.test")
require.Equal(t, 200, res.Code, res.Body.String())
}
var auditCount int64
require.NoError(t, db.Table("audit_logs").Count(&auditCount).Error)
assert.Greater(t, auditCount, int64(0))
groupCacheSnapshot.updated = time.Time{}
res := request("GET", "/api/uptime/status", "", "", "")
require.Equal(t, 200, res.Code, res.Body.String())
for _, secret := range []string{"internal", "secret-alias", "private-channel", "private-key", "disabled-model", "disabled-ability"} {
assert.NotContains(t, res.Body.String(), secret)
}
assert.Contains(t, res.Body.String(), "#688")
fallback := model.DetectionRun{ID:"fallback-completed",RequestKey:"fallback-completed",Kind:"pelican",ChannelID:1,GroupName:"internal",Status:"running",Quality:"normal",HTML:"<html>pelican</html>",CreatedAt:now,Deadline:now+600}
require.NoError(t,db.Create(&fallback).Error)
require.NoError(t,db.Create(&model.DetectionCall{ID:fallback.ID,RunID:fallback.ID,SecretHash:"fallback-only-test",State:"finished"}).Error)
processLocalPelican(context.Background(),fallback)
require.NoError(t,db.First(&fallback,"id = ?",fallback.ID).Error)
assert.Equal(t,"succeeded",fallback.Status)
assert.Equal(t,"unknown",fallback.Quality)
assert.Equal(t,"<html>pelican</html>",fallback.HTML)
assert.True(t,fallback.Settled)
require.NoError(t,db.Delete(&model.DetectionCall{},"id = ?",fallback.ID).Error)
require.NoError(t,db.Delete(&model.DetectionRun{},"id = ?",fallback.ID).Error)
visitorSite := request("GET", "/api/group-status", "", "", "")
require.Equal(t, 200, visitorSite.Code)
assert.Contains(t, visitorSite.Body.String(), `"group":"hidden"`)
assert.Contains(t, visitorSite.Body.String(), `"name":"internal"`)
assert.NotContains(t, visitorSite.Body.String(), "#688")
site := request("GET", "/api/group-status", "", tokens[1], "")
require.Equal(t, 200, site.Code)
assert.Contains(t, site.Body.String(), `"group":"internal"`)
assert.Contains(t, site.Body.String(), `"name":"internal"`)
assert.NotContains(t, site.Body.String(), "#688")
assert.Contains(t, site.Body.String(), `"group":"hidden"`)
adminSite := request("GET", "/api/group-status", "", tokens[10], "")
assert.Contains(t, adminSite.Body.String(), `"group":"hidden"`)
assert.Equal(t, 404, request("GET", "/api/detections/artwork/hidden-art", "", "", "").Code)
assert.Equal(t, 200, request("GET", "/api/group-status/artwork/hidden-art", "", "", "").Code)
require.Equal(t, 200, request("POST", "/api/uptime/group/reset", `{"group":"internal"}`, tokens[10], "https://status.test").Code)
require.NoError(t, db.Model(&model.DetectionRun{}).Where("id = ?", "in-flight").Updates(map[string]any{"settled": true, "status": "succeeded", "quality": "normal", "finished_at": now + 10}).Error)
groups, err := currentGroupMonitors(context.Background(), true)
require.NoError(t, err)
byGroup := map[string]Monitor{}
for _, m := range groups[0].Monitors {
byGroup[m.Group] = m
}
assert.Empty(t, byGroup["internal"].Fingerprint.History)
assert.Nil(t, byGroup["internal"].Fingerprint.Rate)
assert.Empty(t, byGroup["internal"].Pelican.History)
assert.Len(t, byGroup["shared"].Fingerprint.History, 1)
assert.Len(t, byGroup["shared"].Pelican.History, 2)
assert.True(t, byGroup["hidden"].Hidden)
assert.Equal(t, []string{fingerprint.Model}, byGroup["internal"].Models)
require.NoError(t, model.SaveStatusPageGroup(context.Background(), "internal", &model.StatusPageGroup{Alias: "#689", Hidden: true}))
settings, err := model.GetStatusPageGroups(context.Background())
require.NoError(t, err)
assert.GreaterOrEqual(t, settings["internal"].ResetBefore, now)
assert.Equal(t, "#689", settings["internal"].Alias)
require.NoError(t, model.SaveStatusPageGroup(context.Background(), "shared", &model.StatusPageGroup{Hidden: true}))
assert.Equal(t, 404, request("GET", "/api/detections/artwork/old-work", "", "", "").Code)
var retained model.DetectionRun
require.NoError(t, db.First(&retained, "id = ?", "old-work").Error)
assert.Equal(t, 17, retained.Charged)
assert.Equal(t, "<h1>art</h1>", retained.HTML)
require.NoError(t, db.Create(&model.DetectionRun{ID: "new-work", RequestKey: "new-work", ChannelID: 1, Kind: "pelican", Status: "succeeded", Quality: "normal", CreatedAt: settings["internal"].ResetBefore + 1, Settled: true}).Error)
groups, err = currentGroupMonitors(context.Background(), true)
require.NoError(t, err)
for _, m := range groups[0].Monitors {
if m.Group == "internal" {
require.Len(t, m.Pelican.History, 1)
assert.Equal(t, "new-work", m.Pelican.History[0].ID)
}
}
})
}
}
+11 -4
View File
@@ -96,7 +96,7 @@ func Login(c *gin.Context) {
return
}
setupLogin(&user, c)
setupLogin(&user, nil, c)
}
// loginMethodFromContext 根据请求路径推导登录方式,用于登录审计日志。
@@ -144,9 +144,11 @@ func recordLoginAudit(user *model.User, c *gin.Context) {
}
// setupLogin evaluates the shared login policy after primary authentication.
// Only a completed Passkey ceremony may go directly to session issuance.
func setupLogin(user *model.User, c *gin.Context) {
challenge, err := service.StartLoginVerification(user, loginMethodFromContext(c))
// Only a completed Passkey ceremony may go directly to session issuance. A
// pending legacy GitHub binding rewrite travels inside the challenge and is
// written only when the verification completes.
func setupLogin(user *model.User, migration *service.LegacyGitHubMigration, c *gin.Context) {
challenge, err := service.StartLoginVerification(user, loginMethodFromContext(c), migration)
if err != nil {
writeSecurityOperationError(c, err)
return
@@ -1278,6 +1280,7 @@ type UpdateUserSettingRequest struct {
UpstreamModelUpdateNotifyEnabled *bool `json:"upstream_model_update_notify_enabled,omitempty"`
AcceptUnsetModelRatioModel bool `json:"accept_unset_model_ratio_model"`
RecordIpLog bool `json:"record_ip_log"`
AstraCostNoticeAck *bool `json:"astra_cost_notice_ack,omitempty"`
}
func UpdateUserSetting(c *gin.Context) {
@@ -1380,6 +1383,10 @@ func UpdateUserSetting(c *gin.Context) {
UpstreamModelUpdateNotifyEnabled: upstreamModelUpdateNotifyEnabled,
AcceptUnsetRatioModel: req.AcceptUnsetModelRatioModel,
RecordIpLog: req.RecordIpLog,
AstraCostNoticeAck: existingSettings.AstraCostNoticeAck,
}
if req.AstraCostNoticeAck != nil {
settings.AstraCostNoticeAck = *req.AstraCostNoticeAck
}
// 如果是webhook类型,添加webhook相关设置
+1 -1
View File
@@ -121,7 +121,7 @@ func WeChatAuth(c *gin.Context) {
})
return
}
setupLogin(&user, c)
setupLogin(&user, nil, c)
}
type wechatBindRequest struct {