fix(models): enforce account mappings for remote Codex discovery

This commit is contained in:
lyy0709
2026-10-01 03:30:56 -07:00
parent d6adebd22d
commit 71e97d1353
5 changed files with 80 additions and 7 deletions
@@ -1001,3 +1001,32 @@ func TestCodexModelsPinnedAccountsETagMatchReturns304(t *testing.T) {
require.Equal(t, http.StatusNotModified, second.Code, second.Body.String())
require.Empty(t, second.Body.Bytes())
}
func TestCodexModelsDiscoveryProjectsWildcardMappingWithoutAllowlist(t *testing.T) {
for _, mode := range []string{"ordinary", "pinned", "fallback"} {
t.Run(mode, func(t *testing.T) {
account := newPinnedCodexAccount(1, service.StatusActive, true, false)
account.Credentials["model_mapping"] = map[string]any{"gpt-*": "gpt-5.4"}
upstream := &codexModelsPinnedHTTPUpstream{bodies: map[int64]string{
1: `{"models":[{"slug":"gpt-5.4"},{"slug":"blocked"},{"slug":"codex-auto-review"},{"slug":"gpt-image-2"}]}`,
}}
handler := newPinnedCodexTestHandler([]service.Account{account}, upstream, 3)
group := &service.Group{ID: 84, Platform: service.PlatformOpenAI,
ModelAllowlist: service.GroupModelAllowlist{Models: []string{"blocked"}},
}
if mode != "ordinary" {
group.CodexModelsManifestConfig = service.GroupCodexModelsManifestConfig{Enabled: true, AccountIDs: []int64{1}}
}
if mode == "fallback" {
group.CodexModelsManifestConfig.AccountIDs = []int64{999}
group.CodexModelsManifestConfig.FallbackToScheduler = true
}
first := performCodexModelsRequestForGroup(t, handler, group, "")
require.Equal(t, http.StatusOK, first.Code, first.Body.String())
require.Equal(t, []string{"gpt-5.4"}, codexHandlerManifestSlugs(t, first))
require.NotEmpty(t, first.Header().Get("ETag"))
second := performCodexModelsRequestForGroup(t, handler, group, first.Header().Get("ETag"))
require.Equal(t, http.StatusNotModified, second.Code)
})
}
}
@@ -8,6 +8,7 @@ import (
"fmt"
"log/slog"
"net/http"
"slices"
"strings"
"sync"
"time"
@@ -91,6 +92,13 @@ func mergeCodexModelsManifestBodies(bodies [][]byte) ([]byte, error) {
}
}
if len(bodies) == 1 {
_, original, err := modelCatalogEntries(bodies[0], "models")
if err == nil && slices.EqualFunc(original, mergedModels, func(a, b json.RawMessage) bool { return bytes.Equal(a, b) }) {
return bodies[0], nil
}
}
encodedModels, err := json.Marshal(mergedModels)
if err != nil {
return nil, fmt.Errorf("encode merged codex models: %w", err)
@@ -139,6 +147,9 @@ func (s *OpenAIGatewayService) FetchPinnedCodexModelsManifest(ctx context.Contex
if err != nil {
return nil, nil, fmt.Errorf("merge pinned codex models manifests: %w", err)
}
if len(results) == 1 && bytes.Equal(merged, results[0].response.Body) {
return results[0].response, results[0].account, nil
}
return &OpenAIModelsResponse{Body: merged, ETag: codexModelsManifestBodyETag(merged)}, results[0].account, nil
}
@@ -63,11 +63,11 @@ func TestMergeCodexModelsManifestBodiesHandlesSluglessEntries(t *testing.T) {
}
func TestMergeCodexModelsManifestBodiesSingleBody(t *testing.T) {
body := `{"models":[{"slug":"model-a"}]}`
body := `{ "unknown": true, "models": [ { "slug": "model-a" } ] }`
merged, err := mergeCodexModelsManifestBodies([][]byte{[]byte(body)})
require.NoError(t, err)
require.JSONEq(t, body, string(merged))
require.Equal(t, body, string(merged))
}
func TestMergeCodexModelsManifestBodiesRejectsInvalidInput(t *testing.T) {
@@ -80,3 +80,9 @@ func TestMergeCodexModelsManifestBodiesRejectsInvalidInput(t *testing.T) {
_, err = mergeCodexModelsManifestBodies([][]byte{[]byte(`{}`), []byte(`{"models":{}`)})
require.Error(t, err, "models 非数组必须报错")
}
func TestMergeCodexModelsManifestBodiesSingleBodyDeduplicates(t *testing.T) {
merged, err := mergeCodexModelsManifestBodies([][]byte{[]byte(`{"models":[{"slug":"a"},{"slug":"a","extra":true}]}`)})
require.NoError(t, err)
require.JSONEq(t, `{"models":[{"slug":"a"}]}`, string(merged))
}
+16 -5
View File
@@ -6,6 +6,7 @@ import (
"encoding/json"
"fmt"
"net/http"
"slices"
"sort"
"strings"
@@ -229,6 +230,10 @@ func projectAccountModelsBody(body []byte, account *Account, group *Group, codex
continue
}
seen[id] = struct{}{}
if id == target {
projected = append(projected, raw)
continue
}
var entry map[string]json.RawMessage
if err := json.Unmarshal(raw, &entry); err != nil {
return nil, err
@@ -243,6 +248,9 @@ func projectAccountModelsBody(body []byte, account *Account, group *Group, codex
}
projected = append(projected, encoded)
}
if slices.EqualFunc(entries, projected, func(a, b json.RawMessage) bool { return bytes.Equal(a, b) }) {
return body, nil
}
envelope[field], err = json.Marshal(projected)
if err != nil {
return nil, err
@@ -250,18 +258,21 @@ func projectAccountModelsBody(body []byte, account *Account, group *Group, codex
return json.Marshal(envelope)
}
// ApplyPinnedCodexModelsMapping is used by pinned discovery and its scheduler
// fallback. The ordinary (non-pinned) Codex path retains its local catalog policy.
// ApplyPinnedCodexModelsMapping projects all remotely discovered Codex catalogs,
// including ordinary discovery and pinned scheduler fallback. Locally generated
// catalogs retain their existing policy and do not pass through this function.
func ApplyPinnedCodexModelsMapping(response *OpenAIModelsResponse, account *Account, group *Group) error {
if group == nil || group.Platform != PlatformOpenAI || !group.CodexModelsManifestConfig.Enabled {
if group == nil || group.Platform != PlatformOpenAI {
return nil
}
body, err := projectAccountModelsBody(response.Body, account, group, true)
if err != nil {
return err
}
response.Body = body
response.ETag = codexModelsManifestBodyETag(body)
if !bytes.Equal(response.Body, body) {
response.Body = body
response.ETag = codexModelsManifestBodyETag(body)
}
return nil
}
@@ -3,6 +3,7 @@ package service
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"strings"
@@ -346,3 +347,18 @@ func TestOpenAIModelsCacheSeparatesRepresentationsForIdenticalRequests(t *testin
require.JSONEq(t, manifestBody, string(manifest.Body))
require.EqualValues(t, 2, calls.Load())
}
func TestApplyCodexModelsMappingPreservesUnchangedBodyAndETag(t *testing.T) {
for _, pinned := range []bool{false, true} {
t.Run(fmt.Sprint(pinned), func(t *testing.T) {
account := newCodexModelsAPIKeyTestAccount("https://models.example/v1")
account.Credentials["model_mapping"] = map[string]any{"gpt-5.4": "gpt-5.4"}
body := []byte(`{ "unknown": { "keep": true }, "models": [ { "slug": "gpt-5.4", "custom": 123 } ] }`)
response := &OpenAIModelsResponse{Body: body, ETag: `W/"upstream-validator"`}
group := &Group{Platform: PlatformOpenAI, CodexModelsManifestConfig: GroupCodexModelsManifestConfig{Enabled: pinned}}
require.NoError(t, ApplyPinnedCodexModelsMapping(response, account, group))
require.Equal(t, body, response.Body)
require.Equal(t, `W/"upstream-validator"`, response.ETag)
})
}
}