Skip to content

Commit b428b81

Browse files
easonliang28easonliang28
committed
refactor(api): leverage abort-signal utils and RequestConfigBuilder for gemini/mistral/lite-llm
Address review feedback (edelauna, #1303): - catch blocks: isRequestAborted(error, signal) catches SDK-native abort errors that surface before the signal flag propagates (error-name and exact-message branches), and createAbortError normalizes the thrown error to the series abort contract - createMessage bridge: RequestConfigBuilder.addMergedSignal (AbortSignal.any) replaces the manual AbortController + addEventListener/removeEventListener plumbing; the finally cleanup blocks are gone - pre-abort fast-fail uses the throwIfAborted helper - specs: message assertions updated to the helper messages, listener-identity assertions inverted to assert no manual listener management, and new regression tests pin the SDK-native-abort branch per provider
1 parent 4e9a602 commit b428b81

6 files changed

Lines changed: 151 additions & 151 deletions

File tree

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

Lines changed: 31 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -478,9 +478,25 @@ describe("GeminiHandler", () => {
478478
const error = await handler
479479
.completePrompt("Test prompt", { abortSignal: controller.signal })
480480
.catch((e: unknown) => e)
481-
expect(error).toBeInstanceOf(DOMException)
481+
expect(error).toBeInstanceOf(Error)
482+
expect((error as Error).name).toBe("AbortError")
483+
expect((error as Error).message).toBe("The Gemini request was aborted")
484+
})
485+
486+
it("should surface a standard AbortError when the SDK throws an abort error before the signal flag propagates", async () => {
487+
// The signal flag has not propagated yet, but the SDK rejects with its
488+
// own abort error: isRequestAborted must catch the error-name branch.
489+
const controller = new AbortController()
490+
vi.mocked(handler["client"].models.generateContent).mockRejectedValue(
491+
Object.assign(new Error("Request was aborted."), { name: "AbortError" }),
492+
)
493+
494+
const error = await handler
495+
.completePrompt("Test prompt", { abortSignal: controller.signal })
496+
.catch((e: unknown) => e)
497+
expect(error).toBeInstanceOf(Error)
482498
expect((error as Error).name).toBe("AbortError")
483-
expect((error as Error).message).toBe("Gemini completion aborted")
499+
expect((error as Error).message).toBe("The Gemini request was aborted")
484500
})
485501

486502
it("should surface the wrapped provider error when the request fails without options", async () => {
@@ -1067,7 +1083,7 @@ describe("GeminiHandler", () => {
10671083
const error = await collectStream(stream).catch((e: unknown) => e)
10681084
expect(error).toBeInstanceOf(Error)
10691085
expect((error as Error).name).toBe("AbortError")
1070-
expect((error as Error).message).toBe("Gemini request aborted")
1086+
expect((error as Error).message).toBe("This operation was aborted")
10711087
expect(handler["client"].models.generateContentStream).not.toHaveBeenCalled()
10721088
})
10731089

@@ -1115,31 +1131,29 @@ describe("GeminiHandler", () => {
11151131
})
11161132
expect(error).toBeInstanceOf(Error)
11171133
expect((error as Error).name).toBe("AbortError")
1118-
expect((error as Error).message).toBe("Gemini request aborted")
1134+
expect((error as Error).message).toBe("The Gemini request was aborted")
11191135
expect(capturedSignal).toBeDefined()
1120-
// The in-flight request must run against a request-local signal, not the
1121-
// external one forwarded by reference.
1136+
// The in-flight request must run against a request-local merged signal, not
1137+
// the external one forwarded by reference.
11221138
expect(capturedSignal).not.toBe(controller.signal)
11231139
expect(capturedSignal?.aborted).toBe(true)
1124-
// The bridge registers a once-only listener on the external signal and
1125-
// detaches it when the request settles. Target the last "abort"
1126-
// registration (the bridge's listener) and assert the exact reference
1127-
// so a bridge that removes a different callback cannot pass.
1128-
const abortAddCalls = addEventListenerSpy.mock.calls.filter(([event]) => event === "abort")
1129-
const addedListener = abortAddCalls[abortAddCalls.length - 1]?.[1]
1130-
expect(typeof addedListener).toBe("function")
1131-
expect(addEventListenerSpy).toHaveBeenCalledWith("abort", addedListener, { once: true })
1132-
expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", addedListener)
1140+
// The bridge (RequestConfigBuilder.addMergedSignal) uses AbortSignal.any,
1141+
// so the external signal must never be managed with manual listeners.
1142+
expect(addEventListenerSpy).not.toHaveBeenCalled()
1143+
expect(removeEventListenerSpy).not.toHaveBeenCalled()
11331144
})
11341145

1135-
it("should not set config.abortSignal when no external signal is provided", async () => {
1146+
it("should set a live request-local config.abortSignal when no external signal is provided", async () => {
11361147
const stub = vi.fn().mockReturnValue((async function* () {})())
11371148
handler["client"].models.generateContentStream = stub
11381149

11391150
await collectStream(handler.createMessage("You are a helpful assistant", messages))
11401151

11411152
const config = stub.mock.calls[0][0].config
1142-
expect(config.abortSignal).toBeUndefined()
1153+
// The request-local controller signal is always present so the provider
1154+
// keeps its own abort handle; with no external signal it never aborts.
1155+
expect(config.abortSignal).toBeInstanceOf(AbortSignal)
1156+
expect(config.abortSignal?.aborted).toBe(false)
11431157
})
11441158

11451159
it("should wrap a non-abort stream failure with the i18n message and capture telemetry", async () => {

‎src/api/providers/__tests__/lite-llm.spec.ts‎

Lines changed: 22 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1432,9 +1432,23 @@ describe("LiteLLMHandler", () => {
14321432
const error = await handler
14331433
.completePrompt("test prompt", { abortSignal: controller.signal })
14341434
.catch((e: unknown) => e)
1435-
expect(error).toBeInstanceOf(DOMException)
1435+
expect(error).toBeInstanceOf(Error)
1436+
expect((error as Error).name).toBe("AbortError")
1437+
expect((error as Error).message).toBe("The LiteLLM request was aborted")
1438+
})
1439+
1440+
it("should surface a standard AbortError when the SDK throws an abort error before the signal flag propagates", async () => {
1441+
// The signal flag has not propagated yet, but the SDK rejects with its
1442+
// own abort error: isRequestAborted must catch the error-name branch.
1443+
mockCreate.mockRejectedValueOnce(Object.assign(new Error("Request was aborted."), { name: "AbortError" }))
1444+
const controller = new AbortController()
1445+
1446+
const error = await handler
1447+
.completePrompt("test prompt", { abortSignal: controller.signal })
1448+
.catch((e: unknown) => e)
1449+
expect(error).toBeInstanceOf(Error)
14361450
expect((error as Error).name).toBe("AbortError")
1437-
expect((error as Error).message).toBe("LiteLLM completion aborted")
1451+
expect((error as Error).message).toBe("The LiteLLM request was aborted")
14381452
})
14391453

14401454
it("should surface the wrapped provider error when the request fails without options", async () => {
@@ -1474,7 +1488,7 @@ describe("LiteLLMHandler", () => {
14741488
const error = await collectStream(stream).catch((e: unknown) => e)
14751489
expect(error).toBeInstanceOf(Error)
14761490
expect((error as Error).name).toBe("AbortError")
1477-
expect((error as Error).message).toBe("LiteLLM streaming aborted")
1491+
expect((error as Error).message).toBe("This operation was aborted")
14781492
expect(mockCreate).not.toHaveBeenCalled()
14791493
})
14801494

@@ -1532,18 +1546,13 @@ describe("LiteLLMHandler", () => {
15321546
})
15331547
expect(error).toBeInstanceOf(Error)
15341548
expect((error as Error).name).toBe("AbortError")
1535-
expect((error as Error).message).toBe("LiteLLM streaming aborted")
1549+
expect((error as Error).message).toBe("The LiteLLM request was aborted")
15361550
expect(capturedSignal).toBeDefined()
15371551
expect(capturedSignal?.aborted).toBe(true)
1538-
// The bridge registers a once-only listener on the external signal and
1539-
// detaches it when the request settles. Target the last "abort"
1540-
// registration (the bridge's listener) and assert the exact reference
1541-
// so a bridge that removes a different callback cannot pass.
1542-
const abortAddCalls = addEventListenerSpy.mock.calls.filter(([event]) => event === "abort")
1543-
const addedListener = abortAddCalls[abortAddCalls.length - 1]?.[1]
1544-
expect(typeof addedListener).toBe("function")
1545-
expect(addEventListenerSpy).toHaveBeenCalledWith("abort", addedListener, { once: true })
1546-
expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", addedListener)
1552+
// The bridge (RequestConfigBuilder.addMergedSignal) uses AbortSignal.any,
1553+
// so the external signal must never be managed with manual listeners.
1554+
expect(addEventListenerSpy).not.toHaveBeenCalled()
1555+
expect(removeEventListenerSpy).not.toHaveBeenCalled()
15471556
})
15481557

15491558
it("should wrap a non-abort stream failure with the i18n-free provider message and no metadata", async () => {

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

Lines changed: 33 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,7 @@ describe("MistralHandler", () => {
135135
tools: expect.any(Array),
136136
toolChoice: "any",
137137
}),
138+
expect.objectContaining({ fetchOptions: { signal: expect.any(AbortSignal) } }),
138139
)
139140

140141
expect(result.value).toBeDefined()
@@ -501,6 +502,7 @@ describe("MistralHandler", () => {
501502
]),
502503
toolChoice: "any",
503504
}),
505+
expect.objectContaining({ fetchOptions: { signal: expect.any(AbortSignal) } }),
504506
)
505507
})
506508

@@ -518,6 +520,7 @@ describe("MistralHandler", () => {
518520
tools: expect.any(Array),
519521
toolChoice: "any",
520522
}),
523+
expect.objectContaining({ fetchOptions: { signal: expect.any(AbortSignal) } }),
521524
)
522525
})
523526

