openai_alpha_search_test.go 27 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656
  1. package service
  2. import (
  3. "bytes"
  4. "context"
  5. "io"
  6. "net/http"
  7. "net/http/httptest"
  8. "strconv"
  9. "strings"
  10. "sync/atomic"
  11. "testing"
  12. "github.com/Wei-Shaw/sub2api/internal/config"
  13. "github.com/Wei-Shaw/sub2api/internal/pkg/openai"
  14. "github.com/gin-gonic/gin"
  15. "github.com/stretchr/testify/require"
  16. "github.com/tidwall/gjson"
  17. )
  18. type alphaSearchAccountStateRepo struct {
  19. AccountRepository
  20. setErrorCalls int
  21. lastError string
  22. }
  23. func (r *alphaSearchAccountStateRepo) SetError(_ context.Context, _ int64, errorMsg string) error {
  24. r.setErrorCalls++
  25. r.lastError = errorMsg
  26. return nil
  27. }
  28. func alphaSearchResponsesSSE(output string) string {
  29. return "event: response.output_text.delta\n" +
  30. `data: {"type":"response.output_text.delta","delta":` + strconv.Quote(output) + `}` + "\n\n" +
  31. "event: response.output_text.annotation.added\n" +
  32. `data: {"type":"response.output_text.annotation.added","annotation":{"type":"url_citation","url":"https://example.com/news","title":"Example News"}}` + "\n\n" +
  33. "event: response.completed\n" +
  34. `data: {"type":"response.completed","response":{"output":[{"type":"message","content":[{"type":"output_text","text":` + strconv.Quote(output) + `}]}]}}` + "\n\n"
  35. }
  36. func TestForwardAlphaSearchOAuthPreservesWire(t *testing.T) {
  37. gin.SetMode(gin.TestMode)
  38. body := []byte(`{
  39. "id":"search-session",
  40. "model":"gpt-5.6-sol",
  41. "reasoning":{"effort":"max","context":"all_turns"},
  42. "input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"latest news"}]}],
  43. "commands":{"search_query":[{"q":"OpenAI news","recency":1}]},
  44. "settings":{"allowed_callers":["direct"],"external_web_access":true},
  45. "max_output_tokens":2000,
  46. "future_field":{"keep":true}
  47. }`)
  48. recorder := httptest.NewRecorder()
  49. c, _ := gin.CreateTestContext(recorder)
  50. c.Request = httptest.NewRequest(http.MethodPost, "/v1/alpha/search?feature=standalone", bytes.NewReader(body))
  51. c.Request.Header.Set("Content-Type", "application/json")
  52. c.Request.Header.Set("User-Agent", codexCLIUserAgent)
  53. c.Request.Header.Set("Originator", "codex_cli_rs")
  54. c.Request.Header.Set("Version", "0.144.1")
  55. c.Request.Header.Set("X-Codex-Turn-Metadata", `{"session_id":"search-session","turn_id":"search-turn"}`)
  56. upstream := &httpUpstreamRecorder{resp: &http.Response{
  57. StatusCode: http.StatusOK,
  58. Header: http.Header{"Content-Type": []string{"application/json"}},
  59. Body: io.NopCloser(strings.NewReader(`{"encrypted_output":"ciphertext","output":"search result"}`)),
  60. }}
  61. service := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
  62. account := &Account{
  63. ID: 42,
  64. Platform: PlatformOpenAI,
  65. Type: AccountTypeOAuth,
  66. Concurrency: 1,
  67. Credentials: map[string]any{
  68. "access_token": "oauth-token",
  69. "chatgpt_account_id": "chatgpt-account",
  70. },
  71. }
  72. result, err := service.ForwardAlphaSearch(context.Background(), c, account, body)
  73. require.NoError(t, err)
  74. require.NotNil(t, result)
  75. require.Equal(t, 1, result.WebSearchCalls)
  76. require.Equal(t, "gpt-5.6-sol", result.Model)
  77. require.Equal(t, http.StatusOK, recorder.Code)
  78. require.JSONEq(t, `{"encrypted_output":"ciphertext","output":"search result"}`, recorder.Body.String())
  79. require.Equal(t, chatgptCodexAlphaSearchURL+"?feature=standalone", upstream.lastReq.URL.String())
  80. require.Equal(t, "chatgpt.com", upstream.lastReq.Host)
  81. require.Equal(t, "Bearer oauth-token", upstream.lastReq.Header.Get("Authorization"))
  82. require.Equal(t, "chatgpt-account", upstream.lastReq.Header.Get("chatgpt-account-id"))
  83. require.Equal(t, "application/json", upstream.lastReq.Header.Get("Accept"))
  84. require.Equal(t, codexCLIVersion, upstream.lastReq.Header.Get("Version"))
  85. require.Empty(t, upstream.lastReq.Header.Get("OpenAI-Beta"))
  86. require.Equal(t,
  87. scopeCodexAccountIdentityValue(account, 0, "session", "search-session"),
  88. gjson.Get(upstream.lastReq.Header.Get("X-Codex-Turn-Metadata"), "session_id").String(),
  89. )
  90. require.Equal(t,
  91. scopeCodexAccountIdentityValue(account, 0, "turn", "search-turn"),
  92. gjson.Get(upstream.lastReq.Header.Get("X-Codex-Turn-Metadata"), "turn_id").String(),
  93. )
  94. require.JSONEq(t, string(body), string(upstream.lastBody))
  95. }
  96. func TestForwardAlphaSearchPATUsesResponsesWebSearchFallback(t *testing.T) {
  97. gin.SetMode(gin.TestMode)
  98. body := []byte(`{
  99. "id":"search-session",
  100. "model":"gpt-5.6-sol",
  101. "commands":{"search_query":[{"q":"OpenAI news"}]},
  102. "prompt_cache_key":"responses-cache-key",
  103. "prompt_cache_retention":"24h"
  104. }`)
  105. recorder := httptest.NewRecorder()
  106. c, _ := gin.CreateTestContext(recorder)
  107. c.Request = httptest.NewRequest(http.MethodPost, "/v1/alpha/search", bytes.NewReader(body))
  108. c.Request.Header.Set("Content-Type", "application/json")
  109. c.Request.Header.Set("User-Agent", codexCLIUserAgent)
  110. c.Request.Header.Set("Originator", "codex_cli_rs")
  111. c.Request.Header.Set("Version", "0.144.1")
  112. c.Request.Header.Set("OpenAI-Beta", "responses=experimental")
  113. c.Request.Header.Set("Accept-Language", "zh-CN")
  114. c.Request.Header.Set("Authorization", "Bearer client-token")
  115. c.Request.Header.Set("Session_ID", "session-client")
  116. c.Request.Header.Set("Conversation_ID", "conversation-client")
  117. c.Request.Header.Set("X-Codex-Beta-Features", "feature-a")
  118. c.Request.Header.Set("X-Codex-Turn-State", "turn-state")
  119. c.Request.Header.Set(responsesLiteHeaderKey, "true")
  120. c.Request.Header.Set("X-Codex-Turn-Metadata", `{"turn_id":"turn-1"}`)
  121. upstream := &httpUpstreamRecorder{resp: &http.Response{
  122. StatusCode: http.StatusOK,
  123. Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"req-search"}},
  124. Body: io.NopCloser(strings.NewReader(alphaSearchResponsesSSE("search result"))),
  125. }}
  126. service := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
  127. account := &Account{
  128. ID: 43,
  129. Platform: PlatformOpenAI,
  130. Type: AccountTypeOAuth,
  131. Concurrency: 1,
  132. Credentials: map[string]any{
  133. "access_token": "at-test-token",
  134. "auth_mode": OpenAIAuthModePersonalAccessToken,
  135. "chatgpt_account_id": "chatgpt-account",
  136. "chatgpt_account_is_fedramp": true,
  137. },
  138. }
  139. result, err := service.ForwardAlphaSearch(context.Background(), c, account, body)
  140. require.NoError(t, err)
  141. require.NotNil(t, result)
  142. require.Equal(t, 1, result.WebSearchCalls)
  143. require.Equal(t, "/v1/responses", result.UpstreamEndpoint)
  144. require.Equal(t, http.StatusOK, recorder.Code)
  145. require.JSONEq(t, `{"output":"search result","results":[{"type":"text_result","ref_id":"turn0search0","url":"https://example.com/news","title":"Example News"}]}`, recorder.Body.String())
  146. require.Equal(t, chatgptCodexURL, upstream.lastReq.URL.String())
  147. require.Equal(t, "Bearer at-test-token", upstream.lastReq.Header.Get("Authorization"))
  148. require.Equal(t, "chatgpt-account", upstream.lastReq.Header.Get("ChatGPT-Account-ID"))
  149. require.Equal(t, "true", upstream.lastReq.Header.Get("X-OpenAI-Fedramp"))
  150. require.Equal(t, "application/json", upstream.lastReq.Header.Get("Content-Type"))
  151. require.Equal(t, "text/event-stream", upstream.lastReq.Header.Get("Accept"))
  152. require.Equal(t, "responses=experimental", upstream.lastReq.Header.Get("OpenAI-Beta"))
  153. require.Equal(t, codexCLIVersion, upstream.lastReq.Header.Get("Version"))
  154. require.Equal(t,
  155. scopeCodexAccountIdentityValue(account, 0, "turn", "turn-1"),
  156. gjson.Get(upstream.lastReq.Header.Get("X-Codex-Turn-Metadata"), "turn_id").String(),
  157. )
  158. require.Equal(t, openai.CodexDefaultOriginator, upstream.lastReq.Header.Get("Originator"))
  159. require.Empty(t, upstream.lastReq.Header.Get("X-Codex-Beta-Features"))
  160. require.Empty(t, upstream.lastReq.Header.Get("X-Codex-Turn-State"))
  161. require.Empty(t, upstream.lastReq.Header.Get(responsesLiteHeaderKey))
  162. require.Empty(t, upstream.lastReq.Header.Get("Accept-Language"))
  163. require.False(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").Exists())
  164. require.False(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_retention").Exists())
  165. require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(upstream.lastBody, "model").String())
  166. require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool())
  167. require.False(t, gjson.GetBytes(upstream.lastBody, "store").Bool())
  168. require.Equal(t, "web_search", gjson.GetBytes(upstream.lastBody, "tools.0.type").String())
  169. require.Contains(t, gjson.GetBytes(upstream.lastBody, "input.0.content.0.text").String(), `"search_query"`)
  170. }
  171. func TestForwardAlphaSearchPATBackfillsMissingChatGPTAccountMetadata(t *testing.T) {
  172. gin.SetMode(gin.TestMode)
  173. body := []byte(`{"id":"search-session","model":"gpt-5.6-sol","commands":{"search_query":[{"q":"OpenAI news"}]}}`)
  174. recorder := httptest.NewRecorder()
  175. c, _ := gin.CreateTestContext(recorder)
  176. c.Request = httptest.NewRequest(http.MethodPost, "/v1/alpha/search", bytes.NewReader(body))
  177. c.Request.Header.Set("Content-Type", "application/json")
  178. var whoamiCalls int32
  179. whoamiServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  180. atomic.AddInt32(&whoamiCalls, 1)
  181. require.Equal(t, "Bearer at-test-token", r.Header.Get("Authorization"))
  182. require.Equal(t, "application/json", r.Header.Get("Accept"))
  183. w.Header().Set("Content-Type", "application/json")
  184. _, _ = w.Write([]byte(`{
  185. "email":"pat@example.com",
  186. "chatgpt_user_id":"user-123",
  187. "chatgpt_account_id":"acct-123",
  188. "chatgpt_plan_type":"plus",
  189. "chatgpt_account_is_fedramp":true
  190. }`))
  191. }))
  192. defer whoamiServer.Close()
  193. oldWhoamiURL := openAICodexPATWhoamiURL
  194. openAICodexPATWhoamiURL = whoamiServer.URL
  195. defer func() { openAICodexPATWhoamiURL = oldWhoamiURL }()
  196. upstream := &httpUpstreamRecorder{resp: &http.Response{
  197. StatusCode: http.StatusOK,
  198. Header: http.Header{"Content-Type": []string{"application/json"}},
  199. Body: io.NopCloser(strings.NewReader(`{"output":"search result"}`)),
  200. }}
  201. oauthService := NewOpenAIOAuthService(nil, nil)
  202. service := &OpenAIGatewayService{
  203. cfg: &config.Config{},
  204. httpUpstream: upstream,
  205. openAITokenProvider: NewOpenAITokenProvider(nil, nil, oauthService),
  206. }
  207. account := &Account{
  208. ID: 45,
  209. Platform: PlatformOpenAI,
  210. Type: AccountTypeOAuth,
  211. Concurrency: 1,
  212. Credentials: map[string]any{
  213. "access_token": "at-test-token",
  214. "auth_mode": OpenAIAuthModePersonalAccessToken,
  215. },
  216. }
  217. result, err := service.ForwardAlphaSearch(context.Background(), c, account, body)
  218. require.NoError(t, err)
  219. require.NotNil(t, result)
  220. require.Equal(t, int32(1), atomic.LoadInt32(&whoamiCalls))
  221. require.Equal(t, "acct-123", upstream.lastReq.Header.Get("ChatGPT-Account-ID"))
  222. require.Equal(t, "true", upstream.lastReq.Header.Get("X-OpenAI-Fedramp"))
  223. require.Equal(t, "acct-123", account.Credentials["chatgpt_account_id"])
  224. require.Equal(t, "user-123", account.Credentials["chatgpt_user_id"])
  225. require.Equal(t, OpenAIAuthModePersonalAccessToken, account.Credentials["auth_mode"])
  226. }
  227. func TestForwardAlphaSearchAPIKeyMapsModelAndPassesThroughError(t *testing.T) {
  228. gin.SetMode(gin.TestMode)
  229. body := []byte(`{"id":"search-session","model":"gpt-5.6-sol","commands":{"search_query":[{"q":"news"}]}}`)
  230. recorder := httptest.NewRecorder()
  231. c, _ := gin.CreateTestContext(recorder)
  232. c.Request = httptest.NewRequest(http.MethodPost, "/alpha/search", bytes.NewReader(body))
  233. c.Request.Header.Set("Content-Type", "application/json")
  234. upstreamBody := `{"error":{"type":"invalid_request_error","message":"bad search"}}`
  235. upstream := &httpUpstreamRecorder{resp: &http.Response{
  236. StatusCode: http.StatusBadRequest,
  237. Header: http.Header{"Content-Type": []string{"application/json"}},
  238. Body: io.NopCloser(strings.NewReader(upstreamBody)),
  239. }}
  240. service := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
  241. account := &Account{
  242. ID: 7,
  243. Platform: PlatformOpenAI,
  244. Type: AccountTypeAPIKey,
  245. Credentials: map[string]any{
  246. "api_key": "sk-test",
  247. "base_url": "https://compat.example/v4",
  248. "model_mapping": map[string]any{
  249. "gpt-5.6-sol": "upstream-5.6",
  250. },
  251. },
  252. }
  253. result, err := service.ForwardAlphaSearch(context.Background(), c, account, body)
  254. require.NoError(t, err)
  255. // 上游错误透传不是一次成功的搜索:不返回 result、不产生按次计费。
  256. require.Nil(t, result)
  257. require.Equal(t, http.StatusBadRequest, recorder.Code)
  258. require.JSONEq(t, upstreamBody, recorder.Body.String())
  259. require.Equal(t, "https://compat.example/v4/alpha/search", upstream.lastReq.URL.String())
  260. require.Equal(t, "Bearer sk-test", upstream.lastReq.Header.Get("Authorization"))
  261. require.Equal(t, "upstream-5.6", gjson.GetBytes(upstream.lastBody, "model").String())
  262. require.True(t, gjson.GetBytes(upstream.lastBody, "commands.search_query").IsArray())
  263. }
  264. func TestForwardAlphaSearchReturnsFailoverBeforeWriting(t *testing.T) {
  265. gin.SetMode(gin.TestMode)
  266. body := []byte(`{"id":"search-session","model":"gpt-5.6-sol","commands":{}}`)
  267. recorder := httptest.NewRecorder()
  268. c, _ := gin.CreateTestContext(recorder)
  269. c.Request = httptest.NewRequest(http.MethodPost, "/v1/alpha/search", bytes.NewReader(body))
  270. upstream := &httpUpstreamRecorder{resp: &http.Response{
  271. StatusCode: http.StatusTooManyRequests,
  272. Header: http.Header{"Content-Type": []string{"application/json"}},
  273. Body: io.NopCloser(strings.NewReader(`{"error":{"message":"rate limited"}}`)),
  274. }}
  275. service := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
  276. account := &Account{
  277. ID: 8,
  278. Platform: PlatformOpenAI,
  279. Type: AccountTypeAPIKey,
  280. Credentials: map[string]any{
  281. "api_key": "sk-test",
  282. },
  283. }
  284. result, err := service.ForwardAlphaSearch(context.Background(), c, account, body)
  285. require.Nil(t, result)
  286. var failoverErr *UpstreamFailoverError
  287. require.ErrorAs(t, err, &failoverErr)
  288. require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
  289. require.Empty(t, failoverErr.Stage)
  290. require.Empty(t, failoverErr.Scope)
  291. require.Empty(t, failoverErr.Reason)
  292. require.Zero(t, failoverErr.ClientStatusCode)
  293. require.Empty(t, failoverErr.ClientMessage)
  294. require.Nil(t, failoverErr.ResponseHeaders)
  295. require.Equal(t, openAIPlatformAlphaSearchURL, upstream.lastReq.URL.String())
  296. require.False(t, c.Writer.Written())
  297. require.Empty(t, recorder.Body.String())
  298. }
  299. func TestForwardAlphaSearchSetupToken429PreservesRetryAfterCooldown(t *testing.T) {
  300. gin.SetMode(gin.TestMode)
  301. body := []byte(`{"id":"search-session","model":"gpt-5.6-sol","commands":{}}`)
  302. recorder := httptest.NewRecorder()
  303. c, _ := gin.CreateTestContext(recorder)
  304. c.Request = httptest.NewRequest(http.MethodPost, "/v1/alpha/search", bytes.NewReader(body))
  305. upstream := &httpUpstreamRecorder{resp: &http.Response{
  306. StatusCode: http.StatusTooManyRequests,
  307. Header: http.Header{
  308. "Content-Type": []string{"application/json"},
  309. "Retry-After": []string{"1"},
  310. "X-Request-Id": []string{"req_alpha_oauth_429"},
  311. },
  312. Body: io.NopCloser(strings.NewReader(`{"error":{"type":"rate_limit_error","code":"rate_limit_exceeded","message":"rate limited"}}`)),
  313. }}
  314. service := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
  315. account := &Account{
  316. ID: 81,
  317. Platform: PlatformOpenAI,
  318. Type: AccountTypeSetupToken,
  319. Concurrency: 1,
  320. Credentials: map[string]any{
  321. "access_token": "oauth-token",
  322. "chatgpt_account_id": "chatgpt-account",
  323. },
  324. }
  325. result, err := service.ForwardAlphaSearch(context.Background(), c, account, body)
  326. require.Nil(t, result)
  327. var failoverErr *UpstreamFailoverError
  328. require.ErrorAs(t, err, &failoverErr)
  329. require.False(t, failoverErr.RetryableOnSameAccount)
  330. require.Zero(t, failoverErr.SameAccountRetryDelay)
  331. require.True(t, failoverErr.SameAccountRetryDeadline.IsZero())
  332. require.Equal(t, "1", failoverErr.ResponseHeaders.Get("Retry-After"))
  333. require.Equal(t, "req_alpha_oauth_429", failoverErr.ResponseHeaders.Get("x-request-id"))
  334. require.False(t, c.Writer.Written())
  335. }
  336. func TestForwardAlphaSearchAccessStateUsesTypedFailover(t *testing.T) {
  337. gin.SetMode(gin.TestMode)
  338. body := []byte(`{"id":"search-session","model":"gpt-5.6-sol","commands":{}}`)
  339. recorder := httptest.NewRecorder()
  340. c, _ := gin.CreateTestContext(recorder)
  341. c.Request = httptest.NewRequest(http.MethodPost, "/v1/alpha/search", bytes.NewReader(body))
  342. upstream := &httpUpstreamRecorder{resp: &http.Response{
  343. StatusCode: http.StatusForbidden,
  344. Header: http.Header{
  345. "Content-Type": []string{"application/json"},
  346. "X-Request-Id": []string{"req_alpha_access_state"},
  347. },
  348. Body: io.NopCloser(strings.NewReader(`{"error":{"code":"deactivated_workspace","message":"Workspace is deactivated"}}`)),
  349. }}
  350. service := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
  351. account := &Account{
  352. ID: 11,
  353. Platform: PlatformOpenAI,
  354. Type: AccountTypeOAuth,
  355. Concurrency: 1,
  356. Credentials: map[string]any{
  357. "access_token": "oauth-token",
  358. "chatgpt_account_id": "chatgpt-account",
  359. },
  360. }
  361. result, err := service.ForwardAlphaSearch(context.Background(), c, account, body)
  362. require.Nil(t, result)
  363. assertOpenAIAlphaSearchAccessStateFailover(t, err, "req_alpha_access_state")
  364. require.Equal(t, chatgptCodexAlphaSearchURL, upstream.lastReq.URL.String())
  365. require.False(t, c.Writer.Written())
  366. }
  367. func TestForwardAlphaSearchPATFallbackAccessStateUsesTypedFailover(t *testing.T) {
  368. gin.SetMode(gin.TestMode)
  369. body := []byte(`{"id":"search-session","model":"gpt-5.6-sol","commands":{"search_query":[{"q":"news"}]}}`)
  370. recorder := httptest.NewRecorder()
  371. c, _ := gin.CreateTestContext(recorder)
  372. c.Request = httptest.NewRequest(http.MethodPost, "/v1/alpha/search", bytes.NewReader(body))
  373. upstream := &httpUpstreamRecorder{resp: &http.Response{
  374. StatusCode: http.StatusForbidden,
  375. Header: http.Header{
  376. "Content-Type": []string{"application/json"},
  377. "X-Request-Id": []string{"req_alpha_pat_access_state"},
  378. },
  379. Body: io.NopCloser(strings.NewReader(`{"error":{"code":"account_disabled","message":"Account is disabled"}}`)),
  380. }}
  381. service := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
  382. account := &Account{
  383. ID: 12,
  384. Platform: PlatformOpenAI,
  385. Type: AccountTypeOAuth,
  386. Concurrency: 1,
  387. Credentials: map[string]any{
  388. "access_token": "at-test-token",
  389. "auth_mode": OpenAIAuthModePersonalAccessToken,
  390. "chatgpt_account_id": "chatgpt-account",
  391. },
  392. }
  393. result, err := service.ForwardAlphaSearch(context.Background(), c, account, body)
  394. require.Nil(t, result)
  395. assertOpenAIAlphaSearchAccessStateFailover(t, err, "req_alpha_pat_access_state")
  396. require.Equal(t, chatgptCodexURL, upstream.lastReq.URL.String())
  397. require.False(t, c.Writer.Written())
  398. }
  399. func assertOpenAIAlphaSearchAccessStateFailover(t *testing.T, err error, requestID string) {
  400. t.Helper()
  401. var failoverErr *UpstreamFailoverError
  402. require.ErrorAs(t, err, &failoverErr)
  403. require.Equal(t, http.StatusForbidden, failoverErr.StatusCode)
  404. require.Equal(t, GatewayFailureStageAccountAuth, failoverErr.Stage)
  405. require.Equal(t, GatewayFailureScopeAccount, failoverErr.Scope)
  406. require.Equal(t, OpenAIUpstreamAccessStateReason, failoverErr.Reason)
  407. require.Equal(t, NextAccountRetry, failoverErr.NextAccountAction)
  408. require.Equal(t, http.StatusBadGateway, failoverErr.ClientStatusCode)
  409. require.Equal(t, openAIUpstreamAccessUnavailableClientMessage, failoverErr.ClientMessage)
  410. require.False(t, failoverErr.RetryableOnSameAccount)
  411. require.Equal(t, requestID, failoverErr.ResponseHeaders.Get("x-request-id"))
  412. }
  413. func TestForwardAlphaSearchUnauthorizedDoesNotMarkAccountError(t *testing.T) {
  414. gin.SetMode(gin.TestMode)
  415. body := []byte(`{"id":"search-session","model":"gpt-5.6-sol","commands":{"search_query":[{"q":"news"}]}}`)
  416. recorder := httptest.NewRecorder()
  417. c, _ := gin.CreateTestContext(recorder)
  418. c.Request = httptest.NewRequest(http.MethodPost, "/v1/alpha/search", bytes.NewReader(body))
  419. upstream := &httpUpstreamRecorder{resp: &http.Response{
  420. StatusCode: http.StatusUnauthorized,
  421. Header: http.Header{"Content-Type": []string{"application/json"}},
  422. Body: io.NopCloser(strings.NewReader(`{"detail":"Unauthorized"}`)),
  423. }}
  424. repo := &alphaSearchAccountStateRepo{}
  425. cfg := &config.Config{}
  426. service := &OpenAIGatewayService{
  427. cfg: cfg,
  428. httpUpstream: upstream,
  429. accountRepo: repo,
  430. rateLimitService: NewRateLimitService(repo, nil, cfg, nil, nil),
  431. }
  432. account := &Account{
  433. ID: 44,
  434. Platform: PlatformOpenAI,
  435. Type: AccountTypeOAuth,
  436. Concurrency: 1,
  437. Credentials: map[string]any{
  438. // 刻意不设置 auth_mode:覆盖历史上把 at- token 当普通 OAuth 导入的账号。
  439. "access_token": "at-test-token",
  440. "chatgpt_account_id": "chatgpt-account",
  441. },
  442. }
  443. result, err := service.ForwardAlphaSearch(context.Background(), c, account, body)
  444. require.Nil(t, result)
  445. var failoverErr *UpstreamFailoverError
  446. require.ErrorAs(t, err, &failoverErr)
  447. require.Equal(t, http.StatusUnauthorized, failoverErr.StatusCode)
  448. require.Zero(t, repo.setErrorCalls)
  449. require.Empty(t, repo.lastError)
  450. require.False(t, c.Writer.Written())
  451. }
  452. func TestForwardAlphaSearchPATResponsesFallbackUnauthorizedDoesNotMarkAccountError(t *testing.T) {
  453. gin.SetMode(gin.TestMode)
  454. body := []byte(`{"id":"search-session","model":"gpt-5.6-sol","commands":{"search_query":[{"q":"news"}]}}`)
  455. recorder := httptest.NewRecorder()
  456. c, _ := gin.CreateTestContext(recorder)
  457. c.Request = httptest.NewRequest(http.MethodPost, "/v1/alpha/search", bytes.NewReader(body))
  458. upstream := &httpUpstreamRecorder{resp: &http.Response{
  459. StatusCode: http.StatusUnauthorized,
  460. Header: http.Header{"Content-Type": []string{"application/json"}},
  461. Body: io.NopCloser(strings.NewReader(`{"detail":"Unauthorized"}`)),
  462. }}
  463. repo := &alphaSearchAccountStateRepo{}
  464. cfg := &config.Config{}
  465. service := &OpenAIGatewayService{
  466. cfg: cfg,
  467. httpUpstream: upstream,
  468. accountRepo: repo,
  469. rateLimitService: NewRateLimitService(repo, nil, cfg, nil, nil),
  470. }
  471. account := &Account{
  472. ID: 46,
  473. Platform: PlatformOpenAI,
  474. Type: AccountTypeOAuth,
  475. Concurrency: 1,
  476. Credentials: map[string]any{
  477. "access_token": "at-test-token",
  478. "auth_mode": OpenAIAuthModePersonalAccessToken,
  479. "chatgpt_account_id": "chatgpt-account",
  480. },
  481. }
  482. result, err := service.ForwardAlphaSearch(context.Background(), c, account, body)
  483. require.Nil(t, result)
  484. var failoverErr *UpstreamFailoverError
  485. require.ErrorAs(t, err, &failoverErr)
  486. require.Equal(t, http.StatusUnauthorized, failoverErr.StatusCode)
  487. require.Equal(t, chatgptCodexURL, upstream.lastReq.URL.String())
  488. require.Equal(t, "text/event-stream", upstream.lastReq.Header.Get("Accept"))
  489. require.Equal(t, "responses=experimental", upstream.lastReq.Header.Get("OpenAI-Beta"))
  490. require.Zero(t, repo.setErrorCalls)
  491. require.Empty(t, repo.lastError)
  492. require.False(t, c.Writer.Written())
  493. }
  494. // API key 上游(官方平台或第三方中转)不提供 /v1/alpha/search 时返回的
  495. // 404/405 必须触发换号而不是把错误透传给客户端:混合分组里 OAuth 账号可以
  496. // 承接搜索,请求不能死在先被选中的 API key 账号上。端点缺失也不能写账号
  497. // 错误状态——账号本身是健康的。
  498. func TestForwardAlphaSearchAPIKeyEndpointNotFoundFailsOver(t *testing.T) {
  499. gin.SetMode(gin.TestMode)
  500. body := []byte(`{"id":"search-session","model":"gpt-5.6-sol","commands":{"search_query":[{"q":"news"}]}}`)
  501. recorder := httptest.NewRecorder()
  502. c, _ := gin.CreateTestContext(recorder)
  503. c.Request = httptest.NewRequest(http.MethodPost, "/v1/alpha/search", bytes.NewReader(body))
  504. upstream := &httpUpstreamRecorder{resp: &http.Response{
  505. StatusCode: http.StatusNotFound,
  506. Header: http.Header{"Content-Type": []string{"application/json"}},
  507. Body: io.NopCloser(strings.NewReader(`{"error":{"message":"Not Found"}}`)),
  508. }}
  509. repo := &alphaSearchAccountStateRepo{}
  510. cfg := &config.Config{}
  511. service := &OpenAIGatewayService{
  512. cfg: cfg,
  513. httpUpstream: upstream,
  514. accountRepo: repo,
  515. rateLimitService: NewRateLimitService(repo, nil, cfg, nil, nil),
  516. }
  517. account := &Account{
  518. ID: 9,
  519. Platform: PlatformOpenAI,
  520. Type: AccountTypeAPIKey,
  521. Credentials: map[string]any{
  522. "api_key": "sk-test",
  523. "base_url": "https://relay.example",
  524. },
  525. }
  526. result, err := service.ForwardAlphaSearch(context.Background(), c, account, body)
  527. require.Nil(t, result)
  528. var failoverErr *UpstreamFailoverError
  529. require.ErrorAs(t, err, &failoverErr)
  530. require.Equal(t, http.StatusNotFound, failoverErr.StatusCode)
  531. require.Zero(t, repo.setErrorCalls)
  532. require.Empty(t, repo.lastError)
  533. require.False(t, c.Writer.Written())
  534. require.Empty(t, recorder.Body.String())
  535. }
  536. // OAuth 账号的 chatgpt.com 端点固定存在,404 保持原有透传行为不变。
  537. func TestForwardAlphaSearchOAuthNotFoundPassesThrough(t *testing.T) {
  538. gin.SetMode(gin.TestMode)
  539. body := []byte(`{"id":"search-session","model":"gpt-5.6-sol","commands":{"search_query":[{"q":"news"}]}}`)
  540. recorder := httptest.NewRecorder()
  541. c, _ := gin.CreateTestContext(recorder)
  542. c.Request = httptest.NewRequest(http.MethodPost, "/v1/alpha/search", bytes.NewReader(body))
  543. upstreamBody := `{"detail":"Not Found"}`
  544. upstream := &httpUpstreamRecorder{resp: &http.Response{
  545. StatusCode: http.StatusNotFound,
  546. Header: http.Header{"Content-Type": []string{"application/json"}},
  547. Body: io.NopCloser(strings.NewReader(upstreamBody)),
  548. }}
  549. service := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
  550. account := &Account{
  551. ID: 10,
  552. Platform: PlatformOpenAI,
  553. Type: AccountTypeOAuth,
  554. Concurrency: 1,
  555. Credentials: map[string]any{
  556. "access_token": "oauth-token",
  557. "chatgpt_account_id": "chatgpt-account",
  558. },
  559. }
  560. result, err := service.ForwardAlphaSearch(context.Background(), c, account, body)
  561. require.NoError(t, err)
  562. require.Nil(t, result)
  563. require.Equal(t, http.StatusNotFound, recorder.Code)
  564. require.JSONEq(t, upstreamBody, recorder.Body.String())
  565. }
  566. func TestShouldApplyOpenAIAlphaSearchAccountErrorSideEffects(t *testing.T) {
  567. require.False(t, shouldApplyOpenAIAlphaSearchAccountErrorSideEffects(http.StatusUnauthorized))
  568. require.False(t, shouldApplyOpenAIAlphaSearchAccountErrorSideEffects(http.StatusNotFound))
  569. require.False(t, shouldApplyOpenAIAlphaSearchAccountErrorSideEffects(http.StatusMethodNotAllowed))
  570. require.True(t, shouldApplyOpenAIAlphaSearchAccountErrorSideEffects(http.StatusForbidden))
  571. require.True(t, shouldApplyOpenAIAlphaSearchAccountErrorSideEffects(http.StatusTooManyRequests))
  572. }
  573. func TestOpenAIAlphaSearchSchedulingModelUsesCanonicalAccountMapping(t *testing.T) {
  574. account := &Account{Credentials: map[string]any{
  575. "model_mapping": map[string]any{"client-visible": "canonical-upstream"},
  576. }}
  577. require.Equal(t, "canonical-upstream", openAIAlphaSearchSchedulingModel(account, "client-visible"))
  578. require.Equal(t, "unmapped", openAIAlphaSearchSchedulingModel(account, "unmapped"))
  579. }
  580. func TestSanitizeOpenAIAlphaSearchBody_RemovesResponsesOnlyFields(t *testing.T) {
  581. body := []byte(`{"id":"search-session","store":false,"prompt_cache_key":"cache","commands":{"search_query":[{"q":"news"}]}}`)
  582. normalized, err := sanitizeOpenAIAlphaSearchBody(body)
  583. require.NoError(t, err)
  584. require.False(t, gjson.GetBytes(normalized, "store").Exists())
  585. require.False(t, gjson.GetBytes(normalized, "prompt_cache_key").Exists())
  586. require.Equal(t, "news", gjson.GetBytes(normalized, "commands.search_query.0.q").String())
  587. }
  588. func TestIsOpenAIAlphaSearchEndpointUnsupported(t *testing.T) {
  589. apiKey := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  590. oauth := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  591. require.True(t, isOpenAIAlphaSearchEndpointUnsupported(apiKey, http.StatusNotFound))
  592. require.True(t, isOpenAIAlphaSearchEndpointUnsupported(apiKey, http.StatusMethodNotAllowed))
  593. require.False(t, isOpenAIAlphaSearchEndpointUnsupported(apiKey, http.StatusBadRequest))
  594. require.False(t, isOpenAIAlphaSearchEndpointUnsupported(oauth, http.StatusNotFound))
  595. require.False(t, isOpenAIAlphaSearchEndpointUnsupported(nil, http.StatusNotFound))
  596. }