passthrough_relay_internal_test.go 26 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797
  1. package openai_ws_v2
  2. import (
  3. "context"
  4. "errors"
  5. "io"
  6. "net"
  7. "sync/atomic"
  8. "testing"
  9. "time"
  10. coderws "github.com/coder/websocket"
  11. "github.com/stretchr/testify/require"
  12. "github.com/tidwall/gjson"
  13. )
  14. func TestRunEntry_DelegatesRelay(t *testing.T) {
  15. t.Parallel()
  16. clientConn := newPassthroughTestFrameConn(nil, false)
  17. upstreamConn := newPassthroughTestFrameConn([]passthroughTestFrame{
  18. {
  19. msgType: coderws.MessageText,
  20. payload: []byte(`{"type":"response.completed","response":{"id":"resp_entry","usage":{"input_tokens":1,"output_tokens":1}}}`),
  21. },
  22. }, true)
  23. result, relayExit := RunEntry(EntryInput{
  24. Ctx: context.Background(),
  25. ClientConn: clientConn,
  26. UpstreamConn: upstreamConn,
  27. FirstClientMessage: []byte(`{"type":"response.create","model":"gpt-4o","input":[]}`),
  28. })
  29. require.Nil(t, relayExit)
  30. require.Equal(t, "resp_entry", result.RequestID)
  31. }
  32. func TestRunClientToUpstream_ErrorPaths(t *testing.T) {
  33. t.Parallel()
  34. t.Run("read client eof", func(t *testing.T) {
  35. t.Parallel()
  36. exitCh := make(chan relayExitSignal, 1)
  37. runClientToUpstream(
  38. context.Background(),
  39. newPassthroughTestFrameConn(nil, true),
  40. nil,
  41. func(_ coderws.MessageType, _ []byte) error { return nil },
  42. func() {},
  43. nil,
  44. nil,
  45. exitCh,
  46. )
  47. sig := <-exitCh
  48. require.Equal(t, "read_client", sig.stage)
  49. require.True(t, sig.graceful)
  50. })
  51. t.Run("write upstream failed", func(t *testing.T) {
  52. t.Parallel()
  53. exitCh := make(chan relayExitSignal, 1)
  54. runClientToUpstream(
  55. context.Background(),
  56. newPassthroughTestFrameConn([]passthroughTestFrame{
  57. {msgType: coderws.MessageText, payload: []byte(`{"x":1}`)},
  58. }, true),
  59. nil,
  60. func(_ coderws.MessageType, _ []byte) error { return errors.New("boom") },
  61. func() {},
  62. nil,
  63. nil,
  64. exitCh,
  65. )
  66. sig := <-exitCh
  67. require.Equal(t, "write_upstream", sig.stage)
  68. require.False(t, sig.graceful)
  69. })
  70. t.Run("forwarded counter and trace callback", func(t *testing.T) {
  71. t.Parallel()
  72. exitCh := make(chan relayExitSignal, 1)
  73. forwarded := &atomic.Int64{}
  74. traces := make([]RelayTraceEvent, 0, 2)
  75. runClientToUpstream(
  76. context.Background(),
  77. newPassthroughTestFrameConn([]passthroughTestFrame{
  78. {msgType: coderws.MessageText, payload: []byte(`{"x":1}`)},
  79. }, true),
  80. nil,
  81. func(_ coderws.MessageType, _ []byte) error { return nil },
  82. func() {},
  83. forwarded,
  84. func(event RelayTraceEvent) {
  85. traces = append(traces, event)
  86. },
  87. exitCh,
  88. )
  89. sig := <-exitCh
  90. require.Equal(t, "read_client", sig.stage)
  91. require.Equal(t, int64(1), forwarded.Load())
  92. require.NotEmpty(t, traces)
  93. })
  94. }
  95. func TestRunUpstreamToClient_ErrorAndDropPaths(t *testing.T) {
  96. t.Parallel()
  97. t.Run("read upstream eof", func(t *testing.T) {
  98. t.Parallel()
  99. exitCh := make(chan relayExitSignal, 1)
  100. drop := &atomic.Bool{}
  101. drop.Store(false)
  102. runUpstreamToClient(
  103. context.Background(),
  104. newPassthroughTestFrameConn(nil, true),
  105. func(_ coderws.MessageType, _ []byte) error { return nil },
  106. time.Now(),
  107. time.Now,
  108. &relayState{},
  109. nil,
  110. nil,
  111. nil,
  112. nil,
  113. nil,
  114. nil,
  115. drop,
  116. nil,
  117. nil,
  118. func() {},
  119. nil,
  120. exitCh,
  121. )
  122. sig := <-exitCh
  123. require.Equal(t, "read_upstream", sig.stage)
  124. require.True(t, sig.graceful)
  125. })
  126. t.Run("write client failed", func(t *testing.T) {
  127. t.Parallel()
  128. exitCh := make(chan relayExitSignal, 1)
  129. drop := &atomic.Bool{}
  130. drop.Store(false)
  131. runUpstreamToClient(
  132. context.Background(),
  133. newPassthroughTestFrameConn([]passthroughTestFrame{
  134. {msgType: coderws.MessageText, payload: []byte(`{"type":"response.output_text.delta","delta":"x"}`)},
  135. }, true),
  136. func(_ coderws.MessageType, _ []byte) error { return errors.New("write failed") },
  137. time.Now(),
  138. time.Now,
  139. &relayState{},
  140. nil,
  141. nil,
  142. nil,
  143. nil,
  144. nil,
  145. nil,
  146. drop,
  147. nil,
  148. nil,
  149. func() {},
  150. nil,
  151. exitCh,
  152. )
  153. sig := <-exitCh
  154. require.Equal(t, "write_client", sig.stage)
  155. })
  156. t.Run("drop downstream and stop on terminal", func(t *testing.T) {
  157. t.Parallel()
  158. exitCh := make(chan relayExitSignal, 1)
  159. drop := &atomic.Bool{}
  160. drop.Store(true)
  161. dropped := &atomic.Int64{}
  162. runUpstreamToClient(
  163. context.Background(),
  164. newPassthroughTestFrameConn([]passthroughTestFrame{
  165. {
  166. msgType: coderws.MessageText,
  167. payload: []byte(`{"type":"response.completed","response":{"id":"resp_drop","usage":{"input_tokens":1,"output_tokens":1}}}`),
  168. },
  169. }, true),
  170. func(_ coderws.MessageType, _ []byte) error { return nil },
  171. time.Now(),
  172. time.Now,
  173. &relayState{},
  174. nil,
  175. nil,
  176. nil,
  177. nil,
  178. nil,
  179. nil,
  180. drop,
  181. nil,
  182. dropped,
  183. func() {},
  184. nil,
  185. exitCh,
  186. )
  187. sig := <-exitCh
  188. require.Equal(t, "drain_terminal", sig.stage)
  189. require.True(t, sig.graceful)
  190. require.Equal(t, int64(1), dropped.Load())
  191. })
  192. }
  193. func TestRunIdleWatchdog_NoTimeoutWhenDisabled(t *testing.T) {
  194. t.Parallel()
  195. exitCh := make(chan relayExitSignal, 1)
  196. lastActivity := &atomic.Int64{}
  197. lastActivity.Store(time.Now().UnixNano())
  198. ctx, cancel := context.WithCancel(context.Background())
  199. defer cancel()
  200. go runIdleWatchdog(ctx, time.Now, 0, lastActivity, nil, exitCh)
  201. select {
  202. case <-exitCh:
  203. t.Fatal("unexpected idle timeout signal")
  204. case <-time.After(200 * time.Millisecond):
  205. }
  206. }
  207. func TestHelperFunctionsCoverage(t *testing.T) {
  208. t.Parallel()
  209. require.Equal(t, "text", relayMessageTypeString(coderws.MessageText))
  210. require.Equal(t, "binary", relayMessageTypeString(coderws.MessageBinary))
  211. require.Contains(t, relayMessageTypeString(coderws.MessageType(99)), "unknown(")
  212. require.Equal(t, "", relayErrorString(nil))
  213. require.Equal(t, "x", relayErrorString(errors.New("x")))
  214. require.True(t, isDisconnectError(io.EOF))
  215. require.True(t, isDisconnectError(net.ErrClosed))
  216. require.True(t, isDisconnectError(context.Canceled))
  217. require.True(t, isDisconnectError(coderws.CloseError{Code: coderws.StatusGoingAway}))
  218. require.True(t, isDisconnectError(errors.New("broken pipe")))
  219. require.False(t, isDisconnectError(errors.New("unrelated")))
  220. require.True(t, isTokenEvent("response.output_text.delta"))
  221. require.True(t, isTokenEvent("response.output_audio.delta"))
  222. require.False(t, isTokenEvent("response.completed"))
  223. require.False(t, isTokenEvent(""))
  224. require.False(t, isTokenEvent("response.created"))
  225. require.Equal(t, 2*time.Second, minDuration(2*time.Second, 5*time.Second))
  226. require.Equal(t, 2*time.Second, minDuration(5*time.Second, 2*time.Second))
  227. require.Equal(t, 5*time.Second, minDuration(0, 5*time.Second))
  228. require.Equal(t, 2*time.Second, minDuration(2*time.Second, 0))
  229. ch := make(chan relayExitSignal, 1)
  230. ch <- relayExitSignal{stage: "ok"}
  231. sig, ok := waitRelayExit(ch, 10*time.Millisecond)
  232. require.True(t, ok)
  233. require.Equal(t, "ok", sig.stage)
  234. ch <- relayExitSignal{stage: "ok2"}
  235. sig, ok = waitRelayExit(ch, 0)
  236. require.True(t, ok)
  237. require.Equal(t, "ok2", sig.stage)
  238. _, ok = waitRelayExit(ch, 10*time.Millisecond)
  239. require.False(t, ok)
  240. n, ok := parseUsageIntField(gjson.Get(`{"n":3}`, "n"), true)
  241. require.True(t, ok)
  242. require.Equal(t, 3, n)
  243. _, ok = parseUsageIntField(gjson.Get(`{"n":"x"}`, "n"), true)
  244. require.False(t, ok)
  245. n, ok = parseUsageIntField(gjson.Result{}, false)
  246. require.True(t, ok)
  247. require.Equal(t, 0, n)
  248. _, ok = parseUsageIntField(gjson.Result{}, true)
  249. require.False(t, ok)
  250. }
  251. func TestParseUsageAndEnrichCoverage(t *testing.T) {
  252. t.Parallel()
  253. state := &relayState{}
  254. parseUsageAndAccumulate(state, []byte(`{"type":"response.completed","response":{"usage":{"input_tokens":"bad"}}}`), "response.completed", nil)
  255. require.Equal(t, 0, state.usage.InputTokens)
  256. parseUsageAndAccumulate(
  257. state,
  258. []byte(`{"type":"response.completed","response":{"usage":{"input_tokens":9,"output_tokens":"bad","input_tokens_details":{"cached_tokens":2}}}}`),
  259. "response.completed",
  260. nil,
  261. )
  262. require.Equal(t, 0, state.usage.InputTokens, "部分字段解析失败时不应累加 usage")
  263. require.Equal(t, 0, state.usage.OutputTokens)
  264. require.Equal(t, 0, state.usage.CacheReadInputTokens)
  265. parseUsageAndAccumulate(
  266. state,
  267. []byte(`{"type":"response.completed","response":{"usage":{"input_tokens_details":{"cached_tokens":2}}}}`),
  268. "response.completed",
  269. nil,
  270. )
  271. require.Equal(t, 0, state.usage.InputTokens, "必填 usage 字段缺失时不应累加 usage")
  272. require.Equal(t, 0, state.usage.OutputTokens)
  273. require.Equal(t, 0, state.usage.CacheReadInputTokens)
  274. parseUsageAndAccumulate(state, []byte(`{"type":"response.completed","response":{"usage":{"input_tokens":2,"output_tokens":1,"input_tokens_details":{"cached_tokens":1,"cache_write_tokens":4},"output_tokens_details":{"image_tokens":3}}}}`), "response.completed", nil)
  275. finalizeRelayTurnUsage(state)
  276. require.Equal(t, 2, state.usage.InputTokens)
  277. require.Equal(t, 1, state.usage.OutputTokens)
  278. require.Equal(t, 1, state.usage.CacheReadInputTokens)
  279. require.Equal(t, 4, state.usage.CacheCreationInputTokens)
  280. require.Equal(t, 3, state.usage.ImageOutputTokens)
  281. result := &RelayResult{}
  282. enrichResult(result, state, 5*time.Millisecond)
  283. require.Equal(t, state.usage.InputTokens, result.Usage.InputTokens)
  284. require.Equal(t, state.usage.CacheCreationInputTokens, result.Usage.CacheCreationInputTokens)
  285. require.Equal(t, state.usage.ImageOutputTokens, result.Usage.ImageOutputTokens)
  286. require.Equal(t, 5*time.Millisecond, result.Duration)
  287. parseUsageAndAccumulate(state, []byte(`{"type":"response.in_progress","response":{"usage":{"input_tokens":9}}}`), "response.in_progress", nil)
  288. require.Equal(t, 2, state.usage.InputTokens)
  289. enrichResult(nil, state, 0)
  290. }
  291. func TestParseUsageAndAccumulateIncludesIndependentReasoningTokens(t *testing.T) {
  292. t.Parallel()
  293. state := &relayState{}
  294. got := parseUsageAndAccumulate(
  295. state,
  296. []byte(`{"type":"response.completed","response":{"usage":{"input_tokens":32,"output_tokens":9,"total_tokens":151,"output_tokens_details":{"reasoning_tokens":110}}}}`),
  297. "response.completed",
  298. nil,
  299. )
  300. require.Equal(t, 32, got.InputTokens)
  301. require.Equal(t, 119, got.OutputTokens)
  302. state = &relayState{}
  303. got = parseUsageAndAccumulate(
  304. state,
  305. []byte(`{"type":"response.completed","response":{"usage":{"input_tokens":32,"output_tokens":119,"total_tokens":151,"output_tokens_details":{"reasoning_tokens":110}}}}`),
  306. "response.completed",
  307. nil,
  308. )
  309. require.Equal(t, 119, got.OutputTokens, "inclusive Responses output must not double-count reasoning")
  310. }
  311. func TestParseUsageAndAccumulateAcceptsChatUsageAliases(t *testing.T) {
  312. t.Parallel()
  313. state := &relayState{}
  314. got := parseUsageAndAccumulate(
  315. state,
  316. []byte(`{"type":"response.done","response":{"usage":{"prompt_tokens":12,"completion_tokens":6,"prompt_tokens_details":{"cached_tokens":4},"completion_tokens_details":{"image_tokens":2}}}}`),
  317. "response.done",
  318. nil,
  319. )
  320. finalizeRelayTurnUsage(state)
  321. require.Equal(t, 12, got.InputTokens)
  322. require.Equal(t, 6, got.OutputTokens)
  323. require.Equal(t, 4, got.CacheReadInputTokens)
  324. require.Equal(t, 2, got.ImageOutputTokens)
  325. require.Equal(t, got, state.usage)
  326. }
  327. func TestRelayUsageTerminalWithoutUsageKeepsFallback(t *testing.T) {
  328. t.Parallel()
  329. state := &relayState{}
  330. parseUsageAndAccumulate(state, []byte(`{"type":"response.in_progress","usage":{"input_tokens":9,"output_tokens":4}}`), "response.in_progress", nil)
  331. parseUsageAndAccumulate(state, []byte(`{"type":"response.completed","response":{"id":"resp_1"}}`), "response.completed", nil)
  332. turnUsage := finalizeRelayTurnUsage(state)
  333. require.Equal(t, Usage{InputTokens: 9, OutputTokens: 4}, turnUsage)
  334. require.Equal(t, turnUsage, state.usage)
  335. }
  336. func TestRelayUsageTerminalZeroKeepsFallback(t *testing.T) {
  337. t.Parallel()
  338. state := &relayState{}
  339. parseUsageAndAccumulate(state, []byte(`{"type":"response.in_progress","usage":{"input_tokens":9,"output_tokens":4}}`), "response.in_progress", nil)
  340. parseUsageAndAccumulate(state, []byte(`{"type":"response.completed","usage":{"input_tokens":0,"output_tokens":0}}`), "response.completed", nil)
  341. require.Equal(t, Usage{InputTokens: 9, OutputTokens: 4}, finalizeRelayTurnUsage(state))
  342. require.Equal(t, Usage{InputTokens: 9, OutputTokens: 4}, state.usage)
  343. }
  344. func TestRelayUsageTerminalNonZeroReplacesFallbackAsWhole(t *testing.T) {
  345. t.Parallel()
  346. state := &relayState{}
  347. parseUsageAndAccumulate(state, []byte(`{"type":"response.in_progress","usage":{"input_tokens":9,"output_tokens":4,"input_tokens_details":{"cached_tokens":2}}}`), "response.in_progress", nil)
  348. parseUsageAndAccumulate(state, []byte(`{"type":"response.completed","usage":{"input_tokens":3,"output_tokens":0,"input_tokens_details":{"cached_tokens":0}}}`), "response.completed", nil)
  349. require.Equal(t, Usage{InputTokens: 3}, finalizeRelayTurnUsage(state))
  350. require.Equal(t, Usage{InputTokens: 3}, state.usage)
  351. }
  352. func TestObserveUpstreamMessageBareErrorClearsTurnStateAndFinalizesUsageOnce(t *testing.T) {
  353. t.Parallel()
  354. now := time.Unix(200, 0)
  355. state := &relayState{}
  356. state.setPendingTurnStartedAt(now.Add(-2 * time.Second))
  357. timing := openAIWSRelayGetOrInitTurnTiming(state, "resp_active", now.Add(-time.Second))
  358. require.NotNil(t, timing)
  359. state.setPendingTurnStartedAt(now.Add(-500 * time.Millisecond))
  360. parseUsageAndAccumulate(state, []byte(`{"type":"response.in_progress","usage":{"input_tokens":9,"output_tokens":4}}`), "response.in_progress", nil)
  361. observed := observeUpstreamMessage(
  362. state,
  363. []byte(`{"type":"error","usage":{"input_tokens":3,"output_tokens":1},"error":{"message":"failed"}}`),
  364. now.Add(-3*time.Second),
  365. func() time.Time { return now },
  366. nil,
  367. )
  368. require.False(t, observed.terminal, "bare error settlement is deferred in case response.failed follows")
  369. require.Equal(t, "error", observed.eventType)
  370. observed = finalizePendingBareError(state, now)
  371. require.True(t, observed.terminal)
  372. require.Equal(t, Usage{InputTokens: 3, OutputTokens: 1}, observed.usage)
  373. require.Equal(t, observed.usage, state.usage)
  374. require.Equal(t, Usage{}, state.turnUsage)
  375. require.Nil(t, state.activeTurn)
  376. require.Empty(t, state.turnTimingByID)
  377. // The error is bound to the active turn id, so settlement must not consume a
  378. // later pending start that belongs to the next turn.
  379. require.Equal(t, now.Add(-500*time.Millisecond), state.consumePendingTurnStartedAt())
  380. // Re-observing and settling a later terminal without usage must not re-add the prior turn.
  381. observeUpstreamMessage(state, []byte(`{"type":"error","error":{"message":"again"}}`), now, func() time.Time { return now }, nil)
  382. finalizePendingBareError(state, now)
  383. require.Equal(t, Usage{InputTokens: 3, OutputTokens: 1}, state.usage)
  384. }
  385. func TestObserveUpstreamMessageErrorThenFailedSettlesUsageOnce(t *testing.T) {
  386. t.Parallel()
  387. now := time.Unix(300, 0)
  388. state := &relayState{}
  389. state.setPendingTurnStartedAt(now.Add(-time.Second))
  390. openAIWSRelayGetOrInitTurnTiming(state, "resp_1", now)
  391. parseUsageAndAccumulate(state, []byte(`{"type":"response.in_progress","usage":{"input_tokens":9,"output_tokens":4}}`), "response.in_progress", nil)
  392. errorObserved := observeUpstreamMessage(
  393. state,
  394. []byte(`{"type":"error","error":{"code":"server_is_overloaded","message":"overloaded"}}`),
  395. now,
  396. func() time.Time { return now },
  397. nil,
  398. )
  399. require.False(t, errorObserved.terminal)
  400. require.NotNil(t, state.pendingBareError)
  401. require.Equal(t, Usage{}, state.usage)
  402. failedObserved := observeUpstreamMessage(
  403. state,
  404. []byte(`{"type":"response.failed","response":{"id":"resp_1","usage":{"input_tokens":7,"output_tokens":2}}}`),
  405. now,
  406. func() time.Time { return now },
  407. nil,
  408. )
  409. require.True(t, failedObserved.terminal)
  410. require.Equal(t, Usage{InputTokens: 7, OutputTokens: 2}, failedObserved.usage)
  411. require.Equal(t, failedObserved.usage, state.usage)
  412. require.Nil(t, state.pendingBareError)
  413. }
  414. func TestObserveUpstreamMessageBareErrorBeforeNextCompletedKeepsBothTurns(t *testing.T) {
  415. t.Parallel()
  416. now := time.Unix(400, 0)
  417. state := &relayState{}
  418. parseUsageAndAccumulate(state, []byte(`{"type":"response.in_progress","usage":{"input_tokens":5,"output_tokens":1}}`), "response.in_progress", nil)
  419. bare := observeUpstreamMessage(
  420. state,
  421. []byte(`{"type":"error","error":{"message":"first turn failed"}}`),
  422. now,
  423. func() time.Time { return now },
  424. nil,
  425. )
  426. require.False(t, bare.terminal)
  427. bare = finalizePendingBareError(state, now)
  428. require.True(t, bare.terminal)
  429. require.Equal(t, Usage{InputTokens: 5, OutputTokens: 1}, bare.usage)
  430. completed := observeUpstreamMessage(
  431. state,
  432. []byte(`{"type":"response.completed","response":{"id":"resp_next","usage":{"input_tokens":3,"output_tokens":2}}}`),
  433. now,
  434. func() time.Time { return now },
  435. nil,
  436. )
  437. require.True(t, completed.terminal)
  438. require.Equal(t, Usage{InputTokens: 3, OutputTokens: 2}, completed.usage)
  439. require.Equal(t, Usage{InputTokens: 8, OutputTokens: 3}, state.usage)
  440. }
  441. func TestOpenAICacheCreationTokensFromUsageNestedZeroWins(t *testing.T) {
  442. t.Parallel()
  443. usage := gjson.Parse(`{"input_tokens_details":{"cache_write_tokens":0},"cache_creation_input_tokens":19}`)
  444. require.Zero(t, openAICacheCreationTokensFromUsage(usage))
  445. }
  446. func TestEmitTurnCompleteCoverage(t *testing.T) {
  447. t.Parallel()
  448. // 非 terminal 事件不应触发。
  449. called := 0
  450. emitTurnComplete(func(turn RelayTurnResult) {
  451. called++
  452. }, &relayState{requestModel: "gpt-5"}, observedUpstreamEvent{
  453. terminal: false,
  454. eventType: "response.output_text.delta",
  455. responseID: "resp_ignored",
  456. usage: Usage{InputTokens: 1},
  457. })
  458. require.Equal(t, 0, called)
  459. // 非 error 终态缺少 response_id 时不应触发。
  460. emitTurnComplete(func(turn RelayTurnResult) {
  461. called++
  462. }, &relayState{requestModel: "gpt-5"}, observedUpstreamEvent{
  463. terminal: true,
  464. eventType: "response.completed",
  465. })
  466. require.Equal(t, 0, called)
  467. // Bare error legitimately has no response_id; it must still settle the turn
  468. // or a later completed turn would suppress the adapter's aggregate fallback.
  469. var bareError RelayTurnResult
  470. emitTurnComplete(func(turn RelayTurnResult) {
  471. called++
  472. bareError = turn
  473. }, &relayState{requestModel: "gpt-5"}, observedUpstreamEvent{
  474. terminal: true,
  475. eventType: "error",
  476. usage: Usage{InputTokens: 4, OutputTokens: 1},
  477. })
  478. require.Equal(t, 1, called)
  479. require.Empty(t, bareError.RequestID)
  480. require.Equal(t, "error", bareError.TerminalEventType)
  481. require.Equal(t, Usage{InputTokens: 4, OutputTokens: 1}, bareError.Usage)
  482. // terminal 且 response_id 存在,应该触发;state=nil 时 model 为空串。
  483. var got RelayTurnResult
  484. emitTurnComplete(func(turn RelayTurnResult) {
  485. called++
  486. got = turn
  487. }, nil, observedUpstreamEvent{
  488. terminal: true,
  489. eventType: "response.completed",
  490. responseID: "resp_emit",
  491. usage: Usage{InputTokens: 2, OutputTokens: 3},
  492. })
  493. require.Equal(t, 2, called)
  494. require.Equal(t, "resp_emit", got.RequestID)
  495. require.Equal(t, "response.completed", got.TerminalEventType)
  496. require.Equal(t, 2, got.Usage.InputTokens)
  497. require.Equal(t, 3, got.Usage.OutputTokens)
  498. require.Equal(t, "", got.RequestModel)
  499. }
  500. func TestIsDisconnectErrorCoverage_CloseStatusesAndMessageBranches(t *testing.T) {
  501. t.Parallel()
  502. require.True(t, isDisconnectError(coderws.CloseError{Code: coderws.StatusNormalClosure}))
  503. require.True(t, isDisconnectError(coderws.CloseError{Code: coderws.StatusNoStatusRcvd}))
  504. require.True(t, isDisconnectError(coderws.CloseError{Code: coderws.StatusAbnormalClosure}))
  505. require.True(t, isDisconnectError(errors.New("connection reset by peer")))
  506. require.False(t, isDisconnectError(errors.New(" ")))
  507. }
  508. func TestIsTokenEventCoverageBranches(t *testing.T) {
  509. t.Parallel()
  510. require.False(t, isTokenEvent("response.in_progress"))
  511. require.False(t, isTokenEvent("response.output_item.added"))
  512. require.True(t, isTokenEvent("response.output_audio.delta"))
  513. require.True(t, isTokenEvent("response.function_call_arguments.delta"))
  514. require.True(t, isTokenEvent("response.reasoning_summary_text.delta"))
  515. require.True(t, isTokenEvent("response.output_text.done"))
  516. require.True(t, isTokenEvent("response.function_call_arguments.done"))
  517. require.False(t, isTokenEvent("response.output"))
  518. require.False(t, isTokenEvent("response.output_audio.done"))
  519. require.False(t, isTokenEvent("response.content_part.done"))
  520. require.False(t, isTokenEvent("response.output_item.done"))
  521. require.False(t, isTokenEvent("response.output_text.annotation.added"))
  522. require.False(t, isTokenEvent("response.done"))
  523. }
  524. func TestTerminalAndTokenEventSetsAreDisjoint(t *testing.T) {
  525. t.Parallel()
  526. for _, eventType := range []string{
  527. "response.completed",
  528. "response.done",
  529. "response.failed",
  530. "response.incomplete",
  531. "response.cancelled",
  532. "response.canceled",
  533. } {
  534. require.True(t, isTerminalEvent(eventType), eventType)
  535. require.False(t, isTokenEvent(eventType), eventType)
  536. }
  537. }
  538. func TestShouldParseUsageTerminalEvents(t *testing.T) {
  539. t.Parallel()
  540. for _, eventType := range []string{
  541. "response.completed",
  542. "response.done",
  543. "response.failed",
  544. "response.incomplete",
  545. "response.cancelled",
  546. "response.canceled",
  547. } {
  548. require.True(t, shouldParseUsage(eventType), eventType)
  549. }
  550. require.False(t, shouldParseUsage("response.output_text.delta"))
  551. require.False(t, shouldParseUsage(""))
  552. }
  553. func TestRelayTurnTimingHelpersCoverage(t *testing.T) {
  554. t.Parallel()
  555. now := time.Unix(100, 0)
  556. // nil state
  557. require.Nil(t, openAIWSRelayGetOrInitTurnTiming(nil, "resp_nil", now))
  558. _, ok := openAIWSRelayDeleteTurnTiming(nil, "resp_nil")
  559. require.False(t, ok)
  560. state := &relayState{}
  561. timing := openAIWSRelayGetOrInitTurnTiming(state, "resp_a", now)
  562. require.NotNil(t, timing)
  563. require.Equal(t, now, timing.startAt)
  564. // 再次获取返回同一条 timing
  565. timing2 := openAIWSRelayGetOrInitTurnTiming(state, "resp_a", now.Add(5*time.Second))
  566. require.NotNil(t, timing2)
  567. require.Equal(t, now, timing2.startAt)
  568. // 删除存在键
  569. deleted, ok := openAIWSRelayDeleteTurnTiming(state, "resp_a")
  570. require.True(t, ok)
  571. require.Equal(t, now, deleted.startAt)
  572. // 删除不存在键
  573. _, ok = openAIWSRelayDeleteTurnTiming(state, "resp_a")
  574. require.False(t, ok)
  575. }
  576. func TestObserveUpstreamMessage_ResponseModelIsTurnLocalAndTerminalWins(t *testing.T) {
  577. t.Parallel()
  578. state := &relayState{requestModel: "gpt-5.6-sol"}
  579. startAt := time.Unix(0, 0)
  580. now := startAt
  581. nowFn := func() time.Time {
  582. now = now.Add(5 * time.Millisecond)
  583. return now
  584. }
  585. created := observeUpstreamMessage(
  586. state,
  587. []byte(`{"type":"response.created","response":{"id":"resp_1","model":"gpt-5.5"}}`),
  588. startAt,
  589. nowFn,
  590. nil,
  591. )
  592. require.False(t, created.terminal)
  593. completed := observeUpstreamMessage(
  594. state,
  595. []byte(`{"type":"response.completed","response":{"id":"resp_1","model":"gpt-5.4","usage":{"input_tokens":1,"output_tokens":2}}}`),
  596. startAt,
  597. nowFn,
  598. nil,
  599. )
  600. require.True(t, completed.terminal)
  601. require.Equal(t, "gpt-5.4", completed.responseModel)
  602. require.True(t, completed.responseConflict)
  603. var firstTurn RelayTurnResult
  604. emitTurnComplete(func(turn RelayTurnResult) { firstTurn = turn }, state, completed)
  605. require.Equal(t, "gpt-5.4", firstTurn.ResponseModel)
  606. require.True(t, firstTurn.ResponseModelConflict)
  607. observeUpstreamMessage(
  608. state,
  609. []byte(`{"type":"response.created","response":{"id":"resp_2","model":"gpt-5.3"}}`),
  610. startAt,
  611. nowFn,
  612. nil,
  613. )
  614. second := observeUpstreamMessage(
  615. state,
  616. []byte(`{"type":"response.completed","response":{"id":"resp_2","model":"GPT-5.3","usage":{"input_tokens":3,"output_tokens":4}}}`),
  617. startAt,
  618. nowFn,
  619. nil,
  620. )
  621. require.Equal(t, "GPT-5.3", second.responseModel)
  622. require.False(t, second.responseConflict, "the previous turn must not contaminate this turn")
  623. }
  624. func TestObserveUpstreamMessage_ResponseIDFallbackPolicy(t *testing.T) {
  625. t.Parallel()
  626. state := &relayState{requestModel: "gpt-5"}
  627. startAt := time.Unix(0, 0)
  628. now := startAt
  629. nowFn := func() time.Time {
  630. now = now.Add(5 * time.Millisecond)
  631. return now
  632. }
  633. // 非 terminal:仅有顶层 id,不应把 event id 当成 response_id。
  634. observed := observeUpstreamMessage(
  635. state,
  636. []byte(`{"type":"response.output_text.delta","id":"evt_123","delta":"hi"}`),
  637. startAt,
  638. nowFn,
  639. nil,
  640. )
  641. require.False(t, observed.terminal)
  642. require.Equal(t, "", observed.responseID)
  643. // terminal:允许兜底用顶层 id(用于兼容少数字段变体)。
  644. observed = observeUpstreamMessage(
  645. state,
  646. []byte(`{"type":"response.completed","id":"resp_fallback","response":{"usage":{"input_tokens":1,"output_tokens":1}}}`),
  647. startAt,
  648. nowFn,
  649. nil,
  650. )
  651. require.True(t, observed.terminal)
  652. require.Equal(t, "resp_fallback", observed.responseID)
  653. }
  654. func TestObserveUpstreamMessage_ResponseServiceTierOnlyFromTerminalEvents(t *testing.T) {
  655. t.Parallel()
  656. state := &relayState{requestModel: "gpt-5.6-sol"}
  657. startAt := time.Unix(0, 0)
  658. now := startAt
  659. nowFn := func() time.Time {
  660. now = now.Add(5 * time.Millisecond)
  661. return now
  662. }
  663. created := observeUpstreamMessage(
  664. state,
  665. []byte(`{"type":"response.created","response":{"id":"resp_1","model":"gpt-5.6-sol","service_tier":"priority"}}`),
  666. startAt,
  667. nowFn,
  668. nil,
  669. )
  670. require.False(t, created.terminal)
  671. require.Equal(t, "", created.responseServiceTier)
  672. completed := observeUpstreamMessage(
  673. state,
  674. []byte(`{"type":"response.completed","response":{"id":"resp_1","model":"gpt-5.6-sol","service_tier":"default","usage":{"input_tokens":1,"output_tokens":2}}}`),
  675. startAt,
  676. nowFn,
  677. nil,
  678. )
  679. require.True(t, completed.terminal)
  680. require.Equal(t, "default", completed.responseServiceTier, "the echoed priority tier must lose to the terminal declaration")
  681. var turn RelayTurnResult
  682. emitTurnComplete(func(result RelayTurnResult) { turn = result }, state, completed)
  683. require.Equal(t, "default", turn.ResponseServiceTier)
  684. var result RelayResult
  685. enrichResult(&result, state, now.Sub(startAt))
  686. require.Equal(t, "default", result.ResponseServiceTier)
  687. // A later turn without any declaration must not inherit the previous one.
  688. observeUpstreamMessage(state, []byte(`{"type":"response.created","response":{"id":"resp_2","model":"gpt-5.6-sol"}}`), startAt, nowFn, nil)
  689. second := observeUpstreamMessage(
  690. state,
  691. []byte(`{"type":"response.completed","response":{"id":"resp_2","model":"gpt-5.6-sol","usage":{"input_tokens":3,"output_tokens":4}}}`),
  692. startAt,
  693. nowFn,
  694. nil,
  695. )
  696. require.Equal(t, "", second.responseServiceTier)
  697. }