@@ -655,6 +658,7 @@ describe("MistralHandler", () => {
655658
expect.objectContaining({
656659
toolChoice: "any",
657660
}),
661+
expect.objectContaining({ fetchOptions: { signal: expect.any(AbortSignal) } }),
658662
)
659663
})
660664
})
@@ -799,9 +803,23 @@ describe("MistralHandler", () => {
799803
const error = await handler
800804
.completePrompt("Test prompt", { abortSignal: controller.signal })
801805
.catch((e: unknown) => e)
802-
expect(error).toBeInstanceOf(DOMException)
806+
expect(error).toBeInstanceOf(Error)
807+
expect((error as Error).name).toBe("AbortError")
808+
expect((error as Error).message).toBe("The Mistral request was aborted")
809+
})
810+
811+
it("should surface a standard AbortError when the SDK throws an abort error before the signal flag propagates", async () => {
812+
// The signal flag has not propagated yet, but the SDK rejects with its
813+
// own abort error: isRequestAborted must catch the error-name branch.
814+
mockComplete.mockRejectedValueOnce(Object.assign(new Error("Request was aborted."), { name: "AbortError" }))
815+
const controller = new AbortController()
816+
817+
const error = await handler
818+
.completePrompt("Test prompt", { abortSignal: controller.signal })
819+
.catch((e: unknown) => e)
820+
expect(error).toBeInstanceOf(Error)
803821
expect((error as Error).name).toBe("AbortError")
804-
expect((error as Error).message).toBe("Mistral completion aborted")
822+
expect((error as Error).message).toBe("The Mistral request was aborted")
805823
})
806824
})
807825

