Commit 8103b4b1 by Seefs Committed by GitHub

Merge pull request #3009 from seefs001/feature/improve-param-override

feat: improve channel override ui/ux
parents 30061f3c a53e1393
...@@ -366,7 +366,7 @@ func testChannel(channel *model.Channel, testModel string, endpointType string, ...@@ -366,7 +366,7 @@ func testChannel(channel *model.Channel, testModel string, endpointType string,
newAPIError: types.NewError(err, types.ErrorCodeConvertRequestFailed), newAPIError: types.NewError(err, types.ErrorCodeConvertRequestFailed),
} }
} }
jsonData, err := json.Marshal(convertedRequest) jsonData, err := common.Marshal(convertedRequest)
if err != nil { if err != nil {
return testResult{ return testResult{
context: c, context: c,
...@@ -385,8 +385,15 @@ func testChannel(channel *model.Channel, testModel string, endpointType string, ...@@ -385,8 +385,15 @@ func testChannel(channel *model.Channel, testModel string, endpointType string,
//} //}
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info)) jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info)
if err != nil { if err != nil {
if fixedErr, ok := relaycommon.AsParamOverrideReturnError(err); ok {
return testResult{
context: c,
localErr: fixedErr,
newAPIError: relaycommon.NewAPIErrorFromParamOverride(fixedErr),
}
}
return testResult{ return testResult{
context: c, context: c,
localErr: err, localErr: err,
......
...@@ -182,8 +182,11 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { ...@@ -182,8 +182,11 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
ModelName: relayInfo.OriginModelName, ModelName: relayInfo.OriginModelName,
Retry: common.GetPointer(0), Retry: common.GetPointer(0),
} }
relayInfo.RetryIndex = 0
relayInfo.LastError = nil
for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() { for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() {
relayInfo.RetryIndex = retryParam.GetRetry()
channel, channelErr := getChannel(c, relayInfo, retryParam) channel, channelErr := getChannel(c, relayInfo, retryParam)
if channelErr != nil { if channelErr != nil {
logger.LogError(c, channelErr.Error()) logger.LogError(c, channelErr.Error())
...@@ -216,10 +219,12 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { ...@@ -216,10 +219,12 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
} }
if newAPIError == nil { if newAPIError == nil {
relayInfo.LastError = nil
return return
} }
newAPIError = service.NormalizeViolationFeeError(newAPIError) newAPIError = service.NormalizeViolationFeeError(newAPIError)
relayInfo.LastError = newAPIError
processChannelError(c, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()), newAPIError) processChannelError(c, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()), newAPIError)
......
...@@ -348,8 +348,13 @@ func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, mode ...@@ -348,8 +348,13 @@ func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, mode
common.SetContextKey(c, constant.ContextKeyChannelCreateTime, channel.CreatedTime) common.SetContextKey(c, constant.ContextKeyChannelCreateTime, channel.CreatedTime)
common.SetContextKey(c, constant.ContextKeyChannelSetting, channel.GetSetting()) common.SetContextKey(c, constant.ContextKeyChannelSetting, channel.GetSetting())
common.SetContextKey(c, constant.ContextKeyChannelOtherSetting, channel.GetOtherSettings()) common.SetContextKey(c, constant.ContextKeyChannelOtherSetting, channel.GetOtherSettings())
common.SetContextKey(c, constant.ContextKeyChannelParamOverride, channel.GetParamOverride()) paramOverride := channel.GetParamOverride()
common.SetContextKey(c, constant.ContextKeyChannelHeaderOverride, channel.GetHeaderOverride()) headerOverride := channel.GetHeaderOverride()
if mergedParam, applied := service.ApplyChannelAffinityOverrideTemplate(c, paramOverride); applied {
paramOverride = mergedParam
}
common.SetContextKey(c, constant.ContextKeyChannelParamOverride, paramOverride)
common.SetContextKey(c, constant.ContextKeyChannelHeaderOverride, headerOverride)
if nil != channel.OpenAIOrganization && *channel.OpenAIOrganization != "" { if nil != channel.OpenAIOrganization && *channel.OpenAIOrganization != "" {
common.SetContextKey(c, constant.ContextKeyChannelOrganization, *channel.OpenAIOrganization) common.SetContextKey(c, constant.ContextKeyChannelOrganization, *channel.OpenAIOrganization)
} }
......
...@@ -169,12 +169,17 @@ func applyHeaderOverridePlaceholders(template string, c *gin.Context, apiKey str ...@@ -169,12 +169,17 @@ func applyHeaderOverridePlaceholders(template string, c *gin.Context, apiKey str
// Passthrough rules are applied first, then normal overrides are applied, so explicit overrides win. // Passthrough rules are applied first, then normal overrides are applied, so explicit overrides win.
func processHeaderOverride(info *common.RelayInfo, c *gin.Context) (map[string]string, error) { func processHeaderOverride(info *common.RelayInfo, c *gin.Context) (map[string]string, error) {
headerOverride := make(map[string]string) headerOverride := make(map[string]string)
if info == nil {
return headerOverride, nil
}
headerOverrideSource := common.GetEffectiveHeaderOverride(info)
passAll := false passAll := false
var passthroughRegex []*regexp.Regexp var passthroughRegex []*regexp.Regexp
if !info.IsChannelTest { if !info.IsChannelTest {
for k := range info.HeadersOverride { for k := range headerOverrideSource {
key := strings.TrimSpace(k) key := strings.TrimSpace(strings.ToLower(k))
if key == "" { if key == "" {
continue continue
} }
...@@ -183,12 +188,11 @@ func processHeaderOverride(info *common.RelayInfo, c *gin.Context) (map[string]s ...@@ -183,12 +188,11 @@ func processHeaderOverride(info *common.RelayInfo, c *gin.Context) (map[string]s
continue continue
} }
lower := strings.ToLower(key)
var pattern string var pattern string
switch { switch {
case strings.HasPrefix(lower, headerPassthroughRegexPrefix): case strings.HasPrefix(key, headerPassthroughRegexPrefix):
pattern = strings.TrimSpace(key[len(headerPassthroughRegexPrefix):]) pattern = strings.TrimSpace(key[len(headerPassthroughRegexPrefix):])
case strings.HasPrefix(lower, headerPassthroughRegexPrefixV2): case strings.HasPrefix(key, headerPassthroughRegexPrefixV2):
pattern = strings.TrimSpace(key[len(headerPassthroughRegexPrefixV2):]) pattern = strings.TrimSpace(key[len(headerPassthroughRegexPrefixV2):])
default: default:
continue continue
...@@ -229,15 +233,15 @@ func processHeaderOverride(info *common.RelayInfo, c *gin.Context) (map[string]s ...@@ -229,15 +233,15 @@ func processHeaderOverride(info *common.RelayInfo, c *gin.Context) (map[string]s
if value == "" { if value == "" {
continue continue
} }
headerOverride[name] = value headerOverride[strings.ToLower(strings.TrimSpace(name))] = value
} }
} }
for k, v := range info.HeadersOverride { for k, v := range headerOverrideSource {
if isHeaderPassthroughRuleKey(k) { if isHeaderPassthroughRuleKey(k) {
continue continue
} }
key := strings.TrimSpace(k) key := strings.TrimSpace(strings.ToLower(k))
if key == "" { if key == "" {
continue continue
} }
......
...@@ -53,7 +53,7 @@ func TestProcessHeaderOverride_ChannelTestSkipsClientHeaderPlaceholder(t *testin ...@@ -53,7 +53,7 @@ func TestProcessHeaderOverride_ChannelTestSkipsClientHeaderPlaceholder(t *testin
headers, err := processHeaderOverride(info, ctx) headers, err := processHeaderOverride(info, ctx)
require.NoError(t, err) require.NoError(t, err)
_, ok := headers["X-Upstream-Trace"] _, ok := headers["x-upstream-trace"]
require.False(t, ok) require.False(t, ok)
} }
...@@ -77,7 +77,38 @@ func TestProcessHeaderOverride_NonTestKeepsClientHeaderPlaceholder(t *testing.T) ...@@ -77,7 +77,38 @@ func TestProcessHeaderOverride_NonTestKeepsClientHeaderPlaceholder(t *testing.T)
headers, err := processHeaderOverride(info, ctx) headers, err := processHeaderOverride(info, ctx)
require.NoError(t, err) require.NoError(t, err)
require.Equal(t, "trace-123", headers["X-Upstream-Trace"]) require.Equal(t, "trace-123", headers["x-upstream-trace"])
}
func TestProcessHeaderOverride_RuntimeOverrideIsFinalHeaderMap(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
info := &relaycommon.RelayInfo{
IsChannelTest: false,
UseRuntimeHeadersOverride: true,
RuntimeHeadersOverride: map[string]any{
"x-static": "runtime-value",
"x-runtime": "runtime-only",
},
ChannelMeta: &relaycommon.ChannelMeta{
HeadersOverride: map[string]any{
"X-Static": "legacy-value",
"X-Legacy": "legacy-only",
},
},
}
headers, err := processHeaderOverride(info, ctx)
require.NoError(t, err)
require.Equal(t, "runtime-value", headers["x-static"])
require.Equal(t, "runtime-only", headers["x-runtime"])
_, exists := headers["x-legacy"]
require.False(t, exists)
} }
func TestProcessHeaderOverride_PassthroughSkipsAcceptEncoding(t *testing.T) { func TestProcessHeaderOverride_PassthroughSkipsAcceptEncoding(t *testing.T) {
...@@ -101,8 +132,62 @@ func TestProcessHeaderOverride_PassthroughSkipsAcceptEncoding(t *testing.T) { ...@@ -101,8 +132,62 @@ func TestProcessHeaderOverride_PassthroughSkipsAcceptEncoding(t *testing.T) {
headers, err := processHeaderOverride(info, ctx) headers, err := processHeaderOverride(info, ctx)
require.NoError(t, err) require.NoError(t, err)
require.Equal(t, "trace-123", headers["X-Trace-Id"]) require.Equal(t, "trace-123", headers["x-trace-id"])
_, hasAcceptEncoding := headers["Accept-Encoding"] _, hasAcceptEncoding := headers["accept-encoding"]
require.False(t, hasAcceptEncoding) require.False(t, hasAcceptEncoding)
} }
func TestProcessHeaderOverride_PassHeadersTemplateSetsRuntimeHeaders(t *testing.T) {
t.Parallel()
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
ctx.Request.Header.Set("Originator", "Codex CLI")
ctx.Request.Header.Set("Session_id", "sess-123")
info := &relaycommon.RelayInfo{
IsChannelTest: false,
RequestHeaders: map[string]string{
"Originator": "Codex CLI",
"Session_id": "sess-123",
},
ChannelMeta: &relaycommon.ChannelMeta{
ParamOverride: map[string]any{
"operations": []any{
map[string]any{
"mode": "pass_headers",
"value": []any{"Originator", "Session_id", "X-Codex-Beta-Features"},
},
},
},
HeadersOverride: map[string]any{
"X-Static": "legacy-value",
},
},
}
_, err := relaycommon.ApplyParamOverrideWithRelayInfo([]byte(`{"model":"gpt-4.1"}`), info)
require.NoError(t, err)
require.True(t, info.UseRuntimeHeadersOverride)
require.Equal(t, "Codex CLI", info.RuntimeHeadersOverride["originator"])
require.Equal(t, "sess-123", info.RuntimeHeadersOverride["session_id"])
_, exists := info.RuntimeHeadersOverride["x-codex-beta-features"]
require.False(t, exists)
require.Equal(t, "legacy-value", info.RuntimeHeadersOverride["x-static"])
headers, err := processHeaderOverride(info, ctx)
require.NoError(t, err)
require.Equal(t, "Codex CLI", headers["originator"])
require.Equal(t, "sess-123", headers["session_id"])
_, exists = headers["x-codex-beta-features"]
require.False(t, exists)
upstreamReq := httptest.NewRequest(http.MethodPost, "https://example.com/v1/responses", nil)
applyHeaderOverrideToRequest(upstreamReq, headers)
require.Equal(t, "Codex CLI", upstreamReq.Header.Get("Originator"))
require.Equal(t, "sess-123", upstreamReq.Header.Get("Session_id"))
require.Empty(t, upstreamReq.Header.Get("X-Codex-Beta-Features"))
}
...@@ -70,7 +70,6 @@ func applySystemPromptIfNeeded(c *gin.Context, info *relaycommon.RelayInfo, requ ...@@ -70,7 +70,6 @@ func applySystemPromptIfNeeded(c *gin.Context, info *relaycommon.RelayInfo, requ
} }
func chatCompletionsViaResponses(c *gin.Context, info *relaycommon.RelayInfo, adaptor channel.Adaptor, request *dto.GeneralOpenAIRequest) (*dto.Usage, *types.NewAPIError) { func chatCompletionsViaResponses(c *gin.Context, info *relaycommon.RelayInfo, adaptor channel.Adaptor, request *dto.GeneralOpenAIRequest) (*dto.Usage, *types.NewAPIError) {
overrideCtx := relaycommon.BuildParamOverrideContext(info)
chatJSON, err := common.Marshal(request) chatJSON, err := common.Marshal(request)
if err != nil { if err != nil {
return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
...@@ -82,9 +81,9 @@ func chatCompletionsViaResponses(c *gin.Context, info *relaycommon.RelayInfo, ad ...@@ -82,9 +81,9 @@ func chatCompletionsViaResponses(c *gin.Context, info *relaycommon.RelayInfo, ad
} }
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
chatJSON, err = relaycommon.ApplyParamOverride(chatJSON, info.ParamOverride, overrideCtx) chatJSON, err = relaycommon.ApplyParamOverrideWithRelayInfo(chatJSON, info)
if err != nil { if err != nil {
return nil, types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid, types.ErrOptionWithSkipRetry()) return nil, newAPIErrorFromParamOverride(err)
} }
} }
......
...@@ -153,9 +153,9 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ ...@@ -153,9 +153,9 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
// apply param override // apply param override
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info)) jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info)
if err != nil { if err != nil {
return types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid, types.ErrOptionWithSkipRetry()) return newAPIErrorFromParamOverride(err)
} }
} }
......
...@@ -101,6 +101,7 @@ type RelayInfo struct { ...@@ -101,6 +101,7 @@ type RelayInfo struct {
RelayMode int RelayMode int
OriginModelName string OriginModelName string
RequestURLPath string RequestURLPath string
RequestHeaders map[string]string
ShouldIncludeUsage bool ShouldIncludeUsage bool
DisablePing bool // 是否禁止向下游发送自定义 Ping DisablePing bool // 是否禁止向下游发送自定义 Ping
ClientWs *websocket.Conn ClientWs *websocket.Conn
...@@ -144,6 +145,10 @@ type RelayInfo struct { ...@@ -144,6 +145,10 @@ type RelayInfo struct {
SubscriptionAmountUsedAfterPreConsume int64 SubscriptionAmountUsedAfterPreConsume int64
IsClaudeBetaQuery bool // /v1/messages?beta=true IsClaudeBetaQuery bool // /v1/messages?beta=true
IsChannelTest bool // channel test request IsChannelTest bool // channel test request
RetryIndex int
LastError *types.NewAPIError
RuntimeHeadersOverride map[string]interface{}
UseRuntimeHeadersOverride bool
PriceData types.PriceData PriceData types.PriceData
...@@ -461,6 +466,7 @@ func genBaseRelayInfo(c *gin.Context, request dto.Request) *RelayInfo { ...@@ -461,6 +466,7 @@ func genBaseRelayInfo(c *gin.Context, request dto.Request) *RelayInfo {
isFirstResponse: true, isFirstResponse: true,
RelayMode: relayconstant.Path2RelayMode(c.Request.URL.Path), RelayMode: relayconstant.Path2RelayMode(c.Request.URL.Path),
RequestURLPath: c.Request.URL.String(), RequestURLPath: c.Request.URL.String(),
RequestHeaders: cloneRequestHeaders(c),
IsStream: isStream, IsStream: isStream,
StartTime: startTime, StartTime: startTime,
...@@ -493,6 +499,27 @@ func genBaseRelayInfo(c *gin.Context, request dto.Request) *RelayInfo { ...@@ -493,6 +499,27 @@ func genBaseRelayInfo(c *gin.Context, request dto.Request) *RelayInfo {
return info return info
} }
func cloneRequestHeaders(c *gin.Context) map[string]string {
if c == nil || c.Request == nil {
return nil
}
if len(c.Request.Header) == 0 {
return nil
}
headers := make(map[string]string, len(c.Request.Header))
for key := range c.Request.Header {
value := strings.TrimSpace(c.Request.Header.Get(key))
if value == "" {
continue
}
headers[key] = value
}
if len(headers) == 0 {
return nil
}
return headers
}
func GenRelayInfo(c *gin.Context, relayFormat types.RelayFormat, request dto.Request, ws *websocket.Conn) (*RelayInfo, error) { func GenRelayInfo(c *gin.Context, relayFormat types.RelayFormat, request dto.Request, ws *websocket.Conn) (*RelayInfo, error) {
var info *RelayInfo var info *RelayInfo
var err error var err error
......
...@@ -172,9 +172,9 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types ...@@ -172,9 +172,9 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types
// apply param override // apply param override
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info)) jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info)
if err != nil { if err != nil {
return types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid, types.ErrOptionWithSkipRetry()) return newAPIErrorFromParamOverride(err)
} }
} }
......
...@@ -2,7 +2,6 @@ package relay ...@@ -2,7 +2,6 @@ package relay
import ( import (
"bytes" "bytes"
"encoding/json"
"fmt" "fmt"
"net/http" "net/http"
...@@ -46,15 +45,15 @@ func EmbeddingHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError * ...@@ -46,15 +45,15 @@ func EmbeddingHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
} }
relaycommon.AppendRequestConversionFromRequest(info, convertedRequest) relaycommon.AppendRequestConversionFromRequest(info, convertedRequest)
jsonData, err := json.Marshal(convertedRequest) jsonData, err := common.Marshal(convertedRequest)
if err != nil { if err != nil {
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
} }
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info)) jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info)
if err != nil { if err != nil {
return types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid, types.ErrOptionWithSkipRetry()) return newAPIErrorFromParamOverride(err)
} }
} }
......
...@@ -157,9 +157,9 @@ func GeminiHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ ...@@ -157,9 +157,9 @@ func GeminiHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
// apply param override // apply param override
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info)) jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info)
if err != nil { if err != nil {
return types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid, types.ErrOptionWithSkipRetry()) return newAPIErrorFromParamOverride(err)
} }
} }
...@@ -257,14 +257,9 @@ func GeminiEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo) (newAPI ...@@ -257,14 +257,9 @@ func GeminiEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo) (newAPI
// apply param override // apply param override
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
reqMap := make(map[string]interface{}) jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info)
_ = common.Unmarshal(jsonData, &reqMap)
for key, value := range info.ParamOverride {
reqMap[key] = value
}
jsonData, err = common.Marshal(reqMap)
if err != nil { if err != nil {
return types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid, types.ErrOptionWithSkipRetry()) return newAPIErrorFromParamOverride(err)
} }
} }
logger.LogDebug(c, "Gemini embedding request body: "+string(jsonData)) logger.LogDebug(c, "Gemini embedding request body: "+string(jsonData))
......
...@@ -70,9 +70,9 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type ...@@ -70,9 +70,9 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type
// apply param override // apply param override
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info)) jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info)
if err != nil { if err != nil {
return types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid, types.ErrOptionWithSkipRetry()) return newAPIErrorFromParamOverride(err)
} }
} }
......
package relay
import (
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/types"
)
func newAPIErrorFromParamOverride(err error) *types.NewAPIError {
if fixedErr, ok := relaycommon.AsParamOverrideReturnError(err); ok {
return relaycommon.NewAPIErrorFromParamOverride(fixedErr)
}
return types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid, types.ErrOptionWithSkipRetry())
}
...@@ -61,9 +61,9 @@ func RerankHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ ...@@ -61,9 +61,9 @@ func RerankHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ
// apply param override // apply param override
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info)) jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info)
if err != nil { if err != nil {
return types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid, types.ErrOptionWithSkipRetry()) return newAPIErrorFromParamOverride(err)
} }
} }
......
...@@ -96,9 +96,9 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError * ...@@ -96,9 +96,9 @@ func ResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *
// apply param override // apply param override
if len(info.ParamOverride) > 0 { if len(info.ParamOverride) > 0 {
jsonData, err = relaycommon.ApplyParamOverride(jsonData, info.ParamOverride, relaycommon.BuildParamOverrideContext(info)) jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info)
if err != nil { if err != nil {
return types.NewError(err, types.ErrorCodeChannelParamOverrideInvalid, types.ErrOptionWithSkipRetry()) return newAPIErrorFromParamOverride(err)
} }
} }
......
...@@ -45,6 +45,7 @@ type channelAffinityMeta struct { ...@@ -45,6 +45,7 @@ type channelAffinityMeta struct {
TTLSeconds int TTLSeconds int
RuleName string RuleName string
SkipRetry bool SkipRetry bool
ParamTemplate map[string]interface{}
KeySourceType string KeySourceType string
KeySourceKey string KeySourceKey string
KeySourcePath string KeySourcePath string
...@@ -415,6 +416,84 @@ func buildChannelAffinityKeyHint(s string) string { ...@@ -415,6 +416,84 @@ func buildChannelAffinityKeyHint(s string) string {
return s[:4] + "..." + s[len(s)-4:] return s[:4] + "..." + s[len(s)-4:]
} }
func cloneStringAnyMap(src map[string]interface{}) map[string]interface{} {
if len(src) == 0 {
return map[string]interface{}{}
}
dst := make(map[string]interface{}, len(src))
for k, v := range src {
dst[k] = v
}
return dst
}
func mergeChannelOverride(base map[string]interface{}, tpl map[string]interface{}) map[string]interface{} {
if len(base) == 0 && len(tpl) == 0 {
return map[string]interface{}{}
}
if len(tpl) == 0 {
return base
}
out := cloneStringAnyMap(base)
for k, v := range tpl {
out[k] = v
}
return out
}
func appendChannelAffinityTemplateAdminInfo(c *gin.Context, meta channelAffinityMeta) {
if c == nil {
return
}
if len(meta.ParamTemplate) == 0 {
return
}
templateInfo := map[string]interface{}{
"applied": true,
"rule_name": meta.RuleName,
"param_override_keys": len(meta.ParamTemplate),
}
if anyInfo, ok := c.Get(ginKeyChannelAffinityLogInfo); ok {
if info, ok := anyInfo.(map[string]interface{}); ok {
info["override_template"] = templateInfo
c.Set(ginKeyChannelAffinityLogInfo, info)
return
}
}
c.Set(ginKeyChannelAffinityLogInfo, map[string]interface{}{
"reason": meta.RuleName,
"rule_name": meta.RuleName,
"using_group": meta.UsingGroup,
"model": meta.ModelName,
"request_path": meta.RequestPath,
"key_source": meta.KeySourceType,
"key_key": meta.KeySourceKey,
"key_path": meta.KeySourcePath,
"key_hint": meta.KeyHint,
"key_fp": meta.KeyFingerprint,
"override_template": templateInfo,
})
}
// ApplyChannelAffinityOverrideTemplate merges per-rule channel override templates onto the selected channel override config.
func ApplyChannelAffinityOverrideTemplate(c *gin.Context, paramOverride map[string]interface{}) (map[string]interface{}, bool) {
if c == nil {
return paramOverride, false
}
meta, ok := getChannelAffinityMeta(c)
if !ok {
return paramOverride, false
}
if len(meta.ParamTemplate) == 0 {
return paramOverride, false
}
mergedParam := mergeChannelOverride(paramOverride, meta.ParamTemplate)
appendChannelAffinityTemplateAdminInfo(c, meta)
return mergedParam, true
}
func GetPreferredChannelByAffinity(c *gin.Context, modelName string, usingGroup string) (int, bool) { func GetPreferredChannelByAffinity(c *gin.Context, modelName string, usingGroup string) (int, bool) {
setting := operation_setting.GetChannelAffinitySetting() setting := operation_setting.GetChannelAffinitySetting()
if setting == nil || !setting.Enabled { if setting == nil || !setting.Enabled {
...@@ -466,6 +545,7 @@ func GetPreferredChannelByAffinity(c *gin.Context, modelName string, usingGroup ...@@ -466,6 +545,7 @@ func GetPreferredChannelByAffinity(c *gin.Context, modelName string, usingGroup
TTLSeconds: ttlSeconds, TTLSeconds: ttlSeconds,
RuleName: rule.Name, RuleName: rule.Name,
SkipRetry: rule.SkipRetryOnFailure, SkipRetry: rule.SkipRetryOnFailure,
ParamTemplate: cloneStringAnyMap(rule.ParamOverrideTemplate),
KeySourceType: strings.TrimSpace(usedSource.Type), KeySourceType: strings.TrimSpace(usedSource.Type),
KeySourceKey: strings.TrimSpace(usedSource.Key), KeySourceKey: strings.TrimSpace(usedSource.Key),
KeySourcePath: strings.TrimSpace(usedSource.Path), KeySourcePath: strings.TrimSpace(usedSource.Path),
......
package service
import (
"fmt"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/setting/operation_setting"
"github.com/gin-gonic/gin"
"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]interface{}{
"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]interface{}{
"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]interface{}{
"temperature": 0.7,
"max_tokens": 2000,
}
merged, applied := ApplyChannelAffinityOverrideTemplate(ctx, base)
require.True(t, applied)
require.Equal(t, 0.2, 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]interface{})
require.True(t, ok)
overrideInfoAny, ok := info["override_template"]
require.True(t, ok)
overrideInfo, ok := overrideInfoAny.(map[string]interface{})
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 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, "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]interface{}{
"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]interface{}{
"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)
}
...@@ -18,6 +18,8 @@ type ChannelAffinityRule struct { ...@@ -18,6 +18,8 @@ type ChannelAffinityRule struct {
ValueRegex string `json:"value_regex"` ValueRegex string `json:"value_regex"`
TTLSeconds int `json:"ttl_seconds"` TTLSeconds int `json:"ttl_seconds"`
ParamOverrideTemplate map[string]interface{} `json:"param_override_template,omitempty"`
SkipRetryOnFailure bool `json:"skip_retry_on_failure,omitempty"` SkipRetryOnFailure bool `json:"skip_retry_on_failure,omitempty"`
IncludeUsingGroup bool `json:"include_using_group"` IncludeUsingGroup bool `json:"include_using_group"`
...@@ -32,6 +34,44 @@ type ChannelAffinitySetting struct { ...@@ -32,6 +34,44 @@ type ChannelAffinitySetting struct {
Rules []ChannelAffinityRule `json:"rules"` Rules []ChannelAffinityRule `json:"rules"`
} }
var codexCliPassThroughHeaders = []string{
"Originator",
"Session_id",
"User-Agent",
"X-Codex-Beta-Features",
"X-Codex-Turn-Metadata",
}
var claudeCliPassThroughHeaders = []string{
"X-Stainless-Arch",
"X-Stainless-Lang",
"X-Stainless-Os",
"X-Stainless-Package-Version",
"X-Stainless-Retry-Count",
"X-Stainless-Runtime",
"X-Stainless-Runtime-Version",
"X-Stainless-Timeout",
"User-Agent",
"X-App",
"Anthropic-Beta",
"Anthropic-Dangerous-Direct-Browser-Access",
"Anthropic-Version",
}
func buildPassHeaderTemplate(headers []string) map[string]interface{} {
clonedHeaders := make([]string, 0, len(headers))
clonedHeaders = append(clonedHeaders, headers...)
return map[string]interface{}{
"operations": []map[string]interface{}{
{
"mode": "pass_headers",
"value": clonedHeaders,
"keep_origin": true,
},
},
}
}
var channelAffinitySetting = ChannelAffinitySetting{ var channelAffinitySetting = ChannelAffinitySetting{
Enabled: true, Enabled: true,
SwitchOnSuccess: true, SwitchOnSuccess: true,
...@@ -39,32 +79,34 @@ var channelAffinitySetting = ChannelAffinitySetting{ ...@@ -39,32 +79,34 @@ var channelAffinitySetting = ChannelAffinitySetting{
DefaultTTLSeconds: 3600, DefaultTTLSeconds: 3600,
Rules: []ChannelAffinityRule{ Rules: []ChannelAffinityRule{
{ {
Name: "codex trace", Name: "codex cli trace",
ModelRegex: []string{"^gpt-.*$"}, ModelRegex: []string{"^gpt-.*$"},
PathRegex: []string{"/v1/responses"}, PathRegex: []string{"/v1/responses"},
KeySources: []ChannelAffinityKeySource{ KeySources: []ChannelAffinityKeySource{
{Type: "gjson", Path: "prompt_cache_key"}, {Type: "gjson", Path: "prompt_cache_key"},
}, },
ValueRegex: "", ValueRegex: "",
TTLSeconds: 0, TTLSeconds: 0,
SkipRetryOnFailure: false, ParamOverrideTemplate: buildPassHeaderTemplate(codexCliPassThroughHeaders),
IncludeUsingGroup: true, SkipRetryOnFailure: false,
IncludeRuleName: true, IncludeUsingGroup: true,
UserAgentInclude: nil, IncludeRuleName: true,
UserAgentInclude: nil,
}, },
{ {
Name: "claude code trace", Name: "claude cli trace",
ModelRegex: []string{"^claude-.*$"}, ModelRegex: []string{"^claude-.*$"},
PathRegex: []string{"/v1/messages"}, PathRegex: []string{"/v1/messages"},
KeySources: []ChannelAffinityKeySource{ KeySources: []ChannelAffinityKeySource{
{Type: "gjson", Path: "metadata.user_id"}, {Type: "gjson", Path: "metadata.user_id"},
}, },
ValueRegex: "", ValueRegex: "",
TTLSeconds: 0, TTLSeconds: 0,
SkipRetryOnFailure: false, ParamOverrideTemplate: buildPassHeaderTemplate(claudeCliPassThroughHeaders),
IncludeUsingGroup: true, SkipRetryOnFailure: false,
IncludeRuleName: true, IncludeUsingGroup: true,
UserAgentInclude: nil, IncludeRuleName: true,
UserAgentInclude: nil,
}, },
}, },
} }
......
...@@ -10,7 +10,7 @@ ...@@ -10,7 +10,7 @@
"@visactor/react-vchart": "~1.8.8", "@visactor/react-vchart": "~1.8.8",
"@visactor/vchart": "~1.8.8", "@visactor/vchart": "~1.8.8",
"@visactor/vchart-semi-theme": "~1.8.8", "@visactor/vchart-semi-theme": "~1.8.8",
"axios": "1.13.5", "axios": "1.12.0",
"clsx": "^2.1.1", "clsx": "^2.1.1",
"dayjs": "^1.11.11", "dayjs": "^1.11.11",
"history": "^5.3.0", "history": "^5.3.0",
......
This source diff could not be displayed because it is too large. You can view the blob instead.
/*
Copyright (C) 2025 QuantumNous
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as
published by the Free Software Foundation, either version 3 of the
License, or (at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU Affero General Public License for more details.
You should have received a copy of the GNU Affero General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
For commercial licensing, please contact support@quantumnous.com
*/
const buildPassHeadersTemplate = (headers) => ({
operations: [
{
mode: 'pass_headers',
value: [...headers],
keep_origin: true,
},
],
});
export const CODEX_CLI_HEADER_PASSTHROUGH_HEADERS = [
'Originator',
'Session_id',
'User-Agent',
'X-Codex-Beta-Features',
'X-Codex-Turn-Metadata',
];
export const CLAUDE_CLI_HEADER_PASSTHROUGH_HEADERS = [
'X-Stainless-Arch',
'X-Stainless-Lang',
'X-Stainless-Os',
'X-Stainless-Package-Version',
'X-Stainless-Retry-Count',
'X-Stainless-Runtime',
'X-Stainless-Runtime-Version',
'X-Stainless-Timeout',
'User-Agent',
'X-App',
'Anthropic-Beta',
'Anthropic-Dangerous-Direct-Browser-Access',
'Anthropic-Version',
];
export const CODEX_CLI_HEADER_PASSTHROUGH_TEMPLATE = buildPassHeadersTemplate(
CODEX_CLI_HEADER_PASSTHROUGH_HEADERS,
);
export const CLAUDE_CLI_HEADER_PASSTHROUGH_TEMPLATE = buildPassHeadersTemplate(
CLAUDE_CLI_HEADER_PASSTHROUGH_HEADERS,
);
export const CHANNEL_AFFINITY_RULE_TEMPLATES = {
codexCli: {
name: 'codex cli trace',
model_regex: ['^gpt-.*$'],
path_regex: ['/v1/responses'],
key_sources: [{ type: 'gjson', path: 'prompt_cache_key' }],
param_override_template: CODEX_CLI_HEADER_PASSTHROUGH_TEMPLATE,
value_regex: '',
ttl_seconds: 0,
skip_retry_on_failure: false,
include_using_group: true,
include_rule_name: true,
},
claudeCli: {
name: 'claude cli trace',
model_regex: ['^claude-.*$'],
path_regex: ['/v1/messages'],
key_sources: [{ type: 'gjson', path: 'metadata.user_id' }],
param_override_template: CLAUDE_CLI_HEADER_PASSTHROUGH_TEMPLATE,
value_regex: '',
ttl_seconds: 0,
skip_retry_on_failure: false,
include_using_group: true,
include_rule_name: true,
},
};
export const cloneChannelAffinityTemplate = (template) =>
JSON.parse(JSON.stringify(template || {}));
...@@ -24,3 +24,4 @@ export * from './common.constant'; ...@@ -24,3 +24,4 @@ export * from './common.constant';
export * from './dashboard.constants'; export * from './dashboard.constants';
export * from './playground.constants'; export * from './playground.constants';
export * from './redemption.constants'; export * from './redemption.constants';
export * from './channel-affinity-template.constants';
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment