openai_gateway_chat_completions_test.go 47 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138
  1. package service
  2. import (
  3. "bytes"
  4. "context"
  5. "errors"
  6. "io"
  7. "net/http"
  8. "net/http/httptest"
  9. "strings"
  10. "testing"
  11. "time"
  12. "github.com/Wei-Shaw/sub2api/internal/config"
  13. "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
  14. "github.com/gin-gonic/gin"
  15. "github.com/stretchr/testify/require"
  16. "github.com/tidwall/gjson"
  17. )
  18. type openAIChatFailingWriter struct {
  19. gin.ResponseWriter
  20. failAfter int
  21. writes int
  22. }
  23. func (w *openAIChatFailingWriter) Write(p []byte) (int, error) {
  24. if w.writes >= w.failAfter {
  25. return 0, errors.New("write failed: client disconnected")
  26. }
  27. w.writes++
  28. return w.ResponseWriter.Write(p)
  29. }
  30. type openAIChatStreamReadErrorCloser struct {
  31. payload []byte
  32. err error
  33. sent bool
  34. }
  35. func (r *openAIChatStreamReadErrorCloser) Read(p []byte) (int, error) {
  36. if !r.sent {
  37. r.sent = true
  38. return copy(p, r.payload), nil
  39. }
  40. return 0, r.err
  41. }
  42. func (r *openAIChatStreamReadErrorCloser) Close() error { return nil }
  43. func TestHandleChatStreamingResponse_ClassifiesHTTP2ReadError(t *testing.T) {
  44. gin.SetMode(gin.TestMode)
  45. rec := httptest.NewRecorder()
  46. c, _ := gin.CreateTestContext(rec)
  47. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
  48. resp := &http.Response{
  49. StatusCode: http.StatusOK,
  50. Header: http.Header{
  51. "Content-Type": []string{"text/event-stream"},
  52. "x-request-id": []string{"upstream-rid"},
  53. },
  54. Body: &openAIChatStreamReadErrorCloser{
  55. payload: []byte("data: {\"type\":\"response.output_text.delta\",\"delta\":\"partial\"}\n\n"),
  56. err: errors.New("stream error: stream ID 5; INTERNAL_ERROR; received from peer"),
  57. },
  58. }
  59. svc := &OpenAIGatewayService{cfg: &config.Config{}}
  60. result, err := svc.handleChatStreamingResponse(
  61. resp,
  62. c,
  63. &Account{ID: 1, Name: "openai-oauth", Platform: PlatformOpenAI},
  64. "gpt-5.6-sol",
  65. "gpt-5.6-sol",
  66. "gpt-5.6-sol",
  67. time.Now(),
  68. 0,
  69. )
  70. require.Error(t, err)
  71. require.NotNil(t, result)
  72. require.True(t, c.Writer.Written(), "partial output must make replay unsafe")
  73. code, message, ok := OpenAIUpstreamStreamReadErrorDetails(err)
  74. require.True(t, ok)
  75. require.Equal(t, OpenAIUpstreamHTTP2StreamErrorCode, code)
  76. require.Equal(t, "Upstream HTTP/2 stream failed", message)
  77. require.NotContains(t, message, "stream ID")
  78. require.NotContains(t, message, "INTERNAL_ERROR")
  79. }
  80. func TestNormalizeResponsesRequestServiceTier(t *testing.T) {
  81. t.Parallel()
  82. req := &apicompat.ResponsesRequest{ServiceTier: " fast "}
  83. normalizeResponsesRequestServiceTier(req)
  84. require.Equal(t, "priority", req.ServiceTier)
  85. req.ServiceTier = "flex"
  86. normalizeResponsesRequestServiceTier(req)
  87. require.Equal(t, "flex", req.ServiceTier)
  88. // OpenAI 官方合法 tier 应被透传保留。
  89. req.ServiceTier = "auto"
  90. normalizeResponsesRequestServiceTier(req)
  91. require.Equal(t, "auto", req.ServiceTier)
  92. req.ServiceTier = "default"
  93. normalizeResponsesRequestServiceTier(req)
  94. require.Equal(t, "default", req.ServiceTier)
  95. req.ServiceTier = "scale"
  96. normalizeResponsesRequestServiceTier(req)
  97. require.Equal(t, "scale", req.ServiceTier)
  98. // 真未知值仍被剥离。
  99. req.ServiceTier = "turbo"
  100. normalizeResponsesRequestServiceTier(req)
  101. require.Empty(t, req.ServiceTier)
  102. }
  103. func TestNormalizeResponsesBodyServiceTier(t *testing.T) {
  104. t.Parallel()
  105. body, tier, err := normalizeResponsesBodyServiceTier([]byte(`{"model":"gpt-5.1","service_tier":"fast"}`))
  106. require.NoError(t, err)
  107. require.Equal(t, "priority", tier)
  108. require.Equal(t, "priority", gjson.GetBytes(body, "service_tier").String())
  109. body, tier, err = normalizeResponsesBodyServiceTier([]byte(`{"model":"gpt-5.1","service_tier":"flex"}`))
  110. require.NoError(t, err)
  111. require.Equal(t, "flex", tier)
  112. require.Equal(t, "flex", gjson.GetBytes(body, "service_tier").String())
  113. // OpenAI 官方 tier 直接保留在 body 中(透传上游)。
  114. body, tier, err = normalizeResponsesBodyServiceTier([]byte(`{"model":"gpt-5.1","service_tier":"auto"}`))
  115. require.NoError(t, err)
  116. require.Equal(t, "auto", tier)
  117. require.Equal(t, "auto", gjson.GetBytes(body, "service_tier").String())
  118. body, tier, err = normalizeResponsesBodyServiceTier([]byte(`{"model":"gpt-5.1","service_tier":"default"}`))
  119. require.NoError(t, err)
  120. require.Equal(t, "default", tier)
  121. require.Equal(t, "default", gjson.GetBytes(body, "service_tier").String())
  122. body, tier, err = normalizeResponsesBodyServiceTier([]byte(`{"model":"gpt-5.1","service_tier":"scale"}`))
  123. require.NoError(t, err)
  124. require.Equal(t, "scale", tier)
  125. require.Equal(t, "scale", gjson.GetBytes(body, "service_tier").String())
  126. body, tier, err = normalizeResponsesBodyServiceTier([]byte(`{"model":"gpt-5.6-sol","service_tier":"ultrafast"}`))
  127. require.NoError(t, err)
  128. require.Equal(t, "ultrafast", tier)
  129. require.Equal(t, "ultrafast", gjson.GetBytes(body, "service_tier").String())
  130. // 真未知值才会被删除。
  131. body, tier, err = normalizeResponsesBodyServiceTier([]byte(`{"model":"gpt-5.1","service_tier":"turbo"}`))
  132. require.NoError(t, err)
  133. require.Empty(t, tier)
  134. require.False(t, gjson.GetBytes(body, "service_tier").Exists())
  135. }
  136. func TestForwardAsChatCompletions_UnknownModelWithoutMessagesDispatchKeepsRequestedModel(t *testing.T) {
  137. gin.SetMode(gin.TestMode)
  138. rec := httptest.NewRecorder()
  139. c, _ := gin.CreateTestContext(rec)
  140. body := []byte(`{"model":"gpt6","messages":[{"role":"user","content":"hello"}],"stream":false}`)
  141. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  142. c.Request.Header.Set("Content-Type", "application/json")
  143. upstream := &httpUpstreamRecorder{resp: &http.Response{
  144. StatusCode: http.StatusBadRequest,
  145. Header: http.Header{"Content-Type": []string{"application/json"}, "x-request-id": []string{"rid_chat_unknown_model"}},
  146. Body: io.NopCloser(strings.NewReader(`{"error":{"type":"invalid_request_error","message":"model not found"}}`)),
  147. }}
  148. svc := &OpenAIGatewayService{
  149. cfg: &config.Config{},
  150. httpUpstream: upstream,
  151. }
  152. account := &Account{
  153. ID: 1,
  154. Name: "openai-oauth",
  155. Platform: PlatformOpenAI,
  156. Type: AccountTypeOAuth,
  157. Concurrency: 1,
  158. Credentials: map[string]any{
  159. "access_token": "oauth-token",
  160. "chatgpt_account_id": "chatgpt-acc",
  161. },
  162. }
  163. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
  164. require.Error(t, err)
  165. require.Nil(t, result)
  166. require.Equal(t, "gpt6", gjson.GetBytes(upstream.lastBody, "model").String())
  167. require.NotEqual(t, "gpt-5.4", gjson.GetBytes(upstream.lastBody, "model").String())
  168. require.Equal(t, http.StatusBadRequest, rec.Code)
  169. }
  170. func TestForwardAsChatCompletions_APIKeyPropagatesPromptCacheKeyInResponsesBody(t *testing.T) {
  171. gin.SetMode(gin.TestMode)
  172. rec := httptest.NewRecorder()
  173. c, _ := gin.CreateTestContext(rec)
  174. body := []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":"hello"}],"stream":false}`)
  175. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  176. c.Request.Header.Set("Content-Type", "application/json")
  177. c.Set("api_key", &APIKey{ID: 99})
  178. upstream := &httpUpstreamRecorder{resp: &http.Response{
  179. StatusCode: http.StatusBadRequest,
  180. Header: http.Header{"Content-Type": []string{"application/json"}, "x-request-id": []string{"rid_chat_prompt_cache"}},
  181. Body: io.NopCloser(strings.NewReader(`{"error":{"type":"invalid_request_error","message":"stop before response parsing"}}`)),
  182. }}
  183. svc := &OpenAIGatewayService{
  184. cfg: &config.Config{},
  185. httpUpstream: upstream,
  186. }
  187. account := &Account{
  188. ID: 2,
  189. Name: "openai-compatible",
  190. Platform: PlatformOpenAI,
  191. Type: AccountTypeAPIKey,
  192. Concurrency: 1,
  193. Credentials: map[string]any{
  194. "api_key": "sk-compatible",
  195. },
  196. Extra: map[string]any{
  197. "openai_responses_supported": true,
  198. },
  199. }
  200. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "cache-key-123", "gpt-5.4")
  201. require.Error(t, err)
  202. require.Nil(t, result)
  203. require.Equal(t, "cache-key-123", gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String())
  204. require.Equal(t, "gpt-5.4", gjson.GetBytes(upstream.lastBody, "model").String())
  205. require.Equal(t, "https://api.openai.com/v1/responses", upstream.lastReq.URL.String())
  206. require.Equal(t, "Bearer sk-compatible", upstream.lastReq.Header.Get("Authorization"))
  207. require.Equal(t, generateSessionUUID(isolateOpenAISessionID(99, "cache-key-123")), upstream.lastReq.Header.Get("session_id"))
  208. }
  209. func TestForwardAsChatCompletions_APIKeyAutoDerivesStableIsolatedPromptCacheKey(t *testing.T) {
  210. gin.SetMode(gin.TestMode)
  211. response := func() *http.Response {
  212. return &http.Response{
  213. StatusCode: http.StatusBadRequest,
  214. Header: http.Header{"Content-Type": []string{"application/json"}},
  215. Body: io.NopCloser(strings.NewReader(`{"error":{"type":"invalid_request_error","message":"stop before response parsing"}}`)),
  216. }
  217. }
  218. upstream := &httpUpstreamRecorder{responses: []*http.Response{response(), response(), response()}}
  219. svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
  220. account := &Account{
  221. ID: 2, Name: "openai-compatible", Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1,
  222. Credentials: map[string]any{"api_key": "sk-compatible"},
  223. Extra: map[string]any{"openai_responses_supported": true},
  224. }
  225. firstBody := []byte(`{"model":"gpt-5.4","messages":[{"role":"system","content":"be concise"},{"role":"user","content":"hello"}],"stream":false}`)
  226. appendedBody := []byte(`{"model":"gpt-5.4","messages":[{"role":"system","content":"be concise"},{"role":"user","content":"hello"},{"role":"assistant","content":"hi"},{"role":"user","content":"continue"}],"stream":false}`)
  227. forward := func(apiKeyID int64, body []byte) {
  228. rec := httptest.NewRecorder()
  229. c, _ := gin.CreateTestContext(rec)
  230. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  231. c.Request.Header.Set("Content-Type", "application/json")
  232. c.Set("api_key", &APIKey{ID: apiKeyID})
  233. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "gpt-5.4")
  234. require.Error(t, err)
  235. require.Nil(t, result)
  236. }
  237. forward(99, firstBody)
  238. forward(99, appendedBody)
  239. forward(100, appendedBody)
  240. firstKey := gjson.GetBytes(upstream.bodies[0], "prompt_cache_key").String()
  241. appendedKey := gjson.GetBytes(upstream.bodies[1], "prompt_cache_key").String()
  242. otherTenantKey := gjson.GetBytes(upstream.bodies[2], "prompt_cache_key").String()
  243. require.NotEmpty(t, firstKey)
  244. require.Equal(t, firstKey, appendedKey)
  245. require.NotEqual(t, firstKey, otherTenantKey)
  246. require.Equal(t, generateSessionUUID(firstKey), upstream.requests[0].Header.Get("session_id"))
  247. require.Equal(t, upstream.requests[0].Header.Get("session_id"), upstream.requests[1].Header.Get("session_id"))
  248. require.Equal(t, generateSessionUUID(otherTenantKey), upstream.requests[2].Header.Get("session_id"))
  249. require.NotEqual(t, upstream.requests[1].Header.Get("session_id"), upstream.requests[2].Header.Get("session_id"))
  250. }
  251. func TestForwardAsChatCompletions_ResponsesShapeDoesNotAutoDerivePromptCacheKey(t *testing.T) {
  252. gin.SetMode(gin.TestMode)
  253. response := func() *http.Response {
  254. return &http.Response{
  255. StatusCode: http.StatusBadRequest,
  256. Header: http.Header{"Content-Type": []string{"application/json"}},
  257. Body: io.NopCloser(strings.NewReader(`{"error":{"type":"invalid_request_error","message":"stop before response parsing"}}`)),
  258. }
  259. }
  260. upstream := &httpUpstreamRecorder{responses: []*http.Response{response(), response(), response()}}
  261. svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
  262. account := &Account{
  263. ID: 2, Name: "openai-compatible", Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1,
  264. Credentials: map[string]any{"api_key": "sk-compatible"},
  265. Extra: map[string]any{"openai_responses_supported": true},
  266. }
  267. firstBody := []byte(`{"model":"gpt-5.4","input":[{"role":"user","content":[{"type":"input_text","text":"first unrelated input"}]}],"stream":false}`)
  268. secondBody := []byte(`{"model":"gpt-5.4","input":[{"role":"user","content":[{"type":"input_text","text":"second unrelated input"}]}],"stream":false}`)
  269. forward := func(body []byte, promptCacheKey string) {
  270. rec := httptest.NewRecorder()
  271. c, _ := gin.CreateTestContext(rec)
  272. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  273. c.Set("api_key", &APIKey{ID: 99})
  274. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, promptCacheKey, "gpt-5.4")
  275. require.Error(t, err)
  276. require.Nil(t, result)
  277. }
  278. forward(firstBody, "")
  279. forward(secondBody, "")
  280. forward(secondBody, "explicit-responses-key")
  281. require.False(t, gjson.GetBytes(upstream.bodies[0], "prompt_cache_key").Exists())
  282. require.Empty(t, upstream.requests[0].Header.Get("session_id"))
  283. require.False(t, gjson.GetBytes(upstream.bodies[1], "prompt_cache_key").Exists())
  284. require.Empty(t, upstream.requests[1].Header.Get("session_id"))
  285. require.Equal(t, "explicit-responses-key", gjson.GetBytes(upstream.bodies[2], "prompt_cache_key").String())
  286. require.Equal(t, generateSessionUUID(isolateOpenAISessionID(99, "explicit-responses-key")), upstream.requests[2].Header.Get("session_id"))
  287. }
  288. func TestForwardAsChatCompletions_OAuthDoesNotInjectDefaultInstructions(t *testing.T) {
  289. gin.SetMode(gin.TestMode)
  290. rec := httptest.NewRecorder()
  291. c, _ := gin.CreateTestContext(rec)
  292. body := []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":"hello"}],"stream":false}`)
  293. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  294. c.Request.Header.Set("Content-Type", "application/json")
  295. upstream := &httpUpstreamRecorder{resp: &http.Response{
  296. StatusCode: http.StatusBadRequest,
  297. Header: http.Header{"Content-Type": []string{"application/json"}, "x-request-id": []string{"rid_chat_no_default_instructions"}},
  298. Body: io.NopCloser(strings.NewReader(`{"error":{"type":"invalid_request_error","message":"stop before response parsing"}}`)),
  299. }}
  300. svc := &OpenAIGatewayService{
  301. cfg: &config.Config{},
  302. httpUpstream: upstream,
  303. }
  304. account := &Account{
  305. ID: 3,
  306. Name: "openai-oauth",
  307. Platform: PlatformOpenAI,
  308. Type: AccountTypeOAuth,
  309. Concurrency: 1,
  310. Credentials: map[string]any{
  311. "access_token": "oauth-token",
  312. "chatgpt_account_id": "chatgpt-acc",
  313. },
  314. }
  315. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "gpt-5.4")
  316. require.Error(t, err)
  317. require.Nil(t, result)
  318. require.NotNil(t, upstream.lastReq)
  319. require.Equal(t, chatgptCodexURL, upstream.lastReq.URL.String())
  320. require.True(t, gjson.GetBytes(upstream.lastBody, "instructions").Exists())
  321. require.Equal(t, "", gjson.GetBytes(upstream.lastBody, "instructions").String())
  322. require.NotContains(t, string(upstream.lastBody), "Communicate with the user by streaming thinking")
  323. }
  324. func forwardOAuthChatCompletionsForUpstreamBody(t *testing.T, body []byte) []byte {
  325. t.Helper()
  326. gin.SetMode(gin.TestMode)
  327. rec := httptest.NewRecorder()
  328. c, _ := gin.CreateTestContext(rec)
  329. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  330. c.Request.Header.Set("Content-Type", "application/json")
  331. upstream := &httpUpstreamRecorder{resp: &http.Response{
  332. StatusCode: http.StatusBadRequest,
  333. Header: http.Header{"Content-Type": []string{"application/json"}, "x-request-id": []string{"rid_chat_system_promotion"}},
  334. Body: io.NopCloser(strings.NewReader(`{"error":{"type":"invalid_request_error","message":"stop before response parsing"}}`)),
  335. }}
  336. svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
  337. account := &Account{
  338. ID: 4,
  339. Name: "openai-oauth",
  340. Platform: PlatformOpenAI,
  341. Type: AccountTypeOAuth,
  342. Concurrency: 1,
  343. Credentials: map[string]any{
  344. "access_token": "oauth-token",
  345. "chatgpt_account_id": "chatgpt-acc",
  346. },
  347. }
  348. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "gpt-5.4")
  349. require.Error(t, err)
  350. require.Nil(t, result)
  351. require.NotEmpty(t, upstream.lastBody)
  352. return upstream.lastBody
  353. }
  354. func TestForwardAsChatCompletions_OAuthPromotesSystemMessageWithoutDuplication(t *testing.T) {
  355. const systemPrompt = "Unique system prefix for token accounting."
  356. body := []byte(`{"model":"gpt-5.4","messages":[{"role":"system","content":"` + systemPrompt + `"},{"role":"user","content":"hello"}],"stream":false}`)
  357. upstreamBody := forwardOAuthChatCompletionsForUpstreamBody(t, body)
  358. require.Equal(t, systemPrompt, gjson.GetBytes(upstreamBody, "instructions").String())
  359. require.Equal(t, int64(1), gjson.GetBytes(upstreamBody, "input.#").Int())
  360. require.Equal(t, "user", gjson.GetBytes(upstreamBody, "input.0.role").String())
  361. require.Equal(t, 1, strings.Count(string(upstreamBody), systemPrompt))
  362. }
  363. func TestForwardAsChatCompletions_OAuthJsonObjectKeepsSystemMessageInInput(t *testing.T) {
  364. const systemPrompt = "Return JSON only."
  365. body := []byte(`{"model":"gpt-5.4","messages":[{"role":"system","content":"` + systemPrompt + `"},{"role":"user","content":"symbol data"}],"response_format":{"type":" JSON_OBJECT "},"stream":false}`)
  366. upstreamBody := forwardOAuthChatCompletionsForUpstreamBody(t, body)
  367. require.Equal(t, systemPrompt, gjson.GetBytes(upstreamBody, "instructions").String())
  368. require.Equal(t, int64(2), gjson.GetBytes(upstreamBody, "input.#").Int())
  369. require.Equal(t, "developer", gjson.GetBytes(upstreamBody, "input.0.role").String())
  370. require.Equal(t, systemPrompt, gjson.GetBytes(upstreamBody, "input.0.content").String())
  371. require.Equal(t, 2, strings.Count(string(upstreamBody), systemPrompt))
  372. }
  373. func TestForwardAsChatCompletions_OAuthKeepsMixedSystemContentInInput(t *testing.T) {
  374. const systemPrompt = "Inspect this reference image."
  375. body := []byte(`{"model":"gpt-5.4","messages":[{"role":"system","content":[{"type":"text","text":"` + systemPrompt + `"},{"type":"image_url","image_url":{"url":"https://example.com/reference.png"}}]},{"role":"user","content":"hello"}],"stream":false}`)
  376. upstreamBody := forwardOAuthChatCompletionsForUpstreamBody(t, body)
  377. require.Equal(t, systemPrompt, gjson.GetBytes(upstreamBody, "instructions").String())
  378. require.Equal(t, int64(2), gjson.GetBytes(upstreamBody, "input.#").Int())
  379. require.Equal(t, "developer", gjson.GetBytes(upstreamBody, "input.0.role").String())
  380. require.Equal(t, int64(2), gjson.GetBytes(upstreamBody, "input.0.content.#").Int())
  381. require.Equal(t, "input_image", gjson.GetBytes(upstreamBody, "input.0.content.1.type").String())
  382. require.Equal(t, "https://example.com/reference.png", gjson.GetBytes(upstreamBody, "input.0.content.1.image_url").String())
  383. }
  384. func TestForwardAsChatCompletions_ClientDisconnectDrainsUpstreamUsage(t *testing.T) {
  385. gin.SetMode(gin.TestMode)
  386. rec := httptest.NewRecorder()
  387. c, _ := gin.CreateTestContext(rec)
  388. c.Writer = &openAIChatFailingWriter{ResponseWriter: c.Writer, failAfter: 0}
  389. body := []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":"hello"}],"stream":true}`)
  390. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  391. c.Request.Header.Set("Content-Type", "application/json")
  392. upstreamBody := strings.Join([]string{
  393. `data: {"type":"response.created","response":{"id":"resp_1","model":"gpt-5.4","status":"in_progress","output":[]}}`,
  394. "",
  395. `data: {"type":"response.output_text.delta","delta":"ok"}`,
  396. "",
  397. `data: {"type":"response.completed","response":{"id":"resp_1","object":"response","model":"gpt-5.4","status":"completed","output":[{"type":"message","id":"msg_1","role":"assistant","status":"completed","content":[{"type":"output_text","text":"ok"}]}],"usage":{"input_tokens":11,"output_tokens":5,"total_tokens":16,"input_tokens_details":{"cached_tokens":4}}}}`,
  398. "",
  399. "data: [DONE]",
  400. "",
  401. }, "\n")
  402. upstream := &httpUpstreamRecorder{resp: &http.Response{
  403. StatusCode: http.StatusOK,
  404. Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"rid_chat_disconnect"}},
  405. Body: io.NopCloser(strings.NewReader(upstreamBody)),
  406. }}
  407. svc := &OpenAIGatewayService{httpUpstream: upstream}
  408. account := &Account{
  409. ID: 1,
  410. Name: "openai-oauth",
  411. Platform: PlatformOpenAI,
  412. Type: AccountTypeOAuth,
  413. Concurrency: 1,
  414. Credentials: map[string]any{
  415. "access_token": "oauth-token",
  416. "chatgpt_account_id": "chatgpt-acc",
  417. },
  418. }
  419. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "gpt-5.1")
  420. require.NoError(t, err)
  421. require.NotNil(t, result)
  422. require.Equal(t, 11, result.Usage.InputTokens)
  423. require.Equal(t, 5, result.Usage.OutputTokens)
  424. require.Equal(t, 4, result.Usage.CacheReadInputTokens)
  425. }
  426. func TestForwardAsChatCompletions_BufferedContextWindowResponseFailedReturnsErrorWithoutFailover(t *testing.T) {
  427. gin.SetMode(gin.TestMode)
  428. rec := httptest.NewRecorder()
  429. c, _ := gin.CreateTestContext(rec)
  430. body := []byte(`{"model":"gpt-5.5","messages":[{"role":"user","content":"large prompt"}],"stream":false}`)
  431. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  432. c.Request.Header.Set("Content-Type", "application/json")
  433. upstreamBody := strings.Join([]string{
  434. `event: response.failed`,
  435. `data: {"type":"response.failed","response":{"id":"resp_failed","object":"response","model":"gpt-5.5","status":"failed","output":[],"error":{"code":"upstream_error","message":"input exceeds the context window"}}}`,
  436. "",
  437. }, "\n")
  438. upstream := &httpUpstreamRecorder{resp: &http.Response{
  439. StatusCode: http.StatusOK,
  440. Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"rid_chat_failed_buffered"}},
  441. Body: io.NopCloser(strings.NewReader(upstreamBody)),
  442. }}
  443. svc := &OpenAIGatewayService{httpUpstream: upstream}
  444. account := &Account{
  445. ID: 1,
  446. Name: "openai-oauth",
  447. Platform: PlatformOpenAI,
  448. Type: AccountTypeOAuth,
  449. Concurrency: 1,
  450. Credentials: map[string]any{
  451. "access_token": "oauth-token",
  452. "chatgpt_account_id": "chatgpt-acc",
  453. },
  454. }
  455. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "gpt-5.5")
  456. require.Error(t, err)
  457. require.Nil(t, result)
  458. var failoverErr *UpstreamFailoverError
  459. require.False(t, errors.As(err, &failoverErr))
  460. require.True(t, c.Writer.Written())
  461. require.Equal(t, http.StatusBadGateway, rec.Code)
  462. require.Contains(t, rec.Body.String(), "input exceeds the context window")
  463. }
  464. func TestForwardAsChatCompletions_StreamContextWindowResponseFailedReturnsErrorWithoutFailover(t *testing.T) {
  465. gin.SetMode(gin.TestMode)
  466. rec := httptest.NewRecorder()
  467. c, _ := gin.CreateTestContext(rec)
  468. body := []byte(`{"model":"gpt-5.5","messages":[{"role":"user","content":"` + strings.Repeat("large prompt ", 6000) + `"}],"stream":true}`)
  469. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  470. c.Request.Header.Set("Content-Type", "application/json")
  471. upstreamBody := strings.Join([]string{
  472. `data: {"type":"response.created","response":{"id":"resp_failed","model":"gpt-5.5","status":"in_progress","output":[]}}`,
  473. "",
  474. `event: response.failed`,
  475. `data: {"type":"response.failed","response":{"id":"resp_failed","object":"response","model":"gpt-5.5","status":"failed","output":[],"error":{"code":"upstream_error","message":"input exceeds the context window"}}}`,
  476. "",
  477. }, "\n")
  478. upstream := &httpUpstreamRecorder{resp: &http.Response{
  479. StatusCode: http.StatusOK,
  480. Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"rid_chat_failed_stream"}},
  481. Body: io.NopCloser(strings.NewReader(upstreamBody)),
  482. }}
  483. svc := &OpenAIGatewayService{httpUpstream: upstream}
  484. account := &Account{
  485. ID: 1,
  486. Name: "openai-oauth",
  487. Platform: PlatformOpenAI,
  488. Type: AccountTypeOAuth,
  489. Concurrency: 1,
  490. Credentials: map[string]any{
  491. "access_token": "oauth-token",
  492. "chatgpt_account_id": "chatgpt-acc",
  493. },
  494. }
  495. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "gpt-5.5")
  496. require.Error(t, err)
  497. require.NotNil(t, result)
  498. var failoverErr *UpstreamFailoverError
  499. require.False(t, errors.As(err, &failoverErr))
  500. require.True(t, c.Writer.Written())
  501. require.Equal(t, http.StatusBadGateway, rec.Code)
  502. require.Contains(t, rec.Header().Get("Content-Type"), "application/json")
  503. require.Contains(t, rec.Body.String(), "input exceeds the context window")
  504. require.NotContains(t, rec.Body.String(), "[DONE]")
  505. }
  506. func TestForwardAsChatCompletions_StreamBareErrorAfterOutputDoesNotFailOver(t *testing.T) {
  507. gin.SetMode(gin.TestMode)
  508. rec := httptest.NewRecorder()
  509. c, _ := gin.CreateTestContext(rec)
  510. body := []byte(`{"model":"gpt-5.5","messages":[{"role":"user","content":"hello"}],"stream":true}`)
  511. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  512. c.Request.Header.Set("Content-Type", "application/json")
  513. upstreamBody := strings.Join([]string{
  514. `data: {"type":"response.output_text.delta","output_index":0,"content_index":0,"delta":"partial"}`,
  515. "",
  516. `event: error`,
  517. `data: {"type":"error","error":{"type":"server_error","code":"server_error","message":"temporary upstream failure"}}`,
  518. "",
  519. `data: [DONE]`,
  520. "",
  521. }, "\n")
  522. upstream := &httpUpstreamRecorder{resp: &http.Response{
  523. StatusCode: http.StatusOK,
  524. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  525. Body: io.NopCloser(strings.NewReader(upstreamBody)),
  526. }}
  527. svc := &OpenAIGatewayService{httpUpstream: upstream}
  528. account := &Account{
  529. ID: 1, Name: "openai-oauth", Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 1,
  530. Credentials: map[string]any{"access_token": "oauth-token", "chatgpt_account_id": "chatgpt-acc"},
  531. }
  532. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "gpt-5.5")
  533. require.Error(t, err)
  534. require.NotNil(t, result)
  535. var failoverErr *UpstreamFailoverError
  536. require.False(t, errors.As(err, &failoverErr))
  537. require.Contains(t, rec.Body.String(), "partial")
  538. require.Contains(t, rec.Body.String(), "temporary upstream failure")
  539. require.NotContains(t, rec.Body.String(), "[DONE]")
  540. }
  541. func TestForwardAsChatCompletions_StreamCyberPolicyNoFailover(t *testing.T) {
  542. gin.SetMode(gin.TestMode)
  543. rec := httptest.NewRecorder()
  544. c, _ := gin.CreateTestContext(rec)
  545. body := []byte(`{"model":"gpt-5.5","messages":[{"role":"user","content":"` + strings.Repeat("large prompt ", 6000) + `"}],"stream":true}`)
  546. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  547. c.Request.Header.Set("Content-Type", "application/json")
  548. upstreamBody := strings.Join([]string{
  549. `data: {"type":"response.created","response":{"id":"resp_cyber","model":"gpt-5.5","status":"in_progress","output":[]}}`,
  550. "",
  551. `event: response.failed`,
  552. `data: {"type":"response.failed","response":{"id":"resp_cyber","object":"response","model":"gpt-5.5","status":"failed","output":[],"error":{"code":"cyber_policy","message":"flagged for cyber policy"}}}`,
  553. "",
  554. }, "\n")
  555. upstream := &httpUpstreamRecorder{resp: &http.Response{
  556. StatusCode: http.StatusOK,
  557. Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"rid_chat_cyber"}},
  558. Body: io.NopCloser(strings.NewReader(upstreamBody)),
  559. }}
  560. svc := &OpenAIGatewayService{httpUpstream: upstream}
  561. account := &Account{
  562. ID: 1,
  563. Name: "openai-oauth",
  564. Platform: PlatformOpenAI,
  565. Type: AccountTypeOAuth,
  566. Concurrency: 1,
  567. Credentials: map[string]any{
  568. "access_token": "oauth-token",
  569. "chatgpt_account_id": "chatgpt-acc",
  570. },
  571. }
  572. _, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "gpt-5.5")
  573. var failoverErr *UpstreamFailoverError
  574. require.False(t, errors.As(err, &failoverErr), "cyber must NOT trigger failover")
  575. require.NotNil(t, GetOpsCyberPolicy(c), "cyber mark must be set")
  576. respBody := rec.Body.String()
  577. require.Contains(t, respBody, `"error"`)
  578. require.Contains(t, respBody, `"cyber_policy"`)
  579. require.Contains(t, respBody, "data: [DONE]")
  580. }
  581. func TestForwardAsChatCompletions_StreamsUsageWithoutClientStreamOptions(t *testing.T) {
  582. gin.SetMode(gin.TestMode)
  583. rec := httptest.NewRecorder()
  584. c, _ := gin.CreateTestContext(rec)
  585. body := []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":"hello"}],"stream":true}`)
  586. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  587. c.Request.Header.Set("Content-Type", "application/json")
  588. upstreamBody := strings.Join([]string{
  589. `data: {"type":"response.created","response":{"id":"resp_1","model":"gpt-5.4","status":"in_progress","output":[]}}`,
  590. "",
  591. `data: {"type":"response.output_text.delta","delta":"ok"}`,
  592. "",
  593. `data: {"type":"response.completed","response":{"id":"resp_1","object":"response","model":"gpt-5.4","status":"completed","output":[{"type":"message","id":"msg_1","role":"assistant","status":"completed","content":[{"type":"output_text","text":"ok"}]}],"usage":{"input_tokens":13,"output_tokens":7,"total_tokens":20,"input_tokens_details":{"cached_tokens":5}}}}`,
  594. "",
  595. "data: [DONE]",
  596. "",
  597. }, "\n")
  598. upstream := &httpUpstreamRecorder{resp: &http.Response{
  599. StatusCode: http.StatusOK,
  600. Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"rid_chat_usage_no_stream_options"}},
  601. Body: io.NopCloser(strings.NewReader(upstreamBody)),
  602. }}
  603. svc := &OpenAIGatewayService{httpUpstream: upstream}
  604. account := &Account{
  605. ID: 1,
  606. Name: "openai-oauth",
  607. Platform: PlatformOpenAI,
  608. Type: AccountTypeOAuth,
  609. Concurrency: 1,
  610. Credentials: map[string]any{
  611. "access_token": "oauth-token",
  612. "chatgpt_account_id": "chatgpt-acc",
  613. },
  614. }
  615. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "gpt-5.1")
  616. require.NoError(t, err)
  617. require.NotNil(t, result)
  618. require.Equal(t, 13, result.Usage.InputTokens)
  619. require.Equal(t, 7, result.Usage.OutputTokens)
  620. require.Equal(t, 5, result.Usage.CacheReadInputTokens)
  621. responseBody := rec.Body.String()
  622. require.Contains(t, responseBody, `"usage"`)
  623. require.Contains(t, responseBody, `"prompt_tokens":13`)
  624. require.Contains(t, responseBody, `"completion_tokens":7`)
  625. require.Contains(t, responseBody, `"cached_tokens":5`)
  626. }
  627. func TestForwardAsChatCompletions_StreamsTopLevelTerminalUsage(t *testing.T) {
  628. gin.SetMode(gin.TestMode)
  629. rec := httptest.NewRecorder()
  630. c, _ := gin.CreateTestContext(rec)
  631. body := []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":"hello"}],"stream":true}`)
  632. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  633. c.Request.Header.Set("Content-Type", "application/json")
  634. upstreamBody := strings.Join([]string{
  635. `data: {"type":"response.created","response":{"id":"resp_top","model":"gpt-5.4","status":"in_progress","output":[]}}`,
  636. "",
  637. `data: {"type":"response.output_text.delta","delta":"ok"}`,
  638. "",
  639. `data: {"type":"response.completed","response":{"id":"resp_top","object":"response","model":"gpt-5.4","status":"completed","output":[{"type":"message","id":"msg_1","role":"assistant","status":"completed","content":[{"type":"output_text","text":"ok"}]}]},"usage":{"input_tokens":21,"output_tokens":9,"total_tokens":30,"input_tokens_details":{"cached_tokens":4}}}`,
  640. "",
  641. "data: [DONE]",
  642. "",
  643. }, "\n")
  644. upstream := &httpUpstreamRecorder{resp: &http.Response{
  645. StatusCode: http.StatusOK,
  646. Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"rid_chat_top_level_usage"}},
  647. Body: io.NopCloser(strings.NewReader(upstreamBody)),
  648. }}
  649. svc := &OpenAIGatewayService{httpUpstream: upstream}
  650. account := &Account{
  651. ID: 1,
  652. Name: "openai-oauth",
  653. Platform: PlatformOpenAI,
  654. Type: AccountTypeOAuth,
  655. Concurrency: 1,
  656. Credentials: map[string]any{
  657. "access_token": "oauth-token",
  658. "chatgpt_account_id": "chatgpt-acc",
  659. },
  660. }
  661. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "gpt-5.1")
  662. require.NoError(t, err)
  663. require.NotNil(t, result)
  664. require.Equal(t, 21, result.Usage.InputTokens)
  665. require.Equal(t, 9, result.Usage.OutputTokens)
  666. require.Equal(t, 4, result.Usage.CacheReadInputTokens)
  667. responseBody := rec.Body.String()
  668. require.Contains(t, responseBody, `"usage"`)
  669. require.Contains(t, responseBody, `"prompt_tokens":21`)
  670. require.Contains(t, responseBody, `"completion_tokens":9`)
  671. require.Contains(t, responseBody, `"cached_tokens":4`)
  672. }
  673. func TestForwardAsChatCompletions_BufferedTopLevelTerminalUsage(t *testing.T) {
  674. gin.SetMode(gin.TestMode)
  675. rec := httptest.NewRecorder()
  676. c, _ := gin.CreateTestContext(rec)
  677. body := []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":"hello"}],"stream":false}`)
  678. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  679. c.Request.Header.Set("Content-Type", "application/json")
  680. upstreamBody := strings.Join([]string{
  681. `data: {"type":"response.completed","response":{"id":"resp_top_buffered","object":"response","model":"gpt-5.4","status":"completed","output":[{"type":"message","id":"msg_1","role":"assistant","status":"completed","content":[{"type":"output_text","text":"ok"}]}]},"usage":{"input_tokens":18,"output_tokens":6,"total_tokens":24,"input_tokens_details":{"cached_tokens":3}}}`,
  682. "",
  683. "data: [DONE]",
  684. "",
  685. }, "\n")
  686. upstream := &httpUpstreamRecorder{resp: &http.Response{
  687. StatusCode: http.StatusOK,
  688. Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"rid_chat_buffered_top_level_usage"}},
  689. Body: io.NopCloser(strings.NewReader(upstreamBody)),
  690. }}
  691. svc := &OpenAIGatewayService{httpUpstream: upstream}
  692. account := &Account{
  693. ID: 1,
  694. Name: "openai-oauth",
  695. Platform: PlatformOpenAI,
  696. Type: AccountTypeOAuth,
  697. Concurrency: 1,
  698. Credentials: map[string]any{
  699. "access_token": "oauth-token",
  700. "chatgpt_account_id": "chatgpt-acc",
  701. },
  702. }
  703. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "gpt-5.1")
  704. require.NoError(t, err)
  705. require.NotNil(t, result)
  706. require.Equal(t, 18, result.Usage.InputTokens)
  707. require.Equal(t, 6, result.Usage.OutputTokens)
  708. require.Equal(t, 3, result.Usage.CacheReadInputTokens)
  709. responseBody := rec.Body.String()
  710. require.Contains(t, responseBody, `"usage"`)
  711. require.Contains(t, responseBody, `"prompt_tokens":18`)
  712. require.Contains(t, responseBody, `"completion_tokens":6`)
  713. require.Contains(t, responseBody, `"cached_tokens":3`)
  714. }
  715. func TestForwardAsChatCompletions_TerminalUsageWithoutUpstreamCloseReturns(t *testing.T) {
  716. gin.SetMode(gin.TestMode)
  717. rec := httptest.NewRecorder()
  718. c, _ := gin.CreateTestContext(rec)
  719. c.Writer = &openAIChatFailingWriter{ResponseWriter: c.Writer, failAfter: 0}
  720. body := []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":"hello"}],"stream":true}`)
  721. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  722. c.Request.Header.Set("Content-Type", "application/json")
  723. upstreamBody := []byte(`data: {"type":"response.completed","response":{"id":"resp_1","object":"response","model":"gpt-5.4","status":"completed","output":[{"type":"message","id":"msg_1","role":"assistant","status":"completed","content":[{"type":"output_text","text":"ok"}]}],"usage":{"input_tokens":17,"output_tokens":8,"total_tokens":25,"input_tokens_details":{"cached_tokens":6}}}}` + "\n\n")
  724. upstreamStream := newOpenAICompatBlockingReadCloser(upstreamBody)
  725. defer func() {
  726. require.NoError(t, upstreamStream.Close())
  727. }()
  728. upstream := &httpUpstreamRecorder{resp: &http.Response{
  729. StatusCode: http.StatusOK,
  730. Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"rid_chat_terminal_no_close"}},
  731. Body: upstreamStream,
  732. }}
  733. svc := &OpenAIGatewayService{httpUpstream: upstream}
  734. account := &Account{
  735. ID: 1,
  736. Name: "openai-oauth",
  737. Platform: PlatformOpenAI,
  738. Type: AccountTypeOAuth,
  739. Concurrency: 1,
  740. Credentials: map[string]any{
  741. "access_token": "oauth-token",
  742. "chatgpt_account_id": "chatgpt-acc",
  743. },
  744. }
  745. type forwardResult struct {
  746. result *OpenAIForwardResult
  747. err error
  748. }
  749. resultCh := make(chan forwardResult, 1)
  750. go func() {
  751. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "gpt-5.1")
  752. resultCh <- forwardResult{result: result, err: err}
  753. }()
  754. select {
  755. case got := <-resultCh:
  756. require.NoError(t, got.err)
  757. require.NotNil(t, got.result)
  758. require.Equal(t, 17, got.result.Usage.InputTokens)
  759. require.Equal(t, 8, got.result.Usage.OutputTokens)
  760. require.Equal(t, 6, got.result.Usage.CacheReadInputTokens)
  761. case <-time.After(time.Second):
  762. require.Fail(t, "ForwardAsChatCompletions should return after terminal usage event even if upstream keeps the connection open")
  763. }
  764. }
  765. func TestForwardAsChatCompletions_EventNamedTerminalWithoutUpstreamCloseReturns(t *testing.T) {
  766. gin.SetMode(gin.TestMode)
  767. rec := httptest.NewRecorder()
  768. c, _ := gin.CreateTestContext(rec)
  769. body := []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":"hello"}],"stream":true}`)
  770. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  771. c.Request.Header.Set("Content-Type", "application/json")
  772. upstreamBody := []byte(strings.Join([]string{
  773. `event: response.created`,
  774. `data: {"response":{"id":"resp_1","model":"gpt-5.4","status":"in_progress","output":[]}}`,
  775. ``,
  776. `event: response.output_text.delta`,
  777. `data: {"delta":"ok"}`,
  778. ``,
  779. `event: response.completed`,
  780. `data: {"response":{"id":"resp_1","object":"response","model":"gpt-5.4","status":"completed","output":[{"type":"message","id":"msg_1","role":"assistant","status":"completed","content":[{"type":"output_text","text":"ok"}]}],"usage":{"input_tokens":17,"output_tokens":8,"total_tokens":25,"input_tokens_details":{"cached_tokens":6}}}}`,
  781. ``,
  782. ``,
  783. }, "\n"))
  784. upstreamStream := newOpenAICompatBlockingReadCloser(upstreamBody)
  785. defer func() {
  786. require.NoError(t, upstreamStream.Close())
  787. }()
  788. upstream := &httpUpstreamRecorder{resp: &http.Response{
  789. StatusCode: http.StatusOK,
  790. Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"rid_chat_event_named_terminal"}},
  791. Body: upstreamStream,
  792. }}
  793. svc := &OpenAIGatewayService{httpUpstream: upstream}
  794. account := &Account{
  795. ID: 1,
  796. Name: "openai-oauth",
  797. Platform: PlatformOpenAI,
  798. Type: AccountTypeOAuth,
  799. Concurrency: 1,
  800. Credentials: map[string]any{
  801. "access_token": "oauth-token",
  802. "chatgpt_account_id": "chatgpt-acc",
  803. },
  804. }
  805. type forwardResult struct {
  806. result *OpenAIForwardResult
  807. err error
  808. }
  809. resultCh := make(chan forwardResult, 1)
  810. go func() {
  811. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "gpt-5.1")
  812. resultCh <- forwardResult{result: result, err: err}
  813. }()
  814. select {
  815. case got := <-resultCh:
  816. require.NoError(t, got.err)
  817. require.NotNil(t, got.result)
  818. require.Equal(t, 17, got.result.Usage.InputTokens)
  819. require.Equal(t, 8, got.result.Usage.OutputTokens)
  820. require.Equal(t, 6, got.result.Usage.CacheReadInputTokens)
  821. require.Contains(t, rec.Body.String(), `"content":"ok"`)
  822. case <-time.After(time.Second):
  823. require.Fail(t, "ForwardAsChatCompletions should use SSE event names when data payloads omit type")
  824. }
  825. }
  826. func TestForwardAsChatCompletions_EventTypeDoesNotLeakAcrossFrames(t *testing.T) {
  827. gin.SetMode(gin.TestMode)
  828. rec := httptest.NewRecorder()
  829. c, _ := gin.CreateTestContext(rec)
  830. body := []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":"hello"}],"stream":true}`)
  831. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  832. c.Request.Header.Set("Content-Type", "application/json")
  833. upstreamBody := strings.Join([]string{
  834. `event: response.created`,
  835. `data: {"response":{"id":"resp_1","model":"gpt-5.4","status":"in_progress","output":[]}}`,
  836. ``,
  837. `data: {"type":"response.output_text.delta","delta":"ok"}`,
  838. ``,
  839. `event: response.completed`,
  840. `data: {"response":{"id":"resp_1","object":"response","model":"gpt-5.4","status":"completed","output":[{"type":"message","id":"msg_1","role":"assistant","status":"completed","content":[{"type":"output_text","text":"ok"}]}],"usage":{"input_tokens":17,"output_tokens":8,"total_tokens":25,"input_tokens_details":{"cached_tokens":6}}}}`,
  841. ``,
  842. `data: [DONE]`,
  843. ``,
  844. }, "\n")
  845. upstream := &httpUpstreamRecorder{resp: &http.Response{
  846. StatusCode: http.StatusOK,
  847. Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"rid_chat_event_boundary"}},
  848. Body: io.NopCloser(strings.NewReader(upstreamBody)),
  849. }}
  850. svc := &OpenAIGatewayService{httpUpstream: upstream}
  851. account := &Account{
  852. ID: 1,
  853. Name: "openai-oauth",
  854. Platform: PlatformOpenAI,
  855. Type: AccountTypeOAuth,
  856. Concurrency: 1,
  857. Credentials: map[string]any{
  858. "access_token": "oauth-token",
  859. "chatgpt_account_id": "chatgpt-acc",
  860. },
  861. }
  862. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "gpt-5.1")
  863. require.NoError(t, err)
  864. require.NotNil(t, result)
  865. require.Contains(t, rec.Body.String(), `"content":"ok"`)
  866. require.Contains(t, rec.Body.String(), `data: [DONE]`)
  867. }
  868. func TestForwardAsChatCompletions_BufferedTerminalWithoutUpstreamCloseReturns(t *testing.T) {
  869. gin.SetMode(gin.TestMode)
  870. rec := httptest.NewRecorder()
  871. c, _ := gin.CreateTestContext(rec)
  872. body := []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":"hello"}],"stream":false}`)
  873. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  874. c.Request.Header.Set("Content-Type", "application/json")
  875. upstreamBody := []byte(`data: {"type":"response.completed","response":{"id":"resp_1","object":"response","model":"gpt-5.4","status":"completed","output":[{"type":"message","id":"msg_1","role":"assistant","status":"completed","content":[{"type":"output_text","text":"ok"}]}],"usage":{"input_tokens":17,"output_tokens":8,"total_tokens":25,"input_tokens_details":{"cached_tokens":6}}}}` + "\n\n")
  876. upstreamStream := newOpenAICompatBlockingReadCloser(upstreamBody)
  877. defer func() {
  878. require.NoError(t, upstreamStream.Close())
  879. }()
  880. upstream := &httpUpstreamRecorder{resp: &http.Response{
  881. StatusCode: http.StatusOK,
  882. Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"rid_chat_buffered_terminal_no_close"}},
  883. Body: upstreamStream,
  884. }}
  885. svc := &OpenAIGatewayService{httpUpstream: upstream}
  886. account := &Account{
  887. ID: 1,
  888. Name: "openai-oauth",
  889. Platform: PlatformOpenAI,
  890. Type: AccountTypeOAuth,
  891. Concurrency: 1,
  892. Credentials: map[string]any{
  893. "access_token": "oauth-token",
  894. "chatgpt_account_id": "chatgpt-acc",
  895. },
  896. }
  897. type forwardResult struct {
  898. result *OpenAIForwardResult
  899. err error
  900. }
  901. resultCh := make(chan forwardResult, 1)
  902. go func() {
  903. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "gpt-5.1")
  904. resultCh <- forwardResult{result: result, err: err}
  905. }()
  906. select {
  907. case got := <-resultCh:
  908. require.NoError(t, got.err)
  909. require.NotNil(t, got.result)
  910. require.Equal(t, 17, got.result.Usage.InputTokens)
  911. require.Equal(t, 8, got.result.Usage.OutputTokens)
  912. require.Equal(t, 6, got.result.Usage.CacheReadInputTokens)
  913. require.Contains(t, rec.Body.String(), `"finish_reason":"stop"`)
  914. case <-time.After(time.Second):
  915. require.Fail(t, "ForwardAsChatCompletions buffered response should return after terminal usage event even if upstream keeps the connection open")
  916. }
  917. }
  918. func TestForwardAsChatCompletions_DoneSentinelWithoutTerminalReturnsError(t *testing.T) {
  919. gin.SetMode(gin.TestMode)
  920. rec := httptest.NewRecorder()
  921. c, _ := gin.CreateTestContext(rec)
  922. body := []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":"hello"}],"stream":true}`)
  923. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  924. c.Request.Header.Set("Content-Type", "application/json")
  925. upstreamBody := "data: [DONE]\n\n"
  926. upstream := &httpUpstreamRecorder{resp: &http.Response{
  927. StatusCode: http.StatusOK,
  928. Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"rid_chat_missing_terminal"}},
  929. Body: io.NopCloser(strings.NewReader(upstreamBody)),
  930. }}
  931. svc := &OpenAIGatewayService{httpUpstream: upstream}
  932. account := &Account{
  933. ID: 1,
  934. Name: "openai-oauth",
  935. Platform: PlatformOpenAI,
  936. Type: AccountTypeOAuth,
  937. Concurrency: 1,
  938. Credentials: map[string]any{
  939. "access_token": "oauth-token",
  940. "chatgpt_account_id": "chatgpt-acc",
  941. },
  942. }
  943. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "gpt-5.1")
  944. require.Error(t, err)
  945. require.Contains(t, err.Error(), "missing terminal event")
  946. require.NotNil(t, result)
  947. require.Zero(t, result.Usage.InputTokens)
  948. require.Zero(t, result.Usage.OutputTokens)
  949. }
  950. func TestForwardAsChatCompletions_UpstreamRequestIgnoresClientCancel(t *testing.T) {
  951. gin.SetMode(gin.TestMode)
  952. rec := httptest.NewRecorder()
  953. c, _ := gin.CreateTestContext(rec)
  954. reqCtx, cancel := context.WithCancel(context.Background())
  955. body := []byte(`{"model":"gpt-5.4","messages":[{"role":"user","content":"hello"}],"stream":false}`)
  956. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)).WithContext(reqCtx)
  957. c.Request.Header.Set("Content-Type", "application/json")
  958. cancel()
  959. upstreamBody := strings.Join([]string{
  960. `data: {"type":"response.completed","response":{"id":"resp_1","object":"response","model":"gpt-5.4","status":"completed","output":[{"type":"message","id":"msg_1","role":"assistant","status":"completed","content":[{"type":"output_text","text":"ok"}]}],"usage":{"input_tokens":5,"output_tokens":2,"total_tokens":7}}}`,
  961. "",
  962. "data: [DONE]",
  963. "",
  964. }, "\n")
  965. upstream := &httpUpstreamRecorder{resp: &http.Response{
  966. StatusCode: http.StatusOK,
  967. Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"rid_chat_ctx"}},
  968. Body: io.NopCloser(strings.NewReader(upstreamBody)),
  969. }}
  970. svc := &OpenAIGatewayService{httpUpstream: upstream}
  971. account := &Account{
  972. ID: 1,
  973. Name: "openai-oauth",
  974. Platform: PlatformOpenAI,
  975. Type: AccountTypeOAuth,
  976. Concurrency: 1,
  977. Credentials: map[string]any{
  978. "access_token": "oauth-token",
  979. "chatgpt_account_id": "chatgpt-acc",
  980. },
  981. }
  982. result, err := svc.ForwardAsChatCompletions(reqCtx, c, account, body, "", "gpt-5.1")
  983. require.NoError(t, err)
  984. require.NotNil(t, result)
  985. require.NotNil(t, upstream.lastReq)
  986. require.NoError(t, upstream.lastReq.Context().Err())
  987. }
  988. // TestBuildChatStreamErrorSSE verifies F4: the error chunk payload follows the
  989. // OpenAI chat streaming error convention so third-party clients stop retrying.
  990. func TestBuildChatStreamErrorSSE(t *testing.T) {
  991. got := buildChatStreamErrorSSE("cyber_policy", "blocked by policy")
  992. require.True(t, strings.HasPrefix(got, "data: "), "must be an SSE data frame")
  993. payload := strings.TrimSuffix(strings.TrimPrefix(got, "data: "), "\n\n")
  994. require.Equal(t, "invalid_request_error", gjson.Get(payload, "error.type").String())
  995. require.Equal(t, "cyber_policy", gjson.Get(payload, "error.code").String())
  996. require.Equal(t, "blocked by policy", gjson.Get(payload, "error.message").String())
  997. }