openai_codex_models_handler_test.go 35 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003
  1. package handler
  2. import (
  3. "context"
  4. "encoding/json"
  5. "errors"
  6. "fmt"
  7. "io"
  8. "net"
  9. "net/http"
  10. "net/http/httptest"
  11. "sort"
  12. "strings"
  13. "sync"
  14. "testing"
  15. "time"
  16. "github.com/Wei-Shaw/sub2api/internal/config"
  17. middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
  18. "github.com/Wei-Shaw/sub2api/internal/service"
  19. "github.com/gin-gonic/gin"
  20. "github.com/stretchr/testify/require"
  21. )
  22. type codexModelsFailoverAccountRepo struct {
  23. service.AccountRepository
  24. accounts []service.Account
  25. }
  26. func (r codexModelsFailoverAccountRepo) GetByID(_ context.Context, id int64) (*service.Account, error) {
  27. for i := range r.accounts {
  28. if r.accounts[i].ID == id {
  29. account := r.accounts[i]
  30. return &account, nil
  31. }
  32. }
  33. return nil, service.ErrNoAvailableAccounts
  34. }
  35. func (r codexModelsFailoverAccountRepo) ListSchedulableByPlatform(_ context.Context, platform string) ([]service.Account, error) {
  36. accounts := make([]service.Account, 0, len(r.accounts))
  37. for _, account := range r.accounts {
  38. if account.Platform == platform {
  39. accounts = append(accounts, account)
  40. }
  41. }
  42. return accounts, nil
  43. }
  44. func (r codexModelsFailoverAccountRepo) ListSchedulableByGroupID(_ context.Context, _ int64) ([]service.Account, error) {
  45. return append([]service.Account(nil), r.accounts...), nil
  46. }
  47. func (r codexModelsFailoverAccountRepo) ListByGroup(_ context.Context, _ int64) ([]service.Account, error) {
  48. return append([]service.Account(nil), r.accounts...), nil
  49. }
  50. func (r codexModelsFailoverAccountRepo) ListModelAvailabilityCandidates(_ context.Context, _ *int64, _ []string, _ bool) ([]service.Account, error) {
  51. return append([]service.Account(nil), r.accounts...), nil
  52. }
  53. type codexModelsFailoverHTTPUpstream struct {
  54. service.HTTPUpstream
  55. mu sync.Mutex
  56. accountIDs []int64
  57. firstErr error
  58. firstStatus int
  59. firstBody string
  60. statuses map[int64]int
  61. }
  62. func (u *codexModelsFailoverHTTPUpstream) Do(_ *http.Request, _ string, accountID int64, _ int) (*http.Response, error) {
  63. u.mu.Lock()
  64. u.accountIDs = append(u.accountIDs, accountID)
  65. u.mu.Unlock()
  66. status, hasStatus := u.statuses[accountID]
  67. if accountID == 1 || hasStatus {
  68. if u.firstErr != nil {
  69. return nil, u.firstErr
  70. }
  71. if u.firstBody != "" && !hasStatus {
  72. return &http.Response{
  73. StatusCode: http.StatusOK,
  74. Status: "200 OK",
  75. Header: make(http.Header),
  76. Body: io.NopCloser(strings.NewReader(u.firstBody)),
  77. }, nil
  78. }
  79. if !hasStatus {
  80. status = u.firstStatus
  81. }
  82. if status == 0 {
  83. status = http.StatusServiceUnavailable
  84. }
  85. return &http.Response{
  86. StatusCode: status,
  87. Status: http.StatusText(status),
  88. Header: make(http.Header),
  89. Body: io.NopCloser(strings.NewReader(
  90. `{"error":{"message":"No available OpenAI accounts","type":"upstream_error"}}`,
  91. )),
  92. }, nil
  93. }
  94. return &http.Response{
  95. StatusCode: http.StatusOK,
  96. Status: "200 OK",
  97. Header: make(http.Header),
  98. Body: io.NopCloser(strings.NewReader(`{"models":[{"slug":"gpt-5.6-sol"}]}`)),
  99. }, nil
  100. }
  101. func (u *codexModelsFailoverHTTPUpstream) calls() []int64 {
  102. u.mu.Lock()
  103. defer u.mu.Unlock()
  104. return append([]int64(nil), u.accountIDs...)
  105. }
  106. func TestCodexModelsCanceledRequestDoesNotWriteResponse(t *testing.T) {
  107. gin.SetMode(gin.TestMode)
  108. recorder := httptest.NewRecorder()
  109. c, _ := gin.CreateTestContext(recorder)
  110. ctx, cancel := context.WithCancel(context.Background())
  111. cancel()
  112. c.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil).WithContext(ctx)
  113. h := &OpenAIGatewayHandler{}
  114. h.CodexModels(c)
  115. if c.Writer.Written() {
  116. t.Fatalf("canceled request wrote an HTTP response: status=%d body=%q", recorder.Code, recorder.Body.String())
  117. }
  118. }
  119. func TestCodexModelsAppliesLocalFiltersBeforeClientETag(t *testing.T) {
  120. gin.SetMode(gin.TestMode)
  121. groupID := int64(43)
  122. repo := &codexModelsFailoverAccountRepo{accounts: []service.Account{
  123. {
  124. ID: 1,
  125. Name: "custom-openai",
  126. Platform: service.PlatformOpenAI,
  127. Type: service.AccountTypeAPIKey,
  128. Status: service.StatusActive,
  129. Schedulable: true,
  130. Concurrency: 1,
  131. Credentials: map[string]any{
  132. "api_key": "sk-test",
  133. "base_url": "https://upstream.example/v1",
  134. },
  135. },
  136. }}
  137. upstream := &codexModelsFailoverHTTPUpstream{
  138. firstBody: `{"object":"list","data":[{"id":"codex-auto-review"},{"id":"gpt-5.6"}]}`,
  139. }
  140. gatewayService := service.NewOpenAIGatewayService(
  141. repo,
  142. nil, nil, nil, nil, nil, nil, &config.Config{RunMode: config.RunModeSimple}, nil, nil, nil, nil, nil,
  143. upstream,
  144. nil, nil, nil, nil, nil, nil, nil, nil,
  145. )
  146. handler := &OpenAIGatewayHandler{gatewayService: gatewayService}
  147. group := &service.Group{
  148. ID: groupID,
  149. Platform: service.PlatformOpenAI,
  150. ModelAllowlist: service.GroupModelAllowlist{
  151. Enabled: true,
  152. Models: []string{"codex-auto-review", "gpt-5.6"},
  153. },
  154. }
  155. first := performCodexModelsRequestForGroup(t, handler, group, "")
  156. if first.Code != http.StatusOK {
  157. t.Fatalf("first status: got %d, want %d; body=%s", first.Code, http.StatusOK, first.Body.String())
  158. }
  159. if body := first.Body.String(); !strings.Contains(body, "codex-auto-review") || !strings.Contains(body, "gpt-5.6") {
  160. t.Fatalf("first body did not include the explicitly selected models: %s", body)
  161. }
  162. oldETag := first.Header().Get("ETag")
  163. if oldETag == "" {
  164. t.Fatal("first response did not include an ETag")
  165. }
  166. group.ModelAllowlist.Enabled = false
  167. second := performCodexModelsRequestForGroup(t, handler, group, oldETag)
  168. if second.Code != http.StatusOK {
  169. t.Fatalf("second status: got %d, want %d; body=%s", second.Code, http.StatusOK, second.Body.String())
  170. }
  171. if body := second.Body.String(); strings.Contains(body, "codex-auto-review") || !strings.Contains(body, "gpt-5.6") {
  172. t.Fatalf("second body was not the filtered manifest: %s", body)
  173. }
  174. if newETag := second.Header().Get("ETag"); newETag == "" || newETag == oldETag {
  175. t.Fatalf("second ETag: got %q, want a new final-body ETag", newETag)
  176. }
  177. third := performCodexModelsRequestForGroup(t, handler, group, second.Header().Get("ETag"))
  178. if third.Code != http.StatusNotModified {
  179. t.Fatalf("third status: got %d, want %d; body=%s", third.Code, http.StatusNotModified, third.Body.String())
  180. }
  181. if third.Body.Len() != 0 {
  182. t.Fatalf("third body: got %q, want empty", third.Body.String())
  183. }
  184. }
  185. func TestCodexModelsAPIKeyCacheDoesNotLeakGroupFilters(t *testing.T) {
  186. gin.SetMode(gin.TestMode)
  187. repo := &codexModelsFailoverAccountRepo{accounts: []service.Account{
  188. {
  189. ID: 1,
  190. Name: "shared-api-key",
  191. Platform: service.PlatformOpenAI,
  192. Type: service.AccountTypeAPIKey,
  193. Status: service.StatusActive,
  194. Schedulable: true,
  195. Concurrency: 1,
  196. Credentials: map[string]any{
  197. "api_key": "sk-shared",
  198. "base_url": "https://upstream.example/v1",
  199. },
  200. },
  201. }}
  202. upstream := &codexModelsFailoverHTTPUpstream{
  203. firstBody: `{"object":"list","data":[{"id":"model-a"},{"id":"model-b"}]}`,
  204. }
  205. gatewayService := service.NewOpenAIGatewayService(
  206. repo,
  207. nil, nil, nil, nil, nil, nil, &config.Config{RunMode: config.RunModeSimple}, nil, nil, nil, nil, nil,
  208. upstream,
  209. nil, nil, nil, nil, nil, nil, nil, nil,
  210. )
  211. handler := &OpenAIGatewayHandler{gatewayService: gatewayService}
  212. groupA := &service.Group{
  213. ID: 91,
  214. Platform: service.PlatformOpenAI,
  215. ModelAllowlist: service.GroupModelAllowlist{
  216. Enabled: true,
  217. Models: []string{"model-a"},
  218. },
  219. }
  220. groupB := &service.Group{
  221. ID: 92,
  222. Platform: service.PlatformOpenAI,
  223. ModelAllowlist: service.GroupModelAllowlist{
  224. Enabled: true,
  225. Models: []string{"model-b"},
  226. },
  227. }
  228. firstA := performCodexModelsRequestForGroup(t, handler, groupA, "")
  229. require.Equal(t, http.StatusOK, firstA.Code, firstA.Body.String())
  230. require.Equal(t, []string{"model-a"}, codexHandlerManifestSlugs(t, firstA))
  231. firstB := performCodexModelsRequestForGroup(t, handler, groupB, "")
  232. require.Equal(t, http.StatusOK, firstB.Code, firstB.Body.String())
  233. require.Equal(t, []string{"model-b"}, codexHandlerManifestSlugs(t, firstB))
  234. etagA := firstA.Header().Get("ETag")
  235. require.NotEmpty(t, etagA)
  236. var wg sync.WaitGroup
  237. results := make([]*httptest.ResponseRecorder, 8)
  238. for i := range results {
  239. wg.Add(1)
  240. go func(index int) {
  241. defer wg.Done()
  242. if index%2 == 0 {
  243. results[index] = performCodexModelsRequestForGroup(t, handler, groupA, etagA)
  244. return
  245. }
  246. results[index] = performCodexModelsRequestForGroup(t, handler, groupB, "")
  247. }(i)
  248. }
  249. wg.Wait()
  250. sawGroupB := false
  251. for _, recorder := range results {
  252. require.NotNil(t, recorder)
  253. switch recorder.Code {
  254. case http.StatusNotModified:
  255. require.Empty(t, recorder.Body.Bytes())
  256. case http.StatusOK:
  257. slugs := codexHandlerManifestSlugs(t, recorder)
  258. if len(slugs) == 1 && slugs[0] == "model-b" {
  259. sawGroupB = true
  260. continue
  261. }
  262. require.Equal(t, []string{"model-a"}, slugs)
  263. default:
  264. t.Fatalf("unexpected status %d body=%s", recorder.Code, recorder.Body.String())
  265. }
  266. }
  267. require.True(t, sawGroupB)
  268. }
  269. func TestCodexModelsSupplementsConfiguredModelsWithUnmappedAccountDefaults(t *testing.T) {
  270. gin.SetMode(gin.TestMode)
  271. groupID := int64(44)
  272. repo := &codexModelsFailoverAccountRepo{accounts: []service.Account{
  273. {
  274. ID: 1,
  275. Name: "ark-compatible",
  276. Platform: service.PlatformOpenAI,
  277. Type: service.AccountTypeAPIKey,
  278. Status: service.StatusActive,
  279. Schedulable: true,
  280. Priority: 0,
  281. Concurrency: 1,
  282. Credentials: map[string]any{
  283. "api_key": "sk-ark",
  284. "base_url": "https://ark.example/v1",
  285. "model_mapping": map[string]any{
  286. "glm-5.3": "glm-5.3",
  287. },
  288. },
  289. },
  290. {
  291. ID: 2,
  292. Name: "chatgpt-oauth",
  293. Platform: service.PlatformOpenAI,
  294. Type: service.AccountTypeOAuth,
  295. Status: service.StatusActive,
  296. Schedulable: true,
  297. Priority: 1,
  298. Concurrency: 1,
  299. Credentials: map[string]any{
  300. "access_token": "oauth-test",
  301. },
  302. },
  303. }}
  304. upstream := &codexModelsFailoverHTTPUpstream{firstStatus: http.StatusNotFound}
  305. gatewayService := service.NewOpenAIGatewayService(
  306. repo,
  307. nil, nil, nil, nil, nil, nil, &config.Config{RunMode: config.RunModeSimple}, nil, nil, nil, nil, nil,
  308. upstream,
  309. nil, nil, nil, nil, nil, nil, nil, nil,
  310. )
  311. handler := &OpenAIGatewayHandler{gatewayService: gatewayService}
  312. recorder := performCodexModelsRequestForGroup(t, handler, &service.Group{
  313. ID: groupID,
  314. Platform: service.PlatformOpenAI,
  315. }, "")
  316. if got := upstream.calls(); len(got) != 0 {
  317. t.Fatalf("upstream account calls: got %v, want none", got)
  318. }
  319. if recorder.Code != http.StatusOK {
  320. t.Fatalf("status: got %d, want %d; body=%s", recorder.Code, http.StatusOK, recorder.Body.String())
  321. }
  322. var envelope struct {
  323. Models []map[string]any `json:"models"`
  324. }
  325. if err := json.Unmarshal(recorder.Body.Bytes(), &envelope); err != nil {
  326. t.Fatalf("decode body: %v; body=%s", err, recorder.Body.String())
  327. }
  328. require.Contains(t, codexHandlerManifestSlugs(t, recorder), "gpt-5.6-sol")
  329. require.Contains(t, codexHandlerManifestSlugs(t, recorder), "glm-5.3")
  330. for _, model := range envelope.Models {
  331. require.Contains(t, model, "supported_reasoning_levels")
  332. }
  333. }
  334. func TestCodexModelsUnmappedParentAndSparkShadowHonorCustomListAndETag(t *testing.T) {
  335. gin.SetMode(gin.TestMode)
  336. const sparkModel = "gpt-5.3-codex-spark"
  337. parentID := int64(1)
  338. repo := &codexModelsFailoverAccountRepo{accounts: []service.Account{
  339. {
  340. ID: parentID, Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth,
  341. Status: service.StatusActive, Schedulable: true,
  342. },
  343. {
  344. ID: 2, Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth,
  345. Status: service.StatusActive, Schedulable: true,
  346. ParentAccountID: &parentID, QuotaDimension: "spark",
  347. Credentials: map[string]any{"model_mapping": map[string]any{sparkModel: sparkModel}},
  348. },
  349. }}
  350. upstream := &codexModelsFailoverHTTPUpstream{firstStatus: http.StatusNotFound}
  351. gatewayService := service.NewOpenAIGatewayService(
  352. repo,
  353. nil, nil, nil, nil, nil, nil, &config.Config{RunMode: config.RunModeSimple}, nil, nil, nil, nil, nil,
  354. upstream,
  355. nil, nil, nil, nil, nil, nil, nil, nil,
  356. )
  357. handler := &OpenAIGatewayHandler{gatewayService: gatewayService}
  358. group := &service.Group{ID: 45, Platform: service.PlatformOpenAI}
  359. first := performCodexModelsRequestForGroup(t, handler, group, "")
  360. require.Equal(t, http.StatusOK, first.Code, first.Body.String())
  361. slugs := codexHandlerManifestSlugs(t, first)
  362. require.Contains(t, slugs, "gpt-5.6-sol")
  363. require.Contains(t, slugs, sparkModel)
  364. require.NotContains(t, slugs, "gpt-image-2")
  365. require.NotContains(t, slugs, "codex-auto-review")
  366. firstETag := first.Header().Get("ETag")
  367. require.NotEmpty(t, firstETag)
  368. group.ModelAllowlist = service.GroupModelAllowlist{
  369. Enabled: true, Models: []string{"gpt-5.6-sol", sparkModel, "unknown-model"},
  370. }
  371. second := performCodexModelsRequestForGroup(t, handler, group, firstETag)
  372. require.Equal(t, http.StatusOK, second.Code, second.Body.String())
  373. require.ElementsMatch(t, []string{"gpt-5.6-sol", sparkModel}, codexHandlerManifestSlugs(t, second))
  374. require.NotEqual(t, firstETag, second.Header().Get("ETag"))
  375. third := performCodexModelsRequestForGroup(t, handler, group, second.Header().Get("ETag"))
  376. require.Equal(t, http.StatusNotModified, third.Code)
  377. require.Empty(t, third.Body.Bytes())
  378. require.Empty(t, upstream.calls())
  379. }
  380. func TestCompositeCodexModelsReusesExistingManifestSelection(t *testing.T) {
  381. handler, upstream, groupID := newCodexModelsFailoverTestHandler(http.StatusServiceUnavailable)
  382. recorder := performCodexModelsRequestForPlatform(t, handler, groupID, service.PlatformComposite)
  383. if got, want := upstream.calls(), []int64{1, 2}; !equalInt64Slices(got, want) {
  384. t.Fatalf("upstream account calls: got %v, want %v", got, want)
  385. }
  386. if recorder.Code != http.StatusOK {
  387. t.Fatalf("status: got %d, want %d; body=%s", recorder.Code, http.StatusOK, recorder.Body.String())
  388. }
  389. }
  390. func TestCodexModelsFailsOverFromRetryableUpstreamStatus(t *testing.T) {
  391. retryableStatuses := []int{
  392. http.StatusTooManyRequests,
  393. http.StatusInternalServerError,
  394. http.StatusBadGateway,
  395. http.StatusServiceUnavailable,
  396. http.StatusGatewayTimeout,
  397. }
  398. for _, status := range retryableStatuses {
  399. t.Run(http.StatusText(status), func(t *testing.T) {
  400. handler, upstream, groupID := newCodexModelsFailoverTestHandler(status)
  401. recorder := performCodexModelsRequest(t, handler, groupID)
  402. if got, want := upstream.calls(), []int64{1, 2}; !equalInt64Slices(got, want) {
  403. t.Fatalf("upstream account calls: got %v, want %v", got, want)
  404. }
  405. if recorder.Code != http.StatusOK {
  406. t.Fatalf("status: got %d, want %d; body=%s", recorder.Code, http.StatusOK, recorder.Body.String())
  407. }
  408. requireCompleteCodexModelsHandlerResponse(t, recorder, "gpt-5.6-sol")
  409. })
  410. }
  411. }
  412. // Scenario: an API-key upstream without /models is excluded only for this discovery request.
  413. func TestCodexModelsFailsOverWhenAPIKeyModelsEndpointIsUnavailable(t *testing.T) {
  414. for _, status := range []int{http.StatusNotFound, http.StatusMethodNotAllowed} {
  415. t.Run(http.StatusText(status), func(t *testing.T) {
  416. handler, upstream, groupID := newCodexModelsFailoverTestHandler(status)
  417. recorder := performCodexModelsRequest(t, handler, groupID)
  418. if got, want := upstream.calls(), []int64{1, 2}; !equalInt64Slices(got, want) {
  419. t.Fatalf("upstream account calls: got %v, want %v", got, want)
  420. }
  421. if recorder.Code != http.StatusOK {
  422. t.Fatalf("status: got %d, want %d; body=%s", recorder.Code, http.StatusOK, recorder.Body.String())
  423. }
  424. })
  425. }
  426. }
  427. func TestCodexModelsFailsOverFromUpstreamTransportError(t *testing.T) {
  428. handler, upstream, groupID := newCodexModelsFailoverTestHandler(http.StatusServiceUnavailable)
  429. upstream.firstErr = &net.OpError{
  430. Op: "read",
  431. Net: "tcp",
  432. Err: errors.New("connection reset"),
  433. }
  434. recorder := performCodexModelsRequest(t, handler, groupID)
  435. if got, want := upstream.calls(), []int64{1, 2}; !equalInt64Slices(got, want) {
  436. t.Fatalf("upstream account calls: got %v, want %v", got, want)
  437. }
  438. if recorder.Code != http.StatusOK {
  439. t.Fatalf("status: got %d, want %d; body=%s", recorder.Code, http.StatusOK, recorder.Body.String())
  440. }
  441. }
  442. func TestCodexModelsFailsOverFromInvalidManifestEnvelope(t *testing.T) {
  443. handler, upstream, groupID := newCodexModelsFailoverTestHandler(http.StatusOK)
  444. upstream.firstBody = `{"object":"list","data":[]}`
  445. recorder := performCodexModelsRequest(t, handler, groupID)
  446. if got, want := upstream.calls(), []int64{1, 2}; !equalInt64Slices(got, want) {
  447. t.Fatalf("upstream account calls: got %v, want %v", got, want)
  448. }
  449. if recorder.Code != http.StatusOK {
  450. t.Fatalf("status: got %d, want %d; body=%s", recorder.Code, http.StatusOK, recorder.Body.String())
  451. }
  452. requireCompleteCodexModelsHandlerResponse(t, recorder, "gpt-5.6-sol")
  453. }
  454. func TestCodexModelsDoesNotFailOverFromPermanentUpstreamStatus(t *testing.T) {
  455. statuses := []int{
  456. http.StatusBadRequest,
  457. http.StatusUnauthorized,
  458. http.StatusForbidden,
  459. 600,
  460. }
  461. for _, status := range statuses {
  462. t.Run(fmt.Sprintf("status_%d", status), func(t *testing.T) {
  463. handler, upstream, groupID := newCodexModelsFailoverTestHandler(status)
  464. recorder := performCodexModelsRequest(t, handler, groupID)
  465. if got, want := upstream.calls(), []int64{1}; !equalInt64Slices(got, want) {
  466. t.Fatalf("upstream account calls: got %v, want %v", got, want)
  467. }
  468. if recorder.Code != http.StatusBadGateway {
  469. t.Fatalf("status: got %d, want %d; body=%s", recorder.Code, http.StatusBadGateway, recorder.Body.String())
  470. }
  471. })
  472. }
  473. }
  474. func TestCodexModelsDoesNotFailOverFromUpstreamConfigurationError(t *testing.T) {
  475. handler, upstream, groupID := newCodexModelsFailoverTestHandler(http.StatusServiceUnavailable)
  476. upstream.firstErr = errors.New("invalid proxy URL")
  477. recorder := performCodexModelsRequest(t, handler, groupID)
  478. if got, want := upstream.calls(), []int64{1}; !equalInt64Slices(got, want) {
  479. t.Fatalf("upstream account calls: got %v, want %v", got, want)
  480. }
  481. if recorder.Code != http.StatusBadGateway {
  482. t.Fatalf("status: got %d, want %d; body=%s", recorder.Code, http.StatusBadGateway, recorder.Body.String())
  483. }
  484. }
  485. func TestCodexModelsReturnsLastUpstreamErrorWhenAccountsAreExhausted(t *testing.T) {
  486. handler, upstream, groupID := newCodexModelsFailoverTestHandler(http.StatusServiceUnavailable)
  487. upstream.statuses = map[int64]int{
  488. 1: http.StatusServiceUnavailable,
  489. 2: http.StatusGatewayTimeout,
  490. }
  491. recorder := performCodexModelsRequest(t, handler, groupID)
  492. if got, want := upstream.calls(), []int64{1, 2}; !equalInt64Slices(got, want) {
  493. t.Fatalf("upstream account calls: got %v, want %v", got, want)
  494. }
  495. if recorder.Code != http.StatusBadGateway {
  496. t.Fatalf("status: got %d, want %d; body=%s", recorder.Code, http.StatusBadGateway, recorder.Body.String())
  497. }
  498. if body := recorder.Body.String(); !strings.Contains(body, "upstream error 504") {
  499. t.Fatalf("body does not preserve the last upstream error: %s", body)
  500. }
  501. }
  502. func TestCodexModelsHonorsAccountSwitchLimit(t *testing.T) {
  503. handler, upstream, groupID := newCodexModelsFailoverTestHandlerWithAccountCount(http.StatusServiceUnavailable, 4, 2)
  504. upstream.statuses = map[int64]int{
  505. 1: http.StatusServiceUnavailable,
  506. 2: http.StatusBadGateway,
  507. 3: http.StatusGatewayTimeout,
  508. 4: http.StatusInternalServerError,
  509. }
  510. recorder := performCodexModelsRequest(t, handler, groupID)
  511. if got, want := upstream.calls(), []int64{1, 2, 3}; !equalInt64Slices(got, want) {
  512. t.Fatalf("upstream account calls: got %v, want %v", got, want)
  513. }
  514. if recorder.Code != http.StatusBadGateway {
  515. t.Fatalf("status: got %d, want %d; body=%s", recorder.Code, http.StatusBadGateway, recorder.Body.String())
  516. }
  517. if body := recorder.Body.String(); !strings.Contains(body, "upstream error 504") {
  518. t.Fatalf("body does not preserve the limit-ending upstream error: %s", body)
  519. }
  520. }
  521. func newCodexModelsFailoverTestHandler(firstStatus int) (*OpenAIGatewayHandler, *codexModelsFailoverHTTPUpstream, int64) {
  522. return newCodexModelsFailoverTestHandlerWithAccountCount(firstStatus, 2, 3)
  523. }
  524. func newCodexModelsFailoverTestHandlerWithAccountCount(firstStatus, accountCount, maxSwitches int) (*OpenAIGatewayHandler, *codexModelsFailoverHTTPUpstream, int64) {
  525. gin.SetMode(gin.TestMode)
  526. groupID := int64(42)
  527. accounts := make([]service.Account, 0, accountCount)
  528. for i := 1; i <= accountCount; i++ {
  529. accounts = append(accounts, service.Account{
  530. ID: int64(i),
  531. Name: fmt.Sprintf("upstream-%d", i),
  532. Platform: service.PlatformOpenAI,
  533. Type: service.AccountTypeAPIKey,
  534. Status: service.StatusActive,
  535. Schedulable: true,
  536. Priority: i - 1,
  537. Concurrency: 1,
  538. Credentials: map[string]any{
  539. "api_key": fmt.Sprintf("sk-%d", i),
  540. "base_url": fmt.Sprintf("https://upstream-%d.example/v1", i),
  541. },
  542. })
  543. }
  544. upstream := &codexModelsFailoverHTTPUpstream{firstStatus: firstStatus}
  545. cfg := &config.Config{RunMode: config.RunModeSimple}
  546. gatewayService := service.NewOpenAIGatewayService(
  547. codexModelsFailoverAccountRepo{accounts: accounts},
  548. nil, nil, nil, nil, nil, nil, cfg, nil, nil, nil, nil, nil,
  549. upstream,
  550. nil, nil, nil, nil, nil, nil, nil, nil,
  551. )
  552. return &OpenAIGatewayHandler{gatewayService: gatewayService, maxAccountSwitches: maxSwitches}, upstream, groupID
  553. }
  554. func performCodexModelsRequest(t *testing.T, handler *OpenAIGatewayHandler, groupID int64) *httptest.ResponseRecorder {
  555. return performCodexModelsRequestForGroup(t, handler, &service.Group{ID: groupID, Platform: service.PlatformOpenAI}, "")
  556. }
  557. func performCodexModelsRequestForPlatform(t *testing.T, handler *OpenAIGatewayHandler, groupID int64, platform string) *httptest.ResponseRecorder {
  558. return performCodexModelsRequestForGroup(t, handler, &service.Group{ID: groupID, Platform: platform}, "")
  559. }
  560. func performCodexModelsRequestForGroup(t *testing.T, handler *OpenAIGatewayHandler, group *service.Group, etag string) *httptest.ResponseRecorder {
  561. t.Helper()
  562. recorder := httptest.NewRecorder()
  563. c, _ := gin.CreateTestContext(recorder)
  564. c.Request = httptest.NewRequest(http.MethodGet, "/v1/models?client_version=0.144.0", nil)
  565. if etag != "" {
  566. c.Request.Header.Set("If-None-Match", etag)
  567. }
  568. c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
  569. GroupID: &group.ID,
  570. Group: group,
  571. })
  572. handler.CodexModels(c)
  573. return recorder
  574. }
  575. func codexHandlerManifestSlugs(t *testing.T, recorder *httptest.ResponseRecorder) []string {
  576. t.Helper()
  577. var envelope struct {
  578. Models []struct {
  579. Slug string `json:"slug"`
  580. } `json:"models"`
  581. }
  582. if err := json.Unmarshal(recorder.Body.Bytes(), &envelope); err != nil {
  583. t.Fatalf("decode body: %v; body=%s", err, recorder.Body.String())
  584. }
  585. slugs := make([]string, 0, len(envelope.Models))
  586. for _, model := range envelope.Models {
  587. slugs = append(slugs, model.Slug)
  588. }
  589. return slugs
  590. }
  591. func requireCompleteCodexModelsHandlerResponse(t *testing.T, recorder *httptest.ResponseRecorder, slug string) {
  592. t.Helper()
  593. var envelope struct {
  594. Models []map[string]any `json:"models"`
  595. }
  596. if err := json.Unmarshal(recorder.Body.Bytes(), &envelope); err != nil {
  597. t.Fatalf("decode body: %v; body=%s", err, recorder.Body.String())
  598. }
  599. if len(envelope.Models) != 1 {
  600. t.Fatalf("models count: got %d, want 1; body=%s", len(envelope.Models), recorder.Body.String())
  601. }
  602. model := envelope.Models[0]
  603. if got := model["slug"]; got != slug {
  604. t.Fatalf("slug: got %v, want %q", got, slug)
  605. }
  606. if levels, ok := model["supported_reasoning_levels"].([]any); !ok || len(levels) == 0 {
  607. t.Fatalf("supported_reasoning_levels must be populated: %v", model["supported_reasoning_levels"])
  608. }
  609. if messages, ok := model["model_messages"].(map[string]any); !ok || messages["instructions_template"] == "" {
  610. t.Fatalf("model_messages.instructions_template must be populated: %v", model["model_messages"])
  611. }
  612. if policy, ok := model["truncation_policy"].(map[string]any); !ok || len(policy) == 0 {
  613. t.Fatalf("truncation_policy must be populated: %v", model["truncation_policy"])
  614. }
  615. modalities, ok := model["input_modalities"].([]any)
  616. if !ok || len(modalities) != 2 || modalities[0] != "text" || modalities[1] != "image" {
  617. t.Fatalf("known GPT model modalities: got %v, want [text image]", model["input_modalities"])
  618. }
  619. }
  620. func equalInt64Slices(got, want []int64) bool {
  621. if len(got) != len(want) {
  622. return false
  623. }
  624. for i := range got {
  625. if got[i] != want[i] {
  626. return false
  627. }
  628. }
  629. return true
  630. }
  631. // --- 固定账号 manifest 模式 ---
  632. // codexModelsPinnedHTTPUpstream 按账号返回不同 manifest/状态的测试上游。
  633. type codexModelsPinnedHTTPUpstream struct {
  634. service.HTTPUpstream
  635. mu sync.Mutex
  636. calls []int64
  637. bodies map[int64]string
  638. statuses map[int64]int
  639. }
  640. func (u *codexModelsPinnedHTTPUpstream) Do(_ *http.Request, _ string, accountID int64, _ int) (*http.Response, error) {
  641. u.mu.Lock()
  642. u.calls = append(u.calls, accountID)
  643. u.mu.Unlock()
  644. if status, ok := u.statuses[accountID]; ok {
  645. return &http.Response{
  646. StatusCode: status,
  647. Status: http.StatusText(status),
  648. Header: make(http.Header),
  649. Body: io.NopCloser(strings.NewReader(`{"error":{"message":"upstream boom"}}`)),
  650. }, nil
  651. }
  652. body, ok := u.bodies[accountID]
  653. if !ok {
  654. body = `{"models":[]}`
  655. }
  656. return &http.Response{
  657. StatusCode: http.StatusOK,
  658. Status: "200 OK",
  659. Header: make(http.Header),
  660. Body: io.NopCloser(strings.NewReader(body)),
  661. }, nil
  662. }
  663. func (u *codexModelsPinnedHTTPUpstream) accountIDs() []int64 {
  664. u.mu.Lock()
  665. defer u.mu.Unlock()
  666. ids := append([]int64(nil), u.calls...)
  667. sort.Slice(ids, func(i, j int) bool { return ids[i] < ids[j] })
  668. return ids
  669. }
  670. func newPinnedCodexAccount(id int64, status string, schedulable bool, rateLimited bool) service.Account {
  671. account := service.Account{
  672. ID: id,
  673. Name: fmt.Sprintf("pinned-%d", id),
  674. Platform: service.PlatformOpenAI,
  675. Type: service.AccountTypeAPIKey,
  676. Status: status,
  677. Schedulable: schedulable,
  678. Priority: int(id),
  679. Concurrency: 1,
  680. Credentials: map[string]any{
  681. "api_key": fmt.Sprintf("sk-pinned-%d", id),
  682. "base_url": fmt.Sprintf("https://pinned-%d.example/v1", id),
  683. },
  684. }
  685. if rateLimited {
  686. reset := time.Now().Add(10 * time.Minute)
  687. account.RateLimitResetAt = &reset
  688. }
  689. return account
  690. }
  691. func newPinnedCodexTestHandler(accounts []service.Account, upstream *codexModelsPinnedHTTPUpstream, maxSwitches int) *OpenAIGatewayHandler {
  692. gin.SetMode(gin.TestMode)
  693. cfg := &config.Config{RunMode: config.RunModeSimple}
  694. gatewayService := service.NewOpenAIGatewayService(
  695. codexModelsFailoverAccountRepo{accounts: accounts},
  696. nil, nil, nil, nil, nil, nil, cfg, nil, nil, nil, nil, nil,
  697. upstream,
  698. nil, nil, nil, nil, nil, nil, nil, nil,
  699. )
  700. return &OpenAIGatewayHandler{gatewayService: gatewayService, maxAccountSwitches: maxSwitches}
  701. }
  702. func performPinnedCodexModelsRequest(t *testing.T, handler *OpenAIGatewayHandler, group *service.Group, etag string) *httptest.ResponseRecorder {
  703. return performCodexModelsRequestForGroup(t, handler, group, etag)
  704. }
  705. func TestCodexModelsPinnedAccountsMergeUnionWithoutScheduler(t *testing.T) {
  706. // 账号 1 优先级最高(调度器会先选它);固定配置只用 2、3,
  707. // 返回并集且不打账号 1,即可证明没有经过调度器。
  708. accounts := []service.Account{
  709. newPinnedCodexAccount(1, service.StatusActive, true, false),
  710. newPinnedCodexAccount(2, service.StatusActive, true, false),
  711. newPinnedCodexAccount(3, service.StatusActive, true, false),
  712. }
  713. upstream := &codexModelsPinnedHTTPUpstream{bodies: map[int64]string{
  714. 1: `{"models":[{"slug":"scheduler-only"}]}`,
  715. 2: `{"models":[{"slug":"model-a"}]}`,
  716. 3: `{"models":[{"slug":"model-a"},{"slug":"model-c"}]}`,
  717. }}
  718. handler := newPinnedCodexTestHandler(accounts, upstream, 3)
  719. group := &service.Group{
  720. ID: 77,
  721. Platform: service.PlatformOpenAI,
  722. CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{
  723. Enabled: true,
  724. AccountIDs: []int64{2, 3},
  725. },
  726. }
  727. recorder := performPinnedCodexModelsRequest(t, handler, group, "")
  728. require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
  729. require.Equal(t, []string{"model-a", "model-c"}, codexHandlerManifestSlugs(t, recorder))
  730. require.Equal(t, []int64{2, 3}, upstream.accountIDs(), "固定账号模式不得调用调度器或非固定账号")
  731. }
  732. func TestCodexModelsPinnedAccountsUseRateLimitedAccountAndSkipUnavailable(t *testing.T) {
  733. accounts := []service.Account{
  734. newPinnedCodexAccount(2, service.StatusActive, true, true), // 限流中:仍被使用
  735. newPinnedCodexAccount(3, service.StatusActive, false, false), // 调度开关关闭:跳过
  736. newPinnedCodexAccount(4, service.StatusDisabled, true, false), // 停用:跳过
  737. newPinnedCodexAccount(5, service.StatusActive, true, false), // 正常
  738. }
  739. upstream := &codexModelsPinnedHTTPUpstream{bodies: map[int64]string{
  740. 2: `{"models":[{"slug":"from-rate-limited"}]}`,
  741. 5: `{"models":[{"slug":"model-five"}]}`,
  742. }}
  743. handler := newPinnedCodexTestHandler(accounts, upstream, 3)
  744. group := &service.Group{
  745. ID: 78,
  746. Platform: service.PlatformOpenAI,
  747. CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{
  748. Enabled: true,
  749. AccountIDs: []int64{2, 3, 4, 5, 99}, // 99 不在分组:跳过
  750. },
  751. }
  752. recorder := performPinnedCodexModelsRequest(t, handler, group, "")
  753. require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
  754. require.Equal(t, []string{"from-rate-limited", "model-five"}, codexHandlerManifestSlugs(t, recorder))
  755. require.Equal(t, []int64{2, 5}, upstream.accountIDs())
  756. }
  757. func TestCodexModelsPinnedAccountsPartialFailureStillSucceeds(t *testing.T) {
  758. accounts := []service.Account{
  759. newPinnedCodexAccount(2, service.StatusActive, true, false),
  760. newPinnedCodexAccount(3, service.StatusActive, true, false),
  761. }
  762. upstream := &codexModelsPinnedHTTPUpstream{
  763. bodies: map[int64]string{2: `{"models":[{"slug":"model-a"}]}`},
  764. statuses: map[int64]int{3: http.StatusServiceUnavailable},
  765. }
  766. handler := newPinnedCodexTestHandler(accounts, upstream, 3)
  767. group := &service.Group{
  768. ID: 79,
  769. Platform: service.PlatformOpenAI,
  770. CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{
  771. Enabled: true,
  772. AccountIDs: []int64{2, 3},
  773. },
  774. }
  775. recorder := performPinnedCodexModelsRequest(t, handler, group, "")
  776. require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
  777. require.Equal(t, []string{"model-a"}, codexHandlerManifestSlugs(t, recorder))
  778. }
  779. func TestCodexModelsPinnedAccountsAllUnavailableReturns503ByDefault(t *testing.T) {
  780. accounts := []service.Account{
  781. newPinnedCodexAccount(2, service.StatusActive, false, false),
  782. }
  783. upstream := &codexModelsPinnedHTTPUpstream{}
  784. handler := newPinnedCodexTestHandler(accounts, upstream, 3)
  785. group := &service.Group{
  786. ID: 80,
  787. Platform: service.PlatformOpenAI,
  788. CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{
  789. Enabled: true,
  790. AccountIDs: []int64{2},
  791. },
  792. }
  793. recorder := performPinnedCodexModelsRequest(t, handler, group, "")
  794. require.Equal(t, http.StatusServiceUnavailable, recorder.Code, recorder.Body.String())
  795. require.Empty(t, upstream.accountIDs())
  796. }
  797. func TestCodexModelsPinnedAccountsAllFailedReturnsUpstreamError(t *testing.T) {
  798. accounts := []service.Account{
  799. newPinnedCodexAccount(2, service.StatusActive, true, false),
  800. newPinnedCodexAccount(3, service.StatusActive, true, false),
  801. }
  802. upstream := &codexModelsPinnedHTTPUpstream{statuses: map[int64]int{
  803. 2: http.StatusServiceUnavailable,
  804. 3: http.StatusGatewayTimeout,
  805. }}
  806. handler := newPinnedCodexTestHandler(accounts, upstream, 3)
  807. group := &service.Group{
  808. ID: 81,
  809. Platform: service.PlatformOpenAI,
  810. CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{
  811. Enabled: true,
  812. AccountIDs: []int64{2, 3},
  813. },
  814. }
  815. recorder := performPinnedCodexModelsRequest(t, handler, group, "")
  816. require.Equal(t, http.StatusBadGateway, recorder.Code, recorder.Body.String())
  817. require.Contains(t, recorder.Body.String(), "upstream error 504", "全部失败时返回最后一个上游错误")
  818. }
  819. func TestCodexModelsPinnedAccountsFallbackToScheduler(t *testing.T) {
  820. // 固定账号 2 不可用且回退开启:跌入调度器,选中优先级最高的账号 1。
  821. accounts := []service.Account{
  822. newPinnedCodexAccount(1, service.StatusActive, true, false),
  823. newPinnedCodexAccount(2, service.StatusActive, false, false),
  824. }
  825. upstream := &codexModelsPinnedHTTPUpstream{bodies: map[int64]string{
  826. 1: `{"models":[{"slug":"from-scheduler"}]}`,
  827. }}
  828. handler := newPinnedCodexTestHandler(accounts, upstream, 3)
  829. group := &service.Group{
  830. ID: 82,
  831. Platform: service.PlatformOpenAI,
  832. CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{
  833. Enabled: true,
  834. AccountIDs: []int64{2},
  835. FallbackToScheduler: true,
  836. },
  837. }
  838. recorder := performPinnedCodexModelsRequest(t, handler, group, "")
  839. require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
  840. require.Equal(t, []string{"from-scheduler"}, codexHandlerManifestSlugs(t, recorder))
  841. require.Equal(t, []int64{1}, upstream.accountIDs())
  842. }
  843. func TestCodexModelsPinnedAccountsFallbackToSchedulerOnAllFailed(t *testing.T) {
  844. accounts := []service.Account{
  845. newPinnedCodexAccount(1, service.StatusActive, true, false),
  846. newPinnedCodexAccount(2, service.StatusActive, true, false),
  847. }
  848. upstream := &codexModelsPinnedHTTPUpstream{
  849. bodies: map[int64]string{1: `{"models":[{"slug":"from-scheduler"}]}`},
  850. statuses: map[int64]int{2: http.StatusServiceUnavailable},
  851. }
  852. handler := newPinnedCodexTestHandler(accounts, upstream, 3)
  853. group := &service.Group{
  854. ID: 83,
  855. Platform: service.PlatformOpenAI,
  856. CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{
  857. Enabled: true,
  858. AccountIDs: []int64{2},
  859. FallbackToScheduler: true,
  860. },
  861. }
  862. recorder := performPinnedCodexModelsRequest(t, handler, group, "")
  863. require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
  864. require.Equal(t, []string{"from-scheduler"}, codexHandlerManifestSlugs(t, recorder))
  865. }
  866. func TestCodexModelsPinnedAccountsStillApplyCustomModelsListFilter(t *testing.T) {
  867. accounts := []service.Account{
  868. newPinnedCodexAccount(2, service.StatusActive, true, false),
  869. newPinnedCodexAccount(3, service.StatusActive, true, false),
  870. }
  871. upstream := &codexModelsPinnedHTTPUpstream{bodies: map[int64]string{
  872. 2: `{"models":[{"slug":"model-a"}]}`,
  873. 3: `{"models":[{"slug":"model-b"}]}`,
  874. }}
  875. handler := newPinnedCodexTestHandler(accounts, upstream, 3)
  876. group := &service.Group{
  877. ID: 84,
  878. Platform: service.PlatformOpenAI,
  879. ModelAllowlist: service.GroupModelAllowlist{
  880. Enabled: true,
  881. Models: []string{"model-b"},
  882. },
  883. CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{
  884. Enabled: true,
  885. AccountIDs: []int64{2, 3},
  886. },
  887. }
  888. recorder := performPinnedCodexModelsRequest(t, handler, group, "")
  889. require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
  890. require.Equal(t, []string{"model-b"}, codexHandlerManifestSlugs(t, recorder), "分组自定义模型列表过滤仍生效")
  891. }
  892. func TestCodexModelsPinnedAccountsETagMatchReturns304(t *testing.T) {
  893. accounts := []service.Account{
  894. newPinnedCodexAccount(2, service.StatusActive, true, false),
  895. newPinnedCodexAccount(3, service.StatusActive, true, false),
  896. }
  897. upstream := &codexModelsPinnedHTTPUpstream{bodies: map[int64]string{
  898. 2: `{"models":[{"slug":"model-a"}]}`,
  899. 3: `{"models":[{"slug":"model-c"}]}`,
  900. }}
  901. handler := newPinnedCodexTestHandler(accounts, upstream, 3)
  902. group := &service.Group{
  903. ID: 85,
  904. Platform: service.PlatformOpenAI,
  905. CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{
  906. Enabled: true,
  907. AccountIDs: []int64{2, 3},
  908. },
  909. }
  910. first := performPinnedCodexModelsRequest(t, handler, group, "")
  911. require.Equal(t, http.StatusOK, first.Code, first.Body.String())
  912. etag := first.Header().Get("ETag")
  913. require.NotEmpty(t, etag)
  914. second := performPinnedCodexModelsRequest(t, handler, group, etag)
  915. require.Equal(t, http.StatusNotModified, second.Code, second.Body.String())
  916. require.Empty(t, second.Body.Bytes())
  917. }