mirror of
https://github.com/QuantumNous/new-api.git
synced 2026-10-28 08:35:33 +08:00
CaIon
73a471f3a8
Group sensitive-word checks, channel selection, retries and channel health under one "Request policies" settings area with a session rules editor, per-rule session modes (inherit/off/prefer/strict) and a global default. Record routing decisions per attempt into log admin info, expose GET/PATCH /api/option/request_policy with whole-policy validation, and redirect the legacy settings URLs to the new sections.
468 lines
16 KiB
Go
468 lines
16 KiB
Go
package service
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"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/types"
|
|
"github.com/QuantumNous/new-api/setting/operation_setting"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func buildChannelAffinityTemplateContextForTest(meta channelAffinityMeta) *gin.Context {
|
|
rec := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(rec)
|
|
setChannelAffinityContext(ctx, meta)
|
|
return ctx
|
|
}
|
|
|
|
func TestApplyChannelAffinityOverrideTemplate_NoTemplate(t *testing.T) {
|
|
ctx := buildChannelAffinityTemplateContextForTest(channelAffinityMeta{
|
|
RuleName: "rule-no-template",
|
|
})
|
|
base := map[string]any{
|
|
"temperature": 0.7,
|
|
}
|
|
|
|
merged, applied := ApplyChannelAffinityOverrideTemplate(ctx, base)
|
|
require.False(t, applied)
|
|
require.Equal(t, base, merged)
|
|
}
|
|
|
|
func TestApplyChannelAffinityOverrideTemplate_MergeTemplate(t *testing.T) {
|
|
ctx := buildChannelAffinityTemplateContextForTest(channelAffinityMeta{
|
|
RuleName: "rule-with-template",
|
|
ParamTemplate: map[string]any{
|
|
"temperature": 0.2,
|
|
"top_p": 0.95,
|
|
},
|
|
UsingGroup: "default",
|
|
ModelName: "gpt-4.1",
|
|
RequestPath: "/v1/responses",
|
|
KeySourceType: "gjson",
|
|
KeySourcePath: "prompt_cache_key",
|
|
KeyHint: "abcd...wxyz",
|
|
KeyFingerprint: "abcd1234",
|
|
})
|
|
base := map[string]any{
|
|
"temperature": 0.7,
|
|
"max_tokens": 2000,
|
|
}
|
|
|
|
merged, applied := ApplyChannelAffinityOverrideTemplate(ctx, base)
|
|
require.True(t, applied)
|
|
require.Equal(t, 0.7, merged["temperature"])
|
|
require.Equal(t, 0.95, merged["top_p"])
|
|
require.Equal(t, 2000, merged["max_tokens"])
|
|
require.Equal(t, 0.7, base["temperature"])
|
|
|
|
anyInfo, ok := ctx.Get(ginKeyChannelAffinityLogInfo)
|
|
require.True(t, ok)
|
|
info, ok := anyInfo.(map[string]any)
|
|
require.True(t, ok)
|
|
overrideInfoAny, ok := info["override_template"]
|
|
require.True(t, ok)
|
|
overrideInfo, ok := overrideInfoAny.(map[string]any)
|
|
require.True(t, ok)
|
|
require.Equal(t, true, overrideInfo["applied"])
|
|
require.Equal(t, "rule-with-template", overrideInfo["rule_name"])
|
|
require.EqualValues(t, 2, overrideInfo["param_override_keys"])
|
|
}
|
|
|
|
func TestApplyChannelAffinityOverrideTemplate_MergeOperations(t *testing.T) {
|
|
ctx := buildChannelAffinityTemplateContextForTest(channelAffinityMeta{
|
|
RuleName: "rule-with-ops-template",
|
|
ParamTemplate: map[string]any{
|
|
"operations": []map[string]any{
|
|
{
|
|
"mode": "pass_headers",
|
|
"value": []string{"Originator"},
|
|
},
|
|
},
|
|
},
|
|
})
|
|
base := map[string]any{
|
|
"temperature": 0.7,
|
|
"operations": []map[string]any{
|
|
{
|
|
"path": "model",
|
|
"mode": "trim_prefix",
|
|
"value": "openai/",
|
|
},
|
|
},
|
|
}
|
|
|
|
merged, applied := ApplyChannelAffinityOverrideTemplate(ctx, base)
|
|
require.True(t, applied)
|
|
require.Equal(t, 0.7, merged["temperature"])
|
|
|
|
opsAny, ok := merged["operations"]
|
|
require.True(t, ok)
|
|
ops, ok := opsAny.([]any)
|
|
require.True(t, ok)
|
|
require.Len(t, ops, 2)
|
|
|
|
firstOp, ok := ops[0].(map[string]any)
|
|
require.True(t, ok)
|
|
require.Equal(t, "pass_headers", firstOp["mode"])
|
|
|
|
secondOp, ok := ops[1].(map[string]any)
|
|
require.True(t, ok)
|
|
require.Equal(t, "trim_prefix", secondOp["mode"])
|
|
}
|
|
|
|
func TestShouldSkipRetryAfterChannelAffinityFailure(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
ctx func() *gin.Context
|
|
want bool
|
|
}{
|
|
{
|
|
name: "nil context",
|
|
ctx: func() *gin.Context {
|
|
return nil
|
|
},
|
|
want: false,
|
|
},
|
|
{
|
|
name: "explicit skip retry flag in context",
|
|
ctx: func() *gin.Context {
|
|
ctx := buildChannelAffinityTemplateContextForTest(channelAffinityMeta{
|
|
RuleName: "rule-explicit-flag",
|
|
SkipRetry: false,
|
|
UsingGroup: "default",
|
|
ModelName: "gpt-5",
|
|
})
|
|
ctx.Set(ginKeyChannelAffinitySkipRetry, true)
|
|
return ctx
|
|
},
|
|
want: true,
|
|
},
|
|
{
|
|
name: "fallback to matched rule meta",
|
|
ctx: func() *gin.Context {
|
|
return buildChannelAffinityTemplateContextForTest(channelAffinityMeta{
|
|
RuleName: "rule-skip-retry",
|
|
SkipRetry: true,
|
|
UsingGroup: "default",
|
|
ModelName: "gpt-5",
|
|
})
|
|
},
|
|
want: true,
|
|
},
|
|
{
|
|
name: "no flag and no skip retry meta",
|
|
ctx: func() *gin.Context {
|
|
return buildChannelAffinityTemplateContextForTest(channelAffinityMeta{
|
|
RuleName: "rule-no-skip-retry",
|
|
SkipRetry: false,
|
|
UsingGroup: "default",
|
|
ModelName: "gpt-5",
|
|
})
|
|
},
|
|
want: false,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
require.Equal(t, tt.want, ShouldSkipRetryAfterChannelAffinityFailure(tt.ctx()))
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestExtractChannelAffinityValue_RequestHeader(t *testing.T) {
|
|
rec := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(rec)
|
|
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
|
ctx.Request.Header.Set("X-Affinity-Key", " tenant-123 ")
|
|
|
|
value := extractChannelAffinityValue(ctx, operation_setting.ChannelAffinityKeySource{
|
|
Type: "request_header",
|
|
Key: "X-Affinity-Key",
|
|
})
|
|
|
|
require.Equal(t, "tenant-123", value)
|
|
}
|
|
|
|
func TestGetPreferredChannelByAffinity_RequestHeaderKeySource(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
rule := operation_setting.ChannelAffinityRule{
|
|
Name: "header-affinity",
|
|
ModelRegex: []string{"^gpt-.*$"},
|
|
PathRegex: []string{"/v1/responses"},
|
|
KeySources: []operation_setting.ChannelAffinityKeySource{
|
|
{Type: "request_header", Key: "X-Affinity-Key"},
|
|
},
|
|
IncludeRuleName: true,
|
|
IncludeModelName: true,
|
|
}
|
|
|
|
affinityValue := fmt.Sprintf("header-hit-%d", time.Now().UnixNano())
|
|
cacheKeySuffix := buildChannelAffinityCacheKeySuffix(rule, "gpt-5", "default", affinityValue)
|
|
|
|
cache := getChannelAffinityCache()
|
|
require.NoError(t, cache.SetWithTTL(cacheKeySuffix, 9528, time.Minute))
|
|
t.Cleanup(func() {
|
|
_, _ = cache.DeleteMany([]string{cacheKeySuffix})
|
|
})
|
|
|
|
setting := operation_setting.GetChannelAffinitySetting()
|
|
originalRules := setting.Rules
|
|
setting.Rules = append([]operation_setting.ChannelAffinityRule{rule}, originalRules...)
|
|
t.Cleanup(func() {
|
|
setting.Rules = originalRules
|
|
})
|
|
|
|
rec := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(rec)
|
|
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
|
ctx.Request.Header.Set("X-Affinity-Key", affinityValue)
|
|
|
|
channelID, found := GetPreferredChannelByAffinity(ctx, "gpt-5", "default")
|
|
require.True(t, found)
|
|
require.Equal(t, 9528, channelID)
|
|
|
|
meta, ok := getChannelAffinityMeta(ctx)
|
|
require.True(t, ok)
|
|
require.Equal(t, "request_header", meta.KeySourceType)
|
|
require.Equal(t, "X-Affinity-Key", meta.KeySourceKey)
|
|
require.Equal(t, buildChannelAffinityKeyHint(affinityValue), meta.KeyHint)
|
|
}
|
|
|
|
func TestClearCurrentChannelAffinityCache(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
cacheKeySuffix := fmt.Sprintf("codex cli trace:default:clear-current-%d", time.Now().UnixNano())
|
|
cacheKeyFull := channelAffinityCacheNamespace + ":" + cacheKeySuffix
|
|
cache := getChannelAffinityCache()
|
|
require.NoError(t, cache.SetWithTTL(cacheKeySuffix, 9527, time.Minute))
|
|
t.Cleanup(func() {
|
|
_, _ = cache.DeleteMany([]string{cacheKeySuffix})
|
|
})
|
|
|
|
ctx := buildChannelAffinityTemplateContextForTest(channelAffinityMeta{
|
|
CacheKey: cacheKeyFull,
|
|
TTLSeconds: 60,
|
|
RuleName: "codex cli trace",
|
|
SkipRetry: true,
|
|
})
|
|
require.True(t, ShouldSkipRetryAfterChannelAffinityFailure(ctx))
|
|
|
|
deleted := ClearCurrentChannelAffinityCache(ctx)
|
|
require.True(t, deleted)
|
|
_, found, err := cache.Get(cacheKeySuffix)
|
|
require.NoError(t, err)
|
|
require.False(t, found)
|
|
require.False(t, ShouldSkipRetryAfterChannelAffinityFailure(ctx))
|
|
}
|
|
|
|
func TestChannelAffinityHitCodexTemplatePassHeadersEffective(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
setting := operation_setting.GetChannelAffinitySetting()
|
|
require.NotNil(t, setting)
|
|
|
|
var codexRule *operation_setting.ChannelAffinityRule
|
|
for i := range setting.Rules {
|
|
rule := &setting.Rules[i]
|
|
if strings.EqualFold(strings.TrimSpace(rule.Name), "codex cli trace") {
|
|
codexRule = rule
|
|
break
|
|
}
|
|
}
|
|
require.NotNil(t, codexRule)
|
|
|
|
affinityValue := fmt.Sprintf("pc-hit-%d", time.Now().UnixNano())
|
|
cacheKeySuffix := buildChannelAffinityCacheKeySuffix(*codexRule, "gpt-5", "default", affinityValue)
|
|
|
|
cache := getChannelAffinityCache()
|
|
require.NoError(t, cache.SetWithTTL(cacheKeySuffix, 9527, time.Minute))
|
|
t.Cleanup(func() {
|
|
_, _ = cache.DeleteMany([]string{cacheKeySuffix})
|
|
})
|
|
|
|
rec := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(rec)
|
|
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(fmt.Sprintf(`{"prompt_cache_key":"%s"}`, affinityValue)))
|
|
ctx.Request.Header.Set("Content-Type", "application/json")
|
|
|
|
channelID, found := GetPreferredChannelByAffinity(ctx, "gpt-5", "default")
|
|
require.True(t, found)
|
|
require.Equal(t, 9527, channelID)
|
|
|
|
baseOverride := map[string]any{
|
|
"temperature": 0.2,
|
|
}
|
|
mergedOverride, applied := ApplyChannelAffinityOverrideTemplate(ctx, baseOverride)
|
|
require.True(t, applied)
|
|
require.Equal(t, 0.2, mergedOverride["temperature"])
|
|
|
|
info := &relaycommon.RelayInfo{
|
|
RequestHeaders: map[string]string{
|
|
"Originator": "Codex CLI",
|
|
"Session_id": "sess-123",
|
|
"User-Agent": "codex-cli-test",
|
|
},
|
|
ChannelMeta: &relaycommon.ChannelMeta{
|
|
ParamOverride: mergedOverride,
|
|
HeadersOverride: map[string]any{
|
|
"X-Static": "legacy-static",
|
|
},
|
|
},
|
|
}
|
|
|
|
_, err := relaycommon.ApplyParamOverrideWithRelayInfo([]byte(`{"model":"gpt-5"}`), info)
|
|
require.NoError(t, err)
|
|
require.True(t, info.UseRuntimeHeadersOverride)
|
|
|
|
require.Equal(t, "legacy-static", info.RuntimeHeadersOverride["x-static"])
|
|
require.Equal(t, "Codex CLI", info.RuntimeHeadersOverride["originator"])
|
|
require.Equal(t, "sess-123", info.RuntimeHeadersOverride["session_id"])
|
|
require.Equal(t, "codex-cli-test", info.RuntimeHeadersOverride["user-agent"])
|
|
|
|
_, exists := info.RuntimeHeadersOverride["x-codex-beta-features"]
|
|
require.False(t, exists)
|
|
_, exists = info.RuntimeHeadersOverride["x-codex-turn-metadata"]
|
|
require.False(t, exists)
|
|
}
|
|
|
|
func TestMidjourneyPolicyAcceptance(t *testing.T) {
|
|
for _, tc := range []struct {
|
|
name string
|
|
status int
|
|
code int
|
|
properties map[string]any
|
|
requestErr error
|
|
accepted bool
|
|
reason string
|
|
}{
|
|
{name: "submitted", status: 200, code: 1, accepted: true},
|
|
{name: "existing task", status: 200, code: 21, accepted: true},
|
|
{name: "queued", status: 200, code: 22, accepted: true},
|
|
{name: "HTTP 200 with rejected prompt", status: 200, code: 24, reason: "non_retryable_error"},
|
|
{name: "failed existing task", status: 200, code: 21, properties: map[string]any{"status": "FAILURE"}, reason: "task_accepted"},
|
|
{name: "HTTP failure", status: 502, code: 1, reason: "non_retryable_error"},
|
|
{name: "ambiguous transport failure", status: 502, requestErr: errors.New("connection reset"), reason: "non_retryable_error"},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
ctx.Request = httptest.NewRequest(http.MethodPost, "/mj/submit/imagine", strings.NewReader(`{}`))
|
|
ctx.Set("channel_id", 7)
|
|
state := RequestPolicy(ctx)
|
|
state.BeginAttempt(&model.Channel{Id: 7}, "default")
|
|
response := &dto.MidjourneyResponseWithStatusCode{StatusCode: tc.status, Response: dto.MidjourneyResponse{Code: tc.code, Properties: tc.properties}}
|
|
accepted := RecordMidjourneyPolicyResponse(ctx, response, tc.requestErr)
|
|
assert.Equal(t, tc.accepted, accepted)
|
|
assert.False(t, state.Successful, "acceptance alone must not bind before local processing finishes")
|
|
if accepted {
|
|
MarkRequestPolicySuccess(ctx, nil)
|
|
assert.True(t, state.Successful)
|
|
return
|
|
}
|
|
events := state.Events()
|
|
require.Len(t, events, 3)
|
|
assert.Equal(t, PolicyDecision{Action: "failure", Reason: "upstream_failure", Source: "upstream"}, events[1].Decision)
|
|
assert.Equal(t, tc.status, events[1].Status)
|
|
assert.Equal(t, PolicyDecision{Action: "stop", Reason: tc.reason, Source: "system"}, events[2].Decision, "submissions are never replayed")
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestSessionRulesInheritOrOverrideGlobalDefault(t *testing.T) {
|
|
setting := operation_setting.GetChannelAffinitySetting()
|
|
previous := *setting
|
|
t.Cleanup(func() { *setting = previous })
|
|
for _, globalMode := range []string{"", "off", "prefer", "strict"} {
|
|
for _, tc := range []struct {
|
|
name, mode, expected string
|
|
skipRetry bool
|
|
}{
|
|
{name: "inherit", mode: "inherit", expected: globalMode},
|
|
{name: "override-strict", mode: "strict", expected: "strict"},
|
|
{name: "override-prefer", mode: "prefer", expected: "prefer", skipRetry: true},
|
|
{name: "override-off", mode: "off", expected: "off", skipRetry: true},
|
|
{name: "legacy-strict", expected: "strict", skipRetry: true},
|
|
{name: "legacy-prefer", expected: "prefer"},
|
|
} {
|
|
t.Run(globalMode+"/"+tc.name, func(t *testing.T) {
|
|
expected := tc.expected
|
|
if expected == "" {
|
|
expected = "prefer"
|
|
}
|
|
source := "session_rule"
|
|
if tc.mode == "inherit" {
|
|
source = "global"
|
|
}
|
|
rule := operation_setting.ChannelAffinityRule{
|
|
Name: tc.name, ModelRegex: []string{"^test-model$"},
|
|
KeySources: []operation_setting.ChannelAffinityKeySource{{Type: "request_header", Key: "X-Session"}},
|
|
SessionMode: tc.mode, SkipRetryOnFailure: tc.skipRetry,
|
|
ParamOverrideTemplate: map[string]any{"temperature": 0}, IncludeRuleName: true,
|
|
}
|
|
encoded, err := common.Marshal([]operation_setting.ChannelAffinityRule{rule})
|
|
require.NoError(t, err)
|
|
snapshot, err := model.BuildRequestPolicy(map[string]string{"RetryTimes": "2", "channel_affinity_setting.enabled": "true", "channel_affinity_setting.session_mode": globalMode, "channel_affinity_setting.rules": string(encoded)})
|
|
require.NoError(t, err)
|
|
*setting = snapshot.Affinity
|
|
key := t.Name()
|
|
cacheKey := buildChannelAffinityCacheKeySuffix(rule, "test-model", "default", key)
|
|
cache := getChannelAffinityCache()
|
|
require.NoError(t, cache.SetWithTTL(cacheKey, 1, time.Minute))
|
|
t.Cleanup(func() { _, err := cache.DeleteMany([]string{cacheKey}); assert.NoError(t, err) })
|
|
ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
|
ctx.Request.Header.Set("X-Session", key)
|
|
state := RequestPolicy(ctx)
|
|
channelID, found := GetPreferredChannelByAffinity(ctx, "test-model", "default")
|
|
assert.Equal(t, expected != "off", found)
|
|
if found {
|
|
assert.Equal(t, 1, channelID)
|
|
}
|
|
assert.Equal(t, expected, state.SessionMode)
|
|
assert.Equal(t, source, state.SessionModeSource)
|
|
assert.Equal(t, expected == "strict", ShouldSkipRetryAfterChannelAffinityFailure(ctx))
|
|
decision := DecideRelayRetry(ctx, types.NewOpenAIError(errors.New("upstream"), types.ErrorCodeBadResponseStatusCode, http.StatusTooManyRequests), 2)
|
|
if expected == "strict" {
|
|
assert.Equal(t, PolicyDecision{Action: "stop", Reason: "strict_session", Source: source}, decision)
|
|
} else {
|
|
assert.Equal(t, "retry", decision.Action)
|
|
}
|
|
events := state.Events()
|
|
require.Len(t, events, 1)
|
|
assert.Equal(t, PolicyDecision{Action: "match", Reason: "session_rule_matched", Source: "session_rule"}, events[0].Decision)
|
|
assert.Equal(t, tc.name, events[0].Rule)
|
|
transformed, applied := ApplyChannelAffinityOverrideTemplate(ctx, nil)
|
|
assert.True(t, applied, "session behavior does not disable request transforms")
|
|
assert.EqualValues(t, 0, transformed["temperature"])
|
|
if expected == "off" {
|
|
ctx.Set("channel_id", 9)
|
|
RecordChannelAffinity(ctx, 9)
|
|
cached, found, err := cache.Get(cacheKey)
|
|
require.NoError(t, err)
|
|
assert.True(t, found)
|
|
assert.Equal(t, 1, cached, "off does not refresh session bindings")
|
|
_, err = cache.DeleteMany([]string{cacheKey})
|
|
require.NoError(t, err)
|
|
RecordChannelAffinity(ctx, 9)
|
|
_, found, err = cache.Get(cacheKey)
|
|
require.NoError(t, err)
|
|
assert.False(t, found, "off does not establish session bindings")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
}
|