feat: sync deployed v1.0.0-rc.38-custom.42
Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
This commit is contained in:
cneicy committed
1 parent
aa6c3f50f4
commit
ceddfd3d6b
828 files changed
+122253
-13702
No files matched your search
@@ -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")
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user