Skip to content

Commit 6019370

Browse files
1220moritzclaude
andauthored
feat(api): abort signal support for Zoo Gateway completePrompt (#2430)
* feat(api): abort signal support for Zoo Gateway completePrompt Follow the shared abort-signal contract (as in requesty.ts): bridge the caller's signal to a per-request signal, merge the optional timeout, race model lookup and the SDK call against the signal, normalize aborts to AbortError without surfacing a gateway error, drop a late result, and detach the bridge listener on every exit path. Tests cover a pre-aborted signal, a caller abort and a timeout-only abort during model lookup, an abort mid-flight while the SDK ignores the signal, an SDK abort error after the caller aborts, a completion that resolves just before the abort, and removal of the same abort listener. Split out of #2306. Co-Authored-By: Claude Opus 5.5 <[email protected]> * fix(api): check Zoo Gateway cancellation before auth and during error notifications - A pre-aborted signal rejects with an AbortError before the session check, so a cancel wins over a missing session. - The gateway error notification can wait for the user. Race it against the request signal, so an abort or timeout still settles at once with an AbortError. - Tests: no model lookup on a pre-aborted signal (with and without a session), an ordinary lookup failure is kept while the signal is live, a timeout during the SDK call, an abort while a notification waits, and removal of the same listener after success, failure, and abort. Co-Authored-By: Claude Opus 5.5 <[email protected]> * test(api): cover Zoo Gateway abort and notification failure paths - A raw AbortError or the SDK abort error from the model lookup becomes the provider AbortError while the caller signal is live. - A 402 that rejects just before the caller aborts shows no notification. - A failed error notification is logged (with and without a signal) and the gateway error is still thrown. Nothing is logged after an abort. - An SDK failure that is not an Error is rethrown unchanged. Co-Authored-By: Claude Opus 5.5 <[email protected]> * test(api): cover the Zoo Gateway session check and empty choices in completePrompt Co-Authored-By: Claude Opus 5.5 <[email protected]> --------- Co-authored-by: Claude Opus 5.5 <[email protected]>
1 parent d940174 commit 6019370

2 files changed

Lines changed: 368 additions & 20 deletions

File tree

‎src/api/providers/__tests__/zoo-gateway.spec.ts‎

Lines changed: 286 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ vitest.mock("../../../i18n", () => ({
2626
t: (key: string) => key,
2727
}))
2828

29+
import type { MockInstance } from "vitest"
2930
import OpenAI from "openai"
3031

3132
import { zooGatewayDefaultModelId, ZOO_GATEWAY_DEFAULT_TEMPERATURE } from "@roo-code/types"
@@ -500,10 +501,17 @@ describe("ZooGatewayHandler", () => {
500501
})
501502

502503
describe("completePrompt", () => {
504+
let consoleError: MockInstance<typeof console.error>
505+
503506
beforeEach(() => {
504507
mockCreate.mockImplementation(async () => ({
505508
choices: [{ message: { role: "assistant", content: "Test completion response" } }],
506509
}))
510+
consoleError = vitest.spyOn(console, "error").mockImplementation(() => {})
511+
})
512+
513+
afterEach(() => {
514+
consoleError.mockRestore()
507515
})
508516

509517
it("returns completion text from the gateway", async () => {
@@ -520,6 +528,7 @@ describe("ZooGatewayHandler", () => {
520528
temperature: ZOO_GATEWAY_DEFAULT_TEMPERATURE,
521529
max_completion_tokens: 64000,
522530
}),
531+
{},
523532
)
524533
})
525534

@@ -534,14 +543,287 @@ describe("ZooGatewayHandler", () => {
534543
)
535544
})
536545

537-
it("returns an empty string when the model returns no content", async () => {
546+
it("forwards a per-request signal and the timeout to the SDK", async () => {
538547
const handler = new ZooGatewayHandler(mockOptions)
539-
mockCreate.mockImplementation(async () => ({
540-
choices: [{ message: { role: "assistant", content: null } }],
541-
}))
548+
const caller = new AbortController()
549+
550+
await handler.completePrompt("Test", { abortSignal: caller.signal, timeoutMs: 5_000 })
551+
552+
const requestOptions = mockCreate.mock.calls[0][1] as { signal?: AbortSignal; timeout?: number }
553+
expect(requestOptions.signal).toBeInstanceOf(AbortSignal)
554+
expect(requestOptions.signal).not.toBe(caller.signal)
555+
expect(requestOptions.timeout).toBe(5_000)
556+
})
557+
558+
it("sends no signal when the caller gives none", async () => {
559+
const handler = new ZooGatewayHandler(mockOptions)
560+
561+
await handler.completePrompt("Test")
562+
563+
expect(mockCreate.mock.calls[0][1]).toEqual({})
564+
})
565+
566+
it.each([
567+
["with a session", mockOptions],
568+
["without a session", { zooGatewayModelId: "anthropic/claude-sonnet-4" }],
569+
])(
570+
"rejects a pre-aborted signal with an AbortError, %s, before model lookup",
571+
async (_label, options: ApiHandlerOptions) => {
572+
const { getModels } = await import("../fetchers/modelCache")
573+
const handler = new ZooGatewayHandler(options)
574+
const caller = new AbortController()
575+
caller.abort()
576+
577+
await expect(handler.completePrompt("Test", { abortSignal: caller.signal })).rejects.toMatchObject({
578+
name: "AbortError",
579+
})
580+
expect(getModels).not.toHaveBeenCalled()
581+
expect(mockCreate).not.toHaveBeenCalled()
582+
},
583+
)
584+
585+
it("keeps an ordinary model lookup failure while the caller signal is live", async () => {
586+
const { getModels, refreshModels } = await import("../fetchers/modelCache")
587+
const lookupFailure = new Error("model catalog unavailable")
588+
vitest.mocked(getModels).mockRejectedValue(lookupFailure)
589+
vitest.mocked(refreshModels).mockRejectedValue(lookupFailure)
590+
const handler = new ZooGatewayHandler(mockOptions)
591+
592+
await expect(handler.completePrompt("Test", { abortSignal: new AbortController().signal })).rejects.toBe(
593+
lookupFailure,
594+
)
595+
})
596+
597+
it.each([
598+
["a raw AbortError", Object.assign(new Error("The user aborted a request"), { name: "AbortError" })],
599+
["the SDK abort error", new Error("Request was aborted.")],
600+
])(
601+
"normalizes %s from the model lookup to the provider AbortError while the caller signal is live",
602+
async (_label, lookupFailure: Error) => {
603+
const { getModels, refreshModels } = await import("../fetchers/modelCache")
604+
vitest.mocked(getModels).mockRejectedValue(lookupFailure)
605+
vitest.mocked(refreshModels).mockRejectedValue(lookupFailure)
606+
const handler = new ZooGatewayHandler(mockOptions)
607+
608+
await expect(
609+
handler.completePrompt("Test", { abortSignal: new AbortController().signal }),
610+
).rejects.toMatchObject({
611+
name: "AbortError",
612+
message: "The Zoo Gateway request was aborted",
613+
})
614+
expect(mockCreate).not.toHaveBeenCalled()
615+
},
616+
)
617+
618+
it("rejects with an AbortError when the timeout alone expires during the SDK call", async () => {
619+
const handler = new ZooGatewayHandler(mockOptions)
620+
let sdkSignal: AbortSignal | undefined
621+
mockCreate.mockImplementation((_params: unknown, requestOptions: { signal?: AbortSignal }) => {
622+
sdkSignal = requestOptions.signal
623+
return new Promise(() => {})
624+
})
625+
626+
await expect(handler.completePrompt("Test", { timeoutMs: 20 })).rejects.toMatchObject({
627+
name: "AbortError",
628+
message: "The Zoo Gateway request was aborted",
629+
})
630+
expect(sdkSignal?.aborted).toBe(true)
631+
expect(showErrorMessage).not.toHaveBeenCalled()
632+
})
633+
634+
it("settles on abort while a gateway error notification waits for the user", async () => {
635+
const handler = new ZooGatewayHandler(mockOptions)
636+
const caller = new AbortController()
637+
mockCreate.mockRejectedValue(makeApiError(402))
638+
showErrorMessage.mockImplementation(() => new Promise(() => {}))
639+
640+
const pending = handler.completePrompt("Test", { abortSignal: caller.signal })
641+
await vitest.waitFor(() => expect(showErrorMessage).toHaveBeenCalled())
642+
caller.abort()
643+
644+
await expect(pending).rejects.toMatchObject({
645+
name: "AbortError",
646+
message: "The Zoo Gateway request was aborted",
647+
})
648+
// The abort is not a failure of the notification, so nothing is logged.
649+
expect(consoleError).not.toHaveBeenCalled()
650+
})
651+
652+
it.each([
653+
["with a live caller signal", true, new Error("notification failed"), "notification failed"],
654+
["without a caller signal", false, "notification dismissed", "notification dismissed"],
655+
])(
656+
"logs a failed gateway error notification and still rejects with the gateway error, %s",
657+
async (_label, withSignal: boolean, notificationFailure: unknown, logged: string) => {
658+
const handler = new ZooGatewayHandler(mockOptions)
659+
mockCreate.mockRejectedValue(makeApiError(402, { message: "out of credits" }))
660+
showErrorMessage.mockRejectedValue(notificationFailure)
661+
const options = withSignal ? { abortSignal: new AbortController().signal } : undefined
662+
663+
await expect(handler.completePrompt("Test", options)).rejects.toThrow(
664+
"Zoo Gateway completion error: out of credits",
665+
)
666+
expect(consoleError).toHaveBeenCalledWith("Failed to surface Zoo Gateway error:", logged)
667+
},
668+
)
669+
670+
it("rethrows an SDK failure that is not an Error unchanged", async () => {
671+
const handler = new ZooGatewayHandler(mockOptions)
672+
const failure = { reason: "upstream failure" }
673+
mockCreate.mockRejectedValue(failure)
674+
675+
await expect(handler.completePrompt("Test", { abortSignal: new AbortController().signal })).rejects.toBe(
676+
failure,
677+
)
678+
})
679+
680+
it("rejects with an AbortError when the timeout alone expires during model lookup", async () => {
681+
const { getModels } = await import("../fetchers/modelCache")
682+
vitest.mocked(getModels).mockImplementation(() => new Promise(() => {}))
683+
const handler = new ZooGatewayHandler(mockOptions)
684+
685+
await expect(handler.completePrompt("Test", { timeoutMs: 20 })).rejects.toMatchObject({
686+
name: "AbortError",
687+
message: "The Zoo Gateway request was aborted",
688+
})
689+
expect(mockCreate).not.toHaveBeenCalled()
690+
})
691+
692+
it("rejects with an AbortError when the signal aborts during model lookup", async () => {
693+
const { getModels } = await import("../fetchers/modelCache")
694+
vitest.mocked(getModels).mockImplementation(() => new Promise(() => {}))
695+
const handler = new ZooGatewayHandler(mockOptions)
696+
const caller = new AbortController()
697+
698+
const pending = handler.completePrompt("Test", { abortSignal: caller.signal })
699+
caller.abort()
700+
701+
await expect(pending).rejects.toMatchObject({
702+
name: "AbortError",
703+
message: "The Zoo Gateway request was aborted",
704+
})
705+
expect(mockCreate).not.toHaveBeenCalled()
706+
})
707+
708+
it("rejects with an AbortError at once when aborted mid-flight, without a gateway error", async () => {
709+
const handler = new ZooGatewayHandler(mockOptions)
710+
const caller = new AbortController()
711+
let sdkSignal: AbortSignal | undefined
712+
// The SDK call never settles, as when the SDK sleeps between retries.
713+
mockCreate.mockImplementation((_params: unknown, requestOptions: { signal?: AbortSignal }) => {
714+
sdkSignal = requestOptions.signal
715+
return new Promise(() => {})
716+
})
717+
718+
const pending = handler.completePrompt("Test", { abortSignal: caller.signal })
719+
await vitest.waitFor(() => expect(mockCreate).toHaveBeenCalled())
720+
caller.abort()
721+
722+
await expect(pending).rejects.toMatchObject({
723+
name: "AbortError",
724+
message: "The Zoo Gateway request was aborted",
725+
})
726+
expect(sdkSignal?.aborted).toBe(true)
727+
expect(showErrorMessage).not.toHaveBeenCalled()
728+
})
729+
730+
it("reports an SDK abort error after the caller aborts as an AbortError, not a gateway error", async () => {
731+
const handler = new ZooGatewayHandler(mockOptions)
732+
const caller = new AbortController()
733+
mockCreate.mockImplementation(async () => {
734+
caller.abort()
735+
throw new Error("Request was aborted.")
736+
})
737+
738+
await expect(handler.completePrompt("Test", { abortSignal: caller.signal })).rejects.toMatchObject({
739+
name: "AbortError",
740+
})
741+
expect(showErrorMessage).not.toHaveBeenCalled()
742+
})
743+
744+
it("shows no gateway error for a 402 that rejects just before the caller aborts", async () => {
745+
const handler = new ZooGatewayHandler(mockOptions)
746+
const caller = new AbortController()
747+
let rejectCompletion: (reason: unknown) => void = () => {}
748+
const completion = new Promise((_resolve, reject) => {
749+
rejectCompletion = reject
750+
})
751+
mockCreate.mockImplementation(() => completion)
752+
753+
const pending = handler.completePrompt("Test", { abortSignal: caller.signal })
754+
await vitest.waitFor(() => expect(mockCreate).toHaveBeenCalled())
755+
// This callback runs after the handler's own race settles with the 402 but before completePrompt
756+
// resumes, so only the abort check before the notification keeps the 402 notification away.
757+
completion.catch(() => caller.abort())
758+
rejectCompletion(makeApiError(402))
759+
760+
await expect(pending).rejects.toMatchObject({ name: "AbortError" })
761+
expect(showErrorMessage).not.toHaveBeenCalled()
762+
})
763+
764+
it("does not return a completion that resolves just before the caller aborts", async () => {
765+
const handler = new ZooGatewayHandler(mockOptions)
766+
const caller = new AbortController()
767+
let resolveCompletion: (value: unknown) => void = () => {}
768+
const completion = new Promise((resolve) => {
769+
resolveCompletion = resolve
770+
})
771+
mockCreate.mockImplementation(() => completion)
772+
773+
const pending = handler.completePrompt("Test", { abortSignal: caller.signal })
774+
await vitest.waitFor(() => expect(mockCreate).toHaveBeenCalled())
775+
// This callback runs after the handler's own race settles but before completePrompt resumes,
776+
// so only the check after the await can drop the late result.
777+
void completion.then(() => caller.abort())
778+
resolveCompletion({ choices: [{ message: { role: "assistant", content: "late" } }] })
779+
780+
await expect(pending).rejects.toMatchObject({ name: "AbortError" })
781+
})
782+
783+
it.each([
784+
["a completed request", () => {}],
785+
["a failed request", () => mockCreate.mockRejectedValue(new Error("upstream failure"))],
786+
["a request aborted mid-flight", () => mockCreate.mockImplementation(() => new Promise(() => {}))],
787+
])("removes the same abort listener from the caller's signal after %s", async (_label, setup) => {
788+
setup()
789+
const handler = new ZooGatewayHandler(mockOptions)
790+
const caller = new AbortController()
791+
const addListener = vitest.spyOn(caller.signal, "addEventListener")
792+
const removeListener = vitest.spyOn(caller.signal, "removeEventListener")
793+
794+
const pending = handler.completePrompt("Test", { abortSignal: caller.signal }).catch(() => undefined)
795+
await vitest.waitFor(() => expect(mockCreate).toHaveBeenCalled())
796+
caller.abort()
797+
await pending
798+
799+
const abortListeners = addListener.mock.calls.filter(([type]) => type === "abort")
800+
expect(abortListeners).toHaveLength(1)
801+
expect(removeListener).toHaveBeenCalledWith("abort", abortListeners[0][1])
802+
})
803+
804+
it.each([
805+
["no content", [{ message: { role: "assistant", content: null } }]],
806+
["no choices", []],
807+
])("returns an empty string when the model returns %s", async (_label, choices: unknown[]) => {
808+
const handler = new ZooGatewayHandler(mockOptions)
809+
mockCreate.mockImplementation(async () => ({ choices }))
542810

543811
await expect(handler.completePrompt("Test")).resolves.toBe("")
544812
})
813+
814+
it.each([
815+
["with a live caller signal", { abortSignal: new AbortController().signal }],
816+
["without a caller signal", undefined],
817+
])("requires a session before the model lookup, %s", async (_label, options) => {
818+
const { getModels } = await import("../fetchers/modelCache")
819+
const handler = new ZooGatewayHandler({ zooGatewayModelId: "anthropic/claude-sonnet-4" })
820+
821+
await expect(handler.completePrompt("Test", options)).rejects.toThrow(
822+
"Zoo Gateway requires authentication. Please sign in to Zoo Code first.",
823+
)
824+
expect(getModels).not.toHaveBeenCalled()
825+
expect(mockCreate).not.toHaveBeenCalled()
826+
})
545827
})
546828

547829
describe("classifyGatewayApiError", () => {

0 commit comments

Comments
 (0)