github/new-api
1
0
Fork 0
mirror of https://github.com/QuantumNous/new-api.git synced 2026-03-30 00:44:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat：support $keep_only_declared and deduped $append for header token overrides

Browse Source
This commit is contained in:
Seefs
2026-03-09 00:12:53 +08:00
parent 9bb2b6a6ae
commit 6ff5a5dc99
3 changed files with 178 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

43
relay/common/override.go
View File

@@ -847,24 +847,30 @@ func resolveHeaderOverrideValueByMapping(context map[string]interface{}, headerN
return "", false, fmt.Errorf("header value mapping cannot be empty")
}
sourceValue, exists := getHeaderValueFromContext(context, headerName)
if !exists {
return "", false, nil
appendTokens, err := parseHeaderAppendTokens(mapping)
if err != nil {
return "", false, err
}
sourceTokens := splitHeaderListValue(sourceValue)
if len(sourceTokens) == 0 {
return "", false, nil
keepOnlyDeclared := parseHeaderKeepOnlyDeclared(mapping)
sourceValue, exists := getHeaderValueFromContext(context, headerName)
sourceTokens := make([]string, 0)
if exists {
sourceTokens = splitHeaderListValue(sourceValue)
}
wildcardValue, hasWildcard := mapping["*"]
resultTokens := make([]string, 0, len(sourceTokens))
resultTokens := make([]string, 0, len(sourceTokens)+len(appendTokens))
for _, token := range sourceTokens {
replacementRaw, hasReplacement := mapping[token]
if !hasReplacement && hasWildcard {
if !hasReplacement && hasWildcard && !keepOnlyDeclared {
replacementRaw = wildcardValue
hasReplacement = true
}
if !hasReplacement {
if keepOnlyDeclared {
continue
}
resultTokens = append(resultTokens, token)
continue
}
@@ -875,6 +881,7 @@ func resolveHeaderOverrideValueByMapping(context map[string]interface{}, headerN
resultTokens = append(resultTokens, replacementTokens...)
}
resultTokens = append(resultTokens, appendTokens...)
resultTokens = lo.Uniq(resultTokens)
if len(resultTokens) == 0 {
return "", false, nil
@@ -882,6 +889,26 @@ func resolveHeaderOverrideValueByMapping(context map[string]interface{}, headerN
return strings.Join(resultTokens, ","), true, nil
}
func parseHeaderAppendTokens(mapping map[string]interface{}) ([]string, error) {
appendRaw, ok := mapping["$append"]
if !ok {
return nil, nil
}
return parseHeaderReplacementTokens(appendRaw)
}
func parseHeaderKeepOnlyDeclared(mapping map[string]interface{}) bool {
keepOnlyDeclaredRaw, ok := mapping["$keep_only_declared"]
if !ok {
return false
}
keepOnlyDeclared, ok := keepOnlyDeclaredRaw.(bool)
if !ok {
return false
}
return keepOnlyDeclared
}
func parseHeaderReplacementTokens(value interface{}) ([]string, error) {
switch raw := value.(type) {
case nil:

135
relay/common/override_test.go
View File

@@ -1653,6 +1653,141 @@ func TestApplyParamOverrideSetHeaderMapDeleteWholeHeaderWhenAllTokensCleared(t *
}
}
func TestApplyParamOverrideSetHeaderMapAppendsTokens(t *testing.T) {
input := []byte(`{"temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"mode": "set_header",
"path": "anthropic-beta",
"value": map[string]interface{}{
"$append": []interface{}{"context-1m-2025-08-07", "computer-use-2025-01-24"},
},
},
},
}
ctx := map[string]interface{}{
"header_override": map[string]interface{}{
"anthropic-beta": "computer-use-2025-01-24",
},
}
out, err := ApplyParamOverride(input, override, ctx)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"temperature":0.7}`, string(out))
headers, ok := ctx["header_override"].(map[string]interface{})
if !ok {
t.Fatalf("expected header_override context map")
}
if headers["anthropic-beta"] != "computer-use-2025-01-24,context-1m-2025-08-07" {
t.Fatalf("expected anthropic-beta to append new token without duplicates, got: %v", headers["anthropic-beta"])
}
}
func TestApplyParamOverrideSetHeaderMapAppendsTokensWhenHeaderMissing(t *testing.T) {
input := []byte(`{"temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"mode": "set_header",
"path": "anthropic-beta",
"value": map[string]interface{}{
"$append": []interface{}{"context-1m-2025-08-07", "computer-use-2025-01-24"},
},
},
},
}
ctx := map[string]interface{}{}
out, err := ApplyParamOverride(input, override, ctx)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"temperature":0.7}`, string(out))
headers, ok := ctx["header_override"].(map[string]interface{})
if !ok {
t.Fatalf("expected header_override context map")
}
if headers["anthropic-beta"] != "context-1m-2025-08-07,computer-use-2025-01-24" {
t.Fatalf("expected anthropic-beta to be created from appended tokens, got: %v", headers["anthropic-beta"])
}
}
func TestApplyParamOverrideSetHeaderMapKeepOnlyDeclaredDropsUndeclaredTokens(t *testing.T) {
input := []byte(`{"temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"mode": "set_header",
"path": "anthropic-beta",
"value": map[string]interface{}{
"computer-use-2025-01-24": "computer-use-2025-01-24",
"$append": []interface{}{"context-1m-2025-08-07"},
"$keep_only_declared": true,
},
},
},
}
ctx := map[string]interface{}{
"header_override": map[string]interface{}{
"anthropic-beta": "advanced-tool-use-2025-11-20,computer-use-2025-01-24",
},
}
out, err := ApplyParamOverride(input, override, ctx)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"temperature":0.7}`, string(out))
headers, ok := ctx["header_override"].(map[string]interface{})
if !ok {
t.Fatalf("expected header_override context map")
}
if headers["anthropic-beta"] != "computer-use-2025-01-24,context-1m-2025-08-07" {
t.Fatalf("expected anthropic-beta to keep only declared tokens, got: %v", headers["anthropic-beta"])
}
}
func TestApplyParamOverrideSetHeaderMapKeepOnlyDeclaredDeletesHeaderWhenNothingDeclaredMatches(t *testing.T) {
input := []byte(`{"temperature":0.7}`)
override := map[string]interface{}{
"operations": []interface{}{
map[string]interface{}{
"mode": "set_header",
"path": "anthropic-beta",
"value": map[string]interface{}{
"computer-use-2025-01-24": "computer-use-2025-01-24",
"$keep_only_declared": true,
},
},
},
}
ctx := map[string]interface{}{
"header_override": map[string]interface{}{
"anthropic-beta": "advanced-tool-use-2025-11-20",
},
}
out, err := ApplyParamOverride(input, override, ctx)
if err != nil {
t.Fatalf("ApplyParamOverride returned error: %v", err)
}
assertJSONEqual(t, `{"temperature":0.7}`, string(out))
headers, ok := ctx["header_override"].(map[string]interface{})
if !ok {
t.Fatalf("expected header_override context map")
}
if _, exists := headers["anthropic-beta"]; exists {
t.Fatalf("expected anthropic-beta to be deleted when no declared tokens remain, got: %v", headers["anthropic-beta"])
}
}
func TestApplyParamOverrideConditionsObjectShorthand(t *testing.T) {
input := []byte(`{"temperature":0.7}`)
override := map[string]interface{}{