@@ -827,7 +845,7 @@ describe("MistralHandler", () => {
827845
const error = await collectStream(stream).catch((e: unknown) => e)
828846
expect(error).toBeInstanceOf(Error)
829847
expect((error as Error).name).toBe("AbortError")
830-
expect((error as Error).message).toBe("Mistral completion aborted")
848+
expect((error as Error).message).toBe("This operation was aborted")
831849
expect(mockCreate).not.toHaveBeenCalled()
832850
})
833851

@@ -883,21 +901,16 @@ describe("MistralHandler", () => {
883901
})
884902
expect(error).toBeInstanceOf(Error)
885903
expect((error as Error).name).toBe("AbortError")
886-
expect((error as Error).message).toBe("Mistral completion aborted")
904+
expect((error as Error).message).toBe("The Mistral request was aborted")
887905
expect(capturedSignal).toBeDefined()
888-
// The in-flight stream must run against a request-local signal, not the
889-
// external one forwarded by reference.
906+
// The in-flight stream must run against a request-local merged signal, not
907+
// the external one forwarded by reference.
890908
expect(capturedSignal).not.toBe(controller.signal)
891909
expect(capturedSignal?.aborted).toBe(true)
892-
// The bridge registers a once-only listener on the external signal and
893-
// detaches it when the request settles. Target the last "abort"
894-
// registration (the bridge's listener) and assert the exact reference
895-
// so a bridge that removes a different callback cannot pass.
896-
const abortAddCalls = addEventListenerSpy.mock.calls.filter(([event]) => event === "abort")
897-
const addedListener = abortAddCalls[abortAddCalls.length - 1]?.[1]
898-
expect(typeof addedListener).toBe("function")
899-
expect(addEventListenerSpy).toHaveBeenCalledWith("abort", addedListener, { once: true })
900-
expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", addedListener)
910+
// The bridge (RequestConfigBuilder.addMergedSignal) uses AbortSignal.any,
911+
// so the external signal must never be managed with manual listeners.
912+
expect(addEventListenerSpy).not.toHaveBeenCalled()
913+
expect(removeEventListenerSpy).not.toHaveBeenCalled()
901914
})
902915

