Skip to content

Commit 46fb004

Browse files
WebMadWebMad
committed
fix(codex): preserve websocket responses across consumer pauses
1 parent b60bb0b commit 46fb004

2 files changed

Lines changed: 70 additions & 9 deletions

File tree

‎src/api/providers/CodexWebSocketTransport.ts‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -70,18 +70,24 @@ export class CodexWebSocketTransport {
7070
request: CodexWebSocketRequestManager,
7171
response: CodexWebSocketResponseManager,
7272
): AsyncGenerator<CodexResponseEvent> {
73+
const socket = request.socket
7374
request.refreshTimeout()
74-
this.send(request.socket, prepared, headers, prepared.fullContextReason === undefined)
75+
this.send(socket, prepared, headers, prepared.fullContextReason === undefined)
7576
for await (const [data] of request.events) {
76-
request.refreshTimeout()
77+
request.signal.throwIfAborted()
78+
// Consumer work between yields is not socket inactivity.
79+
request.clearTimeout()
7780
const event = parseResponseEvent(String(data))
7881
if (response.accept(event) === "retry") {
79-
this.send(request.socket, prepared, headers, false, "server cache miss")
82+
request.refreshTimeout()
83+
this.send(socket, prepared, headers, false, "server cache miss")
8084
continue
8185
}
82-
if (response.completed) request.clearTimeout()
8386
yield event
8487
if (response.completed) return
88+
request.signal.throwIfAborted()
89+
// Drain queued events even if the socket closed while the consumer was paused.
90+
if (socket.readyState === socket.OPEN) request.refreshTimeout()
8591
}
8692
throw new Error("Codex WebSocket closed before response completed; retry the request")
8793
}

‎src/api/providers/__tests__/CodexWebSocketTransport.spec.ts‎

Lines changed: 60 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -359,14 +359,30 @@ describe("CodexWebSocketTransport", () => {
359359
}
360360
const stream = transport.stream(body(), { ...options(), signal: controller.signal })
361361
expect((await stream.next()).value).toMatchObject({ delta: "Partial" })
362-
controller.abort()
363-
await expect(stream.next()).rejects.toThrow()
362+
const reason = new Error("Stopped")
363+
controller.abort(reason)
364+
await expect(stream.next()).rejects.toBe(reason)
364365
reply = (_request, socket) => complete(socket)
365366
await collectStream(transport.stream(body(), options()))
366367
expect(connections).toBe(2)
367368
expect(requests[1].previous_response_id).toBeUndefined()
368369
})
369370

371+
it("preserves Stop when cancellation races with receiving completion", async () => {
372+
const controller = new AbortController()
373+
const reason = new Error("Stopped")
374+
reply = (_request, socket) => {
375+
const scope = transport["requestScope"]
376+
if (!scope) throw new Error("Missing active request scope")
377+
// The iterator already received this frame, but the consumer has not resumed yet.
378+
scope.requestManager.socket.once("message", () => controller.abort(reason))
379+
complete(socket)
380+
}
381+
await expect(collectStream(transport.stream(body(), { ...options(), signal: controller.signal }))).rejects.toBe(
382+
reason,
383+
)
384+
})
385+
370386
it("disposes an active stream and reconnects with full context on the next request", async () => {
371387
reply = (_request, socket) =>
372388
socket.send(JSON.stringify({ type: "response.output_text.delta", delta: "Partial" }))
@@ -392,17 +408,20 @@ describe("CodexWebSocketTransport", () => {
392408
expect(requests[1].previous_response_id).toBeUndefined()
393409
})
394410

395-
it("disposes request deadlines while the consumer is paused on a partial response", async () => {
411+
it("disposes request deadlines while waiting for more response events", async () => {
396412
vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout"] })
397413
reply = (_request, socket) =>
398414
socket.send(JSON.stringify({ type: "response.output_text.delta", delta: "Partial" }))
399415
const stream = transport.stream(body(), options())
400416
try {
401417
await stream.next()
418+
expect(vi.getTimerCount()).toBe(0)
419+
const next = expect(stream.next()).rejects.toThrow("closed before response completed")
420+
await vi.advanceTimersByTimeAsync(0)
402421
expect(vi.getTimerCount()).toBe(1)
403422
await transport.dispose()
404423
expect(vi.getTimerCount()).toBe(0)
405-
await expect(stream.next()).rejects.toThrow("closed before response completed")
424+
await next
406425
} finally {
407426
await stream.return(undefined)
408427
vi.useRealTimers()
@@ -411,7 +430,9 @@ describe("CodexWebSocketTransport", () => {
411430

412431
it("bounds a silent response with a timeout", async () => {
413432
reply = () => {}
414-
await expect(collectStream(transport.stream(body(), { ...options(), timeoutMs: 100 }))).rejects.toThrow()
433+
await expect(collectStream(transport.stream(body(), { ...options(), timeoutMs: 100 }))).rejects.toThrow(
434+
"Codex WebSocket stream timed out",
435+
)
415436
})
416437

417438
it.each(["null", "[]", "{}", '{"type":42}', "invalid JSON"])(
@@ -459,6 +480,40 @@ describe("CodexWebSocketTransport", () => {
459480
}
460481
})
461482

483+
it.each([false, true])("preserves queued completion during a consumer pause (server closes: %s)", async (close) => {
484+
vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout"] })
485+
reply = (_request, socket) =>
486+
socket.send(JSON.stringify({ type: "response.output_text.delta", delta: "Partial" }))
487+
const stream = transport.stream(body(), { ...options(), timeoutMs: 100 })
488+
try {
489+
expect((await stream.next()).value).toMatchObject({ delta: "Partial" })
490+
const scope = transport["requestScope"]
491+
if (!scope) throw new Error("Missing active request scope")
492+
const received = once(scope.requestManager.socket, "message")
493+
const serverSocket = [...server.clients][0]
494+
complete(serverSocket)
495+
await received
496+
if (close) {
497+
const closed = once(scope.requestManager.socket, "close")
498+
serverSocket.close()
499+
await closed
500+
}
501+
await vi.advanceTimersByTimeAsync(1_000)
502+
expect((await stream.next()).value).toMatchObject({
503+
type: "response.completed",
504+
response: { id: "resp_1", usage: { input_tokens: 1, output_tokens: 2 } },
505+
})
506+
expect((await stream.next()).done).toBe(true)
507+
reply = (_request, socket) => complete(socket)
508+
await collectStream(transport.stream(body(), options()))
509+
expect(connections).toBe(close ? 2 : 1)
510+
expect(requests[1].previous_response_id).toBe(close ? undefined : "resp_1")
511+
} finally {
512+
await stream.return(undefined)
513+
vi.useRealTimers()
514+
}
515+
})
516+
462517
it("rejects overlapping requests without disturbing the active stream", async () => {
463518
reply = (_request, socket) =>
464519
socket.send(JSON.stringify({ type: "response.output_text.delta", delta: "Partial" }))

0 commit comments

Comments
 (0)