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
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:
1 parent
c2b44332d6
commit
a5f7b8e586
828 files changed
+122253
-13702
No files matched your search
@@ -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}})
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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]
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
@@ -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),
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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", "")
|
||||
}
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
@@ -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,
|
||||
|
||||
@@ -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{
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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])
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
|
||||
}
|
||||
|
||||
@@ -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" }
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
@@ -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})
|
||||
}
|
||||
@@ -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}))
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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())
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
@@ -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})
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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})
|
||||
}
|
||||
@@ -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},
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
@@ -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,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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
@@ -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
@@ -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
@@ -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相关设置
|
||||
|
||||
@@ -121,7 +121,7 @@ func WeChatAuth(c *gin.Context) {
|
||||
})
|
||||
return
|
||||
}
|
||||
setupLogin(&user, c)
|
||||
setupLogin(&user, nil, c)
|
||||
}
|
||||
|
||||
type wechatBindRequest struct {
|
||||
|
||||
Reference in new issue
Block a user