account_handler_available_models_test.go 16 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528
  1. package admin
  2. import (
  3. "context"
  4. "encoding/json"
  5. "io"
  6. "net/http"
  7. "net/http/httptest"
  8. "strings"
  9. "testing"
  10. "github.com/Wei-Shaw/sub2api/internal/config"
  11. "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
  12. "github.com/Wei-Shaw/sub2api/internal/service"
  13. "github.com/gin-gonic/gin"
  14. "github.com/stretchr/testify/require"
  15. )
  16. type availableModelsAdminService struct {
  17. *stubAdminService
  18. account service.Account
  19. }
  20. func (s *availableModelsAdminService) GetAccount(_ context.Context, id int64) (*service.Account, error) {
  21. if s.account.ID == id {
  22. acc := s.account
  23. return &acc, nil
  24. }
  25. return s.stubAdminService.GetAccount(context.Background(), id)
  26. }
  27. func setupAvailableModelsRouter(adminSvc service.AdminService) *gin.Engine {
  28. gin.SetMode(gin.TestMode)
  29. router := gin.New()
  30. handler := NewAccountHandler(adminSvc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
  31. router.GET("/api/v1/admin/accounts/:id/models", handler.GetAvailableModels)
  32. return router
  33. }
  34. type syncUpstreamHTTPUpstream struct {
  35. resp *http.Response
  36. responses []*http.Response
  37. err error
  38. }
  39. func (u *syncUpstreamHTTPUpstream) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) {
  40. if u.err != nil {
  41. return nil, u.err
  42. }
  43. if len(u.responses) > 0 {
  44. resp := u.responses[0]
  45. u.responses = u.responses[1:]
  46. return resp, nil
  47. }
  48. return u.resp, nil
  49. }
  50. func (u *syncUpstreamHTTPUpstream) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, profile *tlsfingerprint.Profile) (*http.Response, error) {
  51. return u.Do(req, proxyURL, accountID, accountConcurrency)
  52. }
  53. func setupSyncUpstreamModelsRouter(adminSvc service.AdminService, upstream service.HTTPUpstream) *gin.Engine {
  54. gin.SetMode(gin.TestMode)
  55. router := gin.New()
  56. accountTestSvc := service.NewAccountTestService(
  57. nil,
  58. nil,
  59. nil,
  60. nil,
  61. nil,
  62. upstream,
  63. &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{Enabled: false}}},
  64. nil,
  65. )
  66. handler := NewAccountHandler(adminSvc, nil, nil, nil, nil, nil, nil, nil, accountTestSvc, nil, nil, nil, nil, nil)
  67. router.POST("/api/v1/admin/accounts/:id/models/sync-upstream", handler.SyncUpstreamModels)
  68. router.POST("/api/v1/admin/accounts/models/sync-upstream-preview", handler.SyncUpstreamModelsPreview)
  69. return router
  70. }
  71. func TestAccountHandlerGetAvailableModels_GrokUsesXAIModels(t *testing.T) {
  72. svc := &availableModelsAdminService{
  73. stubAdminService: newStubAdminService(),
  74. account: service.Account{
  75. ID: 44,
  76. Name: "grok-oauth",
  77. Platform: service.PlatformGrok,
  78. Type: service.AccountTypeOAuth,
  79. Status: service.StatusActive,
  80. Credentials: map[string]any{
  81. "model_mapping": map[string]any{
  82. "grok-4.3": "grok-4.3",
  83. },
  84. },
  85. },
  86. }
  87. router := setupAvailableModelsRouter(svc)
  88. rec := httptest.NewRecorder()
  89. req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts/44/models", nil)
  90. router.ServeHTTP(rec, req)
  91. require.Equal(t, http.StatusOK, rec.Code)
  92. var resp struct {
  93. Data []struct {
  94. ID string `json:"id"`
  95. } `json:"data"`
  96. }
  97. require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
  98. require.Len(t, resp.Data, 1)
  99. require.Equal(t, "grok-4.3", resp.Data[0].ID)
  100. }
  101. func TestAccountHandlerGetAvailableModels_GrokDefaultsToXAIModelsWithoutMapping(t *testing.T) {
  102. svc := &availableModelsAdminService{
  103. stubAdminService: newStubAdminService(),
  104. account: service.Account{
  105. ID: 45,
  106. Name: "grok-oauth-defaults",
  107. Platform: service.PlatformGrok,
  108. Type: service.AccountTypeOAuth,
  109. Status: service.StatusActive,
  110. },
  111. }
  112. router := setupAvailableModelsRouter(svc)
  113. rec := httptest.NewRecorder()
  114. req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts/45/models", nil)
  115. router.ServeHTTP(rec, req)
  116. require.Equal(t, http.StatusOK, rec.Code)
  117. var resp struct {
  118. Data []struct {
  119. ID string `json:"id"`
  120. } `json:"data"`
  121. }
  122. require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
  123. require.NotEmpty(t, resp.Data)
  124. var ids []string
  125. for _, model := range resp.Data {
  126. id := model.ID
  127. ids = append(ids, id)
  128. require.NotContains(t, strings.ToLower(id), "claude")
  129. }
  130. require.Contains(t, ids, "grok-4.3")
  131. require.Contains(t, ids, "grok-build-0.1")
  132. }
  133. func TestAccountHandlerGetAvailableModels_OpenAIOAuthUsesExplicitModelMapping(t *testing.T) {
  134. svc := &availableModelsAdminService{
  135. stubAdminService: newStubAdminService(),
  136. account: service.Account{
  137. ID: 42,
  138. Name: "openai-oauth",
  139. Platform: service.PlatformOpenAI,
  140. Type: service.AccountTypeOAuth,
  141. Status: service.StatusActive,
  142. Credentials: map[string]any{
  143. "model_mapping": map[string]any{
  144. "gpt-5": "gpt-5.1",
  145. },
  146. },
  147. },
  148. }
  149. router := setupAvailableModelsRouter(svc)
  150. rec := httptest.NewRecorder()
  151. req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts/42/models", nil)
  152. router.ServeHTTP(rec, req)
  153. require.Equal(t, http.StatusOK, rec.Code)
  154. var resp struct {
  155. Data []struct {
  156. ID string `json:"id"`
  157. } `json:"data"`
  158. }
  159. require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
  160. require.Len(t, resp.Data, 1)
  161. require.Equal(t, "gpt-5", resp.Data[0].ID)
  162. }
  163. func TestAccountHandlerGetAvailableModels_OpenAIOAuthPassthroughFallsBackToDefaults(t *testing.T) {
  164. svc := &availableModelsAdminService{
  165. stubAdminService: newStubAdminService(),
  166. account: service.Account{
  167. ID: 43,
  168. Name: "openai-oauth-passthrough",
  169. Platform: service.PlatformOpenAI,
  170. Type: service.AccountTypeOAuth,
  171. Status: service.StatusActive,
  172. Credentials: map[string]any{
  173. "model_mapping": map[string]any{
  174. "gpt-5": "gpt-5.1",
  175. },
  176. },
  177. Extra: map[string]any{
  178. "openai_passthrough": true,
  179. },
  180. },
  181. }
  182. router := setupAvailableModelsRouter(svc)
  183. rec := httptest.NewRecorder()
  184. req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts/43/models", nil)
  185. router.ServeHTTP(rec, req)
  186. require.Equal(t, http.StatusOK, rec.Code)
  187. var resp struct {
  188. Data []struct {
  189. ID string `json:"id"`
  190. } `json:"data"`
  191. }
  192. require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
  193. require.NotEmpty(t, resp.Data)
  194. require.NotEqual(t, "gpt-5", resp.Data[0].ID)
  195. }
  196. func TestAccountHandlerGetAvailableModels_OpenAIAPIKeyDefaultsToConcreteGPT56Sol(t *testing.T) {
  197. svc := &availableModelsAdminService{
  198. stubAdminService: newStubAdminService(),
  199. account: service.Account{
  200. ID: 46,
  201. Name: "openai-apikey",
  202. Platform: service.PlatformOpenAI,
  203. Type: service.AccountTypeAPIKey,
  204. Status: service.StatusActive,
  205. Credentials: map[string]any{
  206. "api_key": "test-key",
  207. },
  208. },
  209. }
  210. router := setupAvailableModelsRouter(svc)
  211. rec := httptest.NewRecorder()
  212. req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts/46/models", nil)
  213. router.ServeHTTP(rec, req)
  214. require.Equal(t, http.StatusOK, rec.Code)
  215. var resp struct {
  216. Data []struct {
  217. ID string `json:"id"`
  218. } `json:"data"`
  219. }
  220. require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
  221. require.NotEmpty(t, resp.Data)
  222. require.Equal(t, "gpt-5.6-sol", resp.Data[0].ID)
  223. }
  224. func TestAccountHandlerGetAvailableModels_OpenAISparkShadowReturnsMappingModels(t *testing.T) {
  225. parentID := int64(100)
  226. svc := &availableModelsAdminService{
  227. stubAdminService: newStubAdminService(),
  228. account: service.Account{
  229. ID: 44,
  230. Name: "openai-spark-shadow",
  231. Platform: service.PlatformOpenAI,
  232. Type: service.AccountTypeOAuth,
  233. Status: service.StatusActive,
  234. ParentAccountID: &parentID,
  235. QuotaDimension: service.QuotaDimensionSpark,
  236. Credentials: map[string]any{
  237. "model_mapping": map[string]any{
  238. "gpt-5.3-codex-spark": "gpt-5.3-codex-spark",
  239. },
  240. },
  241. },
  242. }
  243. router := setupAvailableModelsRouter(svc)
  244. rec := httptest.NewRecorder()
  245. req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts/44/models", nil)
  246. router.ServeHTTP(rec, req)
  247. require.Equal(t, http.StatusOK, rec.Code)
  248. var resp struct {
  249. Data []struct {
  250. ID string `json:"id"`
  251. } `json:"data"`
  252. }
  253. require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
  254. ids := make([]string, 0, len(resp.Data))
  255. for _, m := range resp.Data {
  256. ids = append(ids, m.ID)
  257. }
  258. require.ElementsMatch(t, []string{
  259. "gpt-5.3-codex-spark",
  260. }, ids, "影子可用模型由 model_mapping 派生(非写死)")
  261. }
  262. func TestAccountHandlerGetAvailableModels_GeminiGoogleOneUsesConservativeCatalog(t *testing.T) {
  263. svc := &availableModelsAdminService{
  264. stubAdminService: newStubAdminService(),
  265. account: service.Account{
  266. ID: 45,
  267. Name: "google-one",
  268. Platform: service.PlatformGemini,
  269. Type: service.AccountTypeOAuth,
  270. Status: service.StatusActive,
  271. Credentials: map[string]any{
  272. "oauth_type": "google_one",
  273. },
  274. },
  275. }
  276. router := setupAvailableModelsRouter(svc)
  277. rec := httptest.NewRecorder()
  278. req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts/45/models", nil)
  279. router.ServeHTTP(rec, req)
  280. require.Equal(t, http.StatusOK, rec.Code)
  281. var resp struct {
  282. Data []struct {
  283. ID string `json:"id"`
  284. } `json:"data"`
  285. }
  286. require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
  287. ids := make([]string, 0, len(resp.Data))
  288. for _, model := range resp.Data {
  289. ids = append(ids, model.ID)
  290. }
  291. require.ElementsMatch(t, []string{"gemini-2.0-flash", "gemini-2.5-flash", "gemini-2.5-pro"}, ids)
  292. require.NotContains(t, ids, "gemini-3.5-flash")
  293. require.NotContains(t, ids, "gemini-2.5-flash-image")
  294. }
  295. func TestAccountHandlerSyncUpstreamModels_ConfigErrorReturnsBadRequest(t *testing.T) {
  296. svc := &availableModelsAdminService{
  297. stubAdminService: newStubAdminService(),
  298. account: service.Account{
  299. ID: 44,
  300. Name: "openai-apikey-missing-key",
  301. Platform: service.PlatformOpenAI,
  302. Type: service.AccountTypeAPIKey,
  303. Status: service.StatusActive,
  304. Credentials: map[string]any{
  305. "base_url": "https://openai.example.com/v1",
  306. },
  307. },
  308. }
  309. router := setupSyncUpstreamModelsRouter(svc, &syncUpstreamHTTPUpstream{})
  310. rec := httptest.NewRecorder()
  311. req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/44/models/sync-upstream", nil)
  312. router.ServeHTTP(rec, req)
  313. require.Equal(t, http.StatusBadRequest, rec.Code)
  314. require.Contains(t, rec.Body.String(), "No OpenAI API key is available")
  315. }
  316. func TestAccountHandlerSyncUpstreamModelsReturnsCapabilityMetadata(t *testing.T) {
  317. svc := &availableModelsAdminService{
  318. stubAdminService: newStubAdminService(),
  319. account: service.Account{
  320. ID: 48, Name: "custom-openai", Platform: service.PlatformOpenAI,
  321. Type: service.AccountTypeAPIKey, Status: service.StatusActive,
  322. Credentials: map[string]any{"api_key": "key", "base_url": "https://provider.example/v1"},
  323. },
  324. }
  325. upstream := &syncUpstreamHTTPUpstream{resp: &http.Response{
  326. StatusCode: http.StatusOK,
  327. Header: http.Header{"Content-Type": []string{"application/json"}},
  328. Body: io.NopCloser(strings.NewReader(`{"models":[{
  329. "id":"custom-thinking-model",
  330. "reasoning":true,
  331. "default_reasoning_level":"high",
  332. "supported_reasoning_levels":["low","high"],
  333. "input_modalities":["text","image"],
  334. "context_window":256000
  335. }]}`)),
  336. }}
  337. router := setupSyncUpstreamModelsRouter(svc, upstream)
  338. rec := httptest.NewRecorder()
  339. req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/48/models/sync-upstream", nil)
  340. router.ServeHTTP(rec, req)
  341. require.Equal(t, http.StatusOK, rec.Code)
  342. var resp struct {
  343. Data service.UpstreamModelCatalog `json:"data"`
  344. }
  345. require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
  346. require.Equal(t, []string{"custom-thinking-model"}, resp.Data.Models)
  347. metadata := resp.Data.Metadata["custom-thinking-model"]
  348. require.NotNil(t, metadata.Reasoning)
  349. require.True(t, *metadata.Reasoning)
  350. require.Equal(t, []string{"low", "high"}, metadata.SupportedReasoningLevels)
  351. require.Equal(t, []string{"text", "image"}, metadata.InputModalities)
  352. }
  353. // Scenario: 创建账号 preview 将具体 mapping 传给 404/405 配置回退。
  354. func TestAccountHandlerSyncUpstreamModelsPreviewUsesProvidedModelMapping(t *testing.T) {
  355. upstream := &syncUpstreamHTTPUpstream{responses: []*http.Response{
  356. {
  357. StatusCode: http.StatusNotFound,
  358. Header: http.Header{"Content-Type": []string{"application/json"}},
  359. Body: io.NopCloser(strings.NewReader(`{"error":"not found"}`)),
  360. },
  361. {
  362. StatusCode: http.StatusOK,
  363. Header: http.Header{"Content-Type": []string{"application/json"}},
  364. Body: io.NopCloser(strings.NewReader(`{
  365. "configured-provider": {
  366. "api": "https://provider.example/v1",
  367. "models": {
  368. "glm-5.3": {
  369. "id": "glm-5.3",
  370. "reasoning": true,
  371. "reasoning_options": [{"type":"effort","values":["low","high"]}],
  372. "modalities": {"input":["text"],"output":["text"]},
  373. "limit": {"context":1000000,"output":131072}
  374. }
  375. }
  376. }
  377. }`)),
  378. },
  379. }}
  380. router := setupSyncUpstreamModelsRouter(newStubAdminService(), upstream)
  381. rec := httptest.NewRecorder()
  382. req := httptest.NewRequest(
  383. http.MethodPost,
  384. "/api/v1/admin/accounts/models/sync-upstream-preview",
  385. strings.NewReader(`{
  386. "platform":"openai",
  387. "type":"apikey",
  388. "base_url":"https://provider.example/v1",
  389. "api_key":"key",
  390. "model_mapping":{"public-glm":"glm-5.3"}
  391. }`),
  392. )
  393. req.Header.Set("Content-Type", "application/json")
  394. router.ServeHTTP(rec, req)
  395. require.Equal(t, http.StatusOK, rec.Code)
  396. var resp struct {
  397. Data service.UpstreamModelCatalog `json:"data"`
  398. }
  399. require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
  400. require.Equal(t, []string{"glm-5.3"}, resp.Data.Models)
  401. require.Equal(t, []string{"low", "high"}, resp.Data.Metadata["glm-5.3"].SupportedReasoningLevels)
  402. }
  403. func TestAccountHandlerSyncUpstreamModels_UpstreamErrorDoesNotExposeBody(t *testing.T) {
  404. svc := &availableModelsAdminService{
  405. stubAdminService: newStubAdminService(),
  406. account: service.Account{
  407. ID: 45,
  408. Name: "openai-apikey-upstream-error",
  409. Platform: service.PlatformOpenAI,
  410. Type: service.AccountTypeAPIKey,
  411. Status: service.StatusActive,
  412. Credentials: map[string]any{
  413. "api_key": "openai-key",
  414. "base_url": "https://openai.example.com/v1",
  415. },
  416. },
  417. }
  418. upstream := &syncUpstreamHTTPUpstream{resp: &http.Response{
  419. StatusCode: http.StatusBadGateway,
  420. Header: http.Header{"Content-Type": []string{"application/json"}},
  421. Body: io.NopCloser(strings.NewReader(`{"error":"SECRET_TOKEN should not be exposed"}`)),
  422. }}
  423. router := setupSyncUpstreamModelsRouter(svc, upstream)
  424. rec := httptest.NewRecorder()
  425. req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/45/models/sync-upstream", nil)
  426. router.ServeHTTP(rec, req)
  427. require.Equal(t, http.StatusBadGateway, rec.Code)
  428. require.Contains(t, rec.Body.String(), "Upstream model list request failed with HTTP 502")
  429. require.NotContains(t, rec.Body.String(), "SECRET_TOKEN")
  430. }
  431. // Scenario: 能力补全失败显示部分成功。
  432. func TestAccountHandlerSyncUpstreamModels_MetadataEnrichmentFailureReturnsWarning(t *testing.T) {
  433. svc := &availableModelsAdminService{
  434. stubAdminService: newStubAdminService(),
  435. account: service.Account{
  436. ID: 46,
  437. Name: "opencode-id-only-model-list",
  438. Platform: service.PlatformOpenAI,
  439. Type: service.AccountTypeAPIKey,
  440. Status: service.StatusActive,
  441. Credentials: map[string]any{
  442. "api_key": "opencode-key",
  443. "base_url": "https://opencode.ai/zen/v1",
  444. },
  445. },
  446. }
  447. upstream := &syncUpstreamHTTPUpstream{responses: []*http.Response{
  448. {
  449. StatusCode: http.StatusOK,
  450. Header: http.Header{"Content-Type": []string{"application/json"}},
  451. Body: io.NopCloser(strings.NewReader(`{"data":[{"id":"x-preview-f-free"}]}`)),
  452. },
  453. {
  454. StatusCode: http.StatusBadGateway,
  455. Header: http.Header{"Content-Type": []string{"application/json"}},
  456. Body: io.NopCloser(strings.NewReader(`{"error":"registry unavailable"}`)),
  457. },
  458. }}
  459. router := setupSyncUpstreamModelsRouter(svc, upstream)
  460. rec := httptest.NewRecorder()
  461. req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/46/models/sync-upstream", nil)
  462. router.ServeHTTP(rec, req)
  463. require.Equal(t, http.StatusOK, rec.Code)
  464. var resp struct {
  465. Data struct {
  466. Models []string `json:"models"`
  467. Warnings []struct {
  468. Code string `json:"code"`
  469. } `json:"warnings"`
  470. } `json:"data"`
  471. }
  472. require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
  473. require.Equal(t, []string{"x-preview-f-free"}, resp.Data.Models)
  474. require.Len(t, resp.Data.Warnings, 1)
  475. require.Equal(t, "upstream_model_metadata_incomplete", resp.Data.Warnings[0].Code)
  476. }