Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions service/channel_affinity.go
Original file line number Diff line number Diff line change
Expand Up @@ -302,6 +302,11 @@ func extractChannelAffinityValue(c *gin.Context, src operation_setting.ChannelAf
return ""
}
return strings.TrimSpace(c.GetString(src.Key))
case "request_header":
if c == nil || c.Request == nil || src.Key == "" {
return ""
}
return strings.TrimSpace(c.Request.Header.Get(src.Key))
case "gjson":
if src.Path == "" {
return ""
Expand Down
60 changes: 60 additions & 0 deletions service/channel_affinity_template_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -176,6 +176,66 @@ func TestShouldSkipRetryAfterChannelAffinityFailure(t *testing.T) {
}
}

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 TestChannelAffinityHitCodexTemplatePassHeadersEffective(t *testing.T) {
gin.SetMode(gin.TestMode)

Expand Down
2 changes: 1 addition & 1 deletion setting/operation_setting/channel_affinity_setting.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ package operation_setting
import "github.com/QuantumNous/new-api/setting/config"

type ChannelAffinityKeySource struct {
Type string `json:"type"` // context_int, context_string, gjson
Type string `json:"type"` // context_int, context_string, request_header, gjson
Key string `json:"key,omitempty"`
Path string `json:"path,omitempty"`
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,7 @@ const KEY_RULES = 'channel_affinity_setting.rules';
const KEY_SOURCE_TYPES = [
{ label: 'context_int', value: 'context_int' },
{ label: 'context_string', value: 'context_string' },
{ label: 'request_header', value: 'request_header' },
{ label: 'gjson', value: 'gjson' },
];

Expand Down Expand Up @@ -659,7 +660,11 @@ export default function SettingsChannelAffinity(props) {
const xs = (keySources || []).map(normalizeKeySource).filter((x) => x.type);
if (xs.length === 0) return { ok: false, message: 'Key 来源不能为空' };
for (const x of xs) {
if (x.type === 'context_int' || x.type === 'context_string') {
if (
x.type === 'context_int' ||
x.type === 'context_string' ||
x.type === 'request_header'
) {
if (!x.key) return { ok: false, message: 'Key 不能为空' };
} else if (x.type === 'gjson') {
if (!x.path) return { ok: false, message: 'Path 不能为空' };
Expand Down Expand Up @@ -1316,7 +1321,7 @@ export default function SettingsChannelAffinity(props) {
</Space>
<Text type='tertiary' size='small'>
{t(
'context_int/context_string 从请求上下文读取;gjson 从入口请求的 JSON body 按 gjson path 读取。',
'context_int/context_string 从请求上下文读取;request_header 从用户请求头读取;gjson 从入口请求的 JSON body 按 gjson path 读取。',
)}
</Text>
<div style={{ marginTop: 8, marginBottom: 8 }}>
Expand Down Expand Up @@ -1358,7 +1363,7 @@ export default function SettingsChannelAffinity(props) {
return (
<Input
placeholder={
isGjson ? 'metadata.conversation_id' : 'user_id'
isGjson ? 'metadata.conversation_id' : 'X-Affinity-Key'
}
aria-label={t('Key 或 Path')}
value={isGjson ? src.path : src.key}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,12 @@ import { Textarea } from '@/components/ui/textarea'
import { RULE_TEMPLATES } from './constants'
import type { AffinityRule, KeySource } from './types'

const KEY_SOURCE_TYPES = ['context_int', 'context_string', 'gjson'] as const
const KEY_SOURCE_TYPES = [
'context_int',
'context_string',
'request_header',
'gjson',
] as const

const CONTEXT_KEY_PRESETS = [
'id',
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ along with this program. If not, see <https://www.gnu.org/licenses/>.
For commercial licensing, please contact support@quantumnous.com
*/
export interface KeySource {
type: 'context_int' | 'context_string' | 'gjson'
type: 'context_int' | 'context_string' | 'request_header' | 'gjson'
key?: string
path?: string
}
Expand Down