903916
it("should wrap a non-abort stream failure when metadata is provided without a signal", async () => {
@@ -910,11 +923,14 @@ describe("MistralHandler", () => {
910923
expect((error as Error).message).toBe("Mistral completion error: boom")
911924
})
912925

913-
it("should not pass a signal to the stream call when no external signal is provided", async () => {
926+
it("should pass a live request-local signal to the stream call when no external signal is provided", async () => {
914927
const stream = handler.createMessage(systemPrompt, messages)
915928
await collectStream(stream)
916929
const streamOptions = mockCreate.mock.calls[0][1]
917-
expect(streamOptions).toBeUndefined()
930+
// The request-local controller signal is always present so the provider
931+
// keeps its own abort handle; with no external signal it never aborts.
932+
expect(streamOptions?.fetchOptions?.signal).toBeInstanceOf(AbortSignal)
933+
expect(streamOptions?.fetchOptions?.signal?.aborted).toBe(false)
918934
})
919935
})
920936
})

‎src/api/providers/gemini.ts‎

Lines changed: 21 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@ import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, Complete
2727
import { BaseProvider } from "./base-provider"
2828
import { NOT_PROVIDED } from "./constants"
2929
import { parseVertexJsonCredentials } from "./utils/vertex-credentials"
30+
import { createAbortError, isRequestAborted, throwIfAborted } from "./utils/abort-signal"
31+
import { RequestConfigBuilder } from "./config-builder/request-config-builder"
3032
import { getRequestTimeoutMs } from "./utils/request-timeout"
3133

3234
type GeminiHandlerOptions = ApiHandlerOptions & {
@@ -404,29 +406,21 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
404406
}
405407
}
406408

407-
// Bridge the external abort signal from Task (metadata.abortSignal) into a
408-
// request-local controller so the in-flight generateContentStream request
409-
// can be cancelled. The @google/genai SDK merges this signal with its own
410-
// timeout handling, which is preserved rather than replaced.
411-
// A pre-aborted signal rejects immediately with an AbortError.
412-
const externalAbortSignal = metadata?.abortSignal
413-
let requestAbortController: AbortController | undefined
414-
let externalAbortListener: (() => void) | undefined
415-
if (externalAbortSignal) {
416-
if (externalAbortSignal.aborted) {
417-
throw new DOMException("Gemini request aborted", "AbortError")
418-
}
419-
const controller = new AbortController()
420-
requestAbortController = controller
421-
const onExternalAbort = () => controller.abort()
422-
externalAbortListener = onExternalAbort
423-
externalAbortSignal.addEventListener("abort", onExternalAbort, { once: true })
424-
}
409+
// Fast-fail if the request was already aborted before building.
410+
throwIfAborted(metadata?.abortSignal)
411+
412+
// The request-local controller is the provider-owned abort handle; the
413+
// signal the SDK receives merges it with the external Task signal
414+
// (AbortSignal.any inside RequestConfigBuilder), so external aborts
415+
// cancel the in-flight request without manual listener management.
416+
const requestBuilder = new RequestConfigBuilder<{ signal?: AbortSignal }>()
417+
requestBuilder.addMergedSignal(new AbortController(), metadata)
418+
const requestSignal = requestBuilder.getOption("signal")
425419

426420
const params: GenerateContentParameters = {
427421
model,
428422
contents,
429-
config: requestAbortController ? { ...config, abortSignal: requestAbortController.signal } : config,
423+
config: requestSignal ? { ...config, abortSignal: requestSignal } : config,
430424
}
431425

432426
try {
@@ -560,10 +554,10 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
560554
}
561555
}
562556
} catch (error) {
563-
// User-initiated abort: surface a standard AbortError rather than the
564-
// wrapped provider error so callers can distinguish cancellation.
565-
if (metadata?.abortSignal?.aborted) {
566-
throw new DOMException("Gemini request aborted", "AbortError")
557+
// Aborted request: covers both the external signal and SDK-native
558+
// abort errors that may surface before the signal flag propagates.
559+
if (isRequestAborted(error, metadata?.abortSignal)) {
560+
throw createAbortError("Gemini")
567561
}
568562
const errorMessage = error instanceof Error ? error.message : String(error)
569563
const apiError = new ApiProviderError(errorMessage, this.providerName, model, "createMessage")
@@ -574,11 +568,6 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
574568
}
575569

576570
throw error
577-
} finally {
578-
// Stryker disable next-line LogicalOperator: externalAbortListener is only assigned when externalAbortSignal is truthy, so && and || evaluate identically here
579-
if (externalAbortSignal && externalAbortListener) {
580-
externalAbortSignal.removeEventListener("abort", externalAbortListener)
581-
}
582571
}
583572
}
584573

@@ -723,10 +712,10 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
723712

724713
return text
725714
} catch (error) {
726-
// User-initiated abort: surface a standard AbortError rather than the
727-
// wrapped provider error so callers can distinguish cancellation.
728-
if (options?.abortSignal?.aborted) {
729-
throw new DOMException("Gemini completion aborted", "AbortError")
715+
// Aborted request: covers both the external signal and SDK-native
716+
// abort errors that may surface before the signal flag propagates.
717+
if (isRequestAborted(error, options?.abortSignal)) {
718+
throw createAbortError("Gemini")
730719
}
731720
const errorMessage = error instanceof Error ? error.message : String(error)
732721
const apiError = new ApiProviderError(errorMessage, this.providerName, model, "completePrompt")

0 commit comments

Comments
 (0)