fix(api): improve timeout and abort signal handling across providers

Author: easonliang28

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

@@ -193,7 +193,7 @@ describe("GeminiHandler", () => {
 				model: GEMINI_MODEL_NAME,
 				contents: [{ role: "user", parts: [{ text: "test prompt" }] }],
 				config: {
-					httpOptions: { signal: controller.signal },
+					httpOptions: { signal: controller.signal, timeout: 10000 },
 					temperature: 1,
 				},
 			})
@@ -206,7 +206,7 @@ describe("GeminiHandler", () => {
 				model: GEMINI_MODEL_NAME,
 				contents: [{ role: "user", parts: [{ text: "test prompt" }] }],
 				config: {
-					httpOptions: undefined,
+					httpOptions: { timeout: 5000 },
 					temperature: 1,
 				},
 			})

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

@@ -508,7 +508,7 @@ describe("MistralHandler", () => {
 			})
 			await handler.completePrompt("test prompt", { abortSignal: controller.signal })
 			expect(mockComplete).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), {
-				signal: controller.signal,
+				fetchOptions: { signal: controller.signal },
 			})
 		})
 
@@ -528,8 +528,8 @@ describe("MistralHandler", () => {
 			})
 			await handler.completePrompt("test prompt", { abortSignal: controller.signal, timeoutMs: 5000 })
 			expect(mockComplete).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), {
-				signal: controller.signal,
-				timeout: 5000,
+				fetchOptions: { signal: controller.signal },
+				timeoutMs: 5000,
 			})
 		})
 
@@ -539,7 +539,7 @@ describe("MistralHandler", () => {
 			})
 			await handler.completePrompt("test prompt", { timeoutMs: 3000 })
 			expect(mockComplete).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), {
-				timeout: 3000,
+				timeoutMs: 3000,
 			})
 		})
 
@@ -548,10 +548,9 @@ describe("MistralHandler", () => {
 				choices: [{ message: { content: "response" } }],
 			})
 			await handler.completePrompt("test prompt", { timeoutMs: 0 })
-			expect(mockComplete).toHaveBeenCalledWith(
-				expect.objectContaining({ model: expect.any(String) }),
-				undefined, // truthy check means 0 is falsy
-			)
+			expect(mockComplete).toHaveBeenCalledWith(expect.objectContaining({ model: expect.any(String) }), {
+				timeoutMs: 0,
+			})
 		})
 	})
 })

src/api/providers/gemini.ts

@@ -589,6 +589,9 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 			if (options?.abortSignal) {
 				httpOpts.signal = options.abortSignal
 			}
+			if (options?.timeoutMs !== undefined) {
+				httpOpts.timeout = options.timeoutMs
+			}
 			if (this.options.googleGeminiBaseUrl) {
 				httpOpts.baseUrl = this.options.googleGeminiBaseUrl
 			}

src/api/providers/mistral.ts

@@ -197,13 +197,13 @@ export class MistralHandler extends BaseProvider implements SingleCompletionHand
 		const { id: model, temperature } = this.getModel()
 
 		try {
-			// Build request options with abortSignal and/or timeout handling
-			const requestOptions: Record<string, unknown> = {}
+			// Build Mistral SDK RequestOptions
+			const requestOptions: Parameters<typeof this.client.chat.complete>[1] = {}
 			if (options?.abortSignal) {
-				requestOptions.signal = options.abortSignal
+				requestOptions.fetchOptions = { signal: options.abortSignal }
 			}
-			if (options?.timeoutMs) {
-				requestOptions.timeout = options.timeoutMs
+			if (options?.timeoutMs !== undefined) {
+				requestOptions.timeoutMs = options.timeoutMs
 			}
 
 			const response = await this.client.chat.complete(
@@ -212,7 +212,7 @@ export class MistralHandler extends BaseProvider implements SingleCompletionHand
 					messages: [{ role: "user", content: prompt }],
 					temperature,
 				},
-				Object.keys(requestOptions).length > 0 ? (requestOptions as any) : undefined,
+				Object.keys(requestOptions).length > 0 ? requestOptions : undefined,
 			)
 
 			const content = response.choices?.[0]?.message.content

src/api/providers/openai-compatible.ts

@@ -208,19 +208,23 @@ export abstract class OpenAICompatibleHandler extends BaseProvider implements Si
 		}
 
 		// Merge abortSignal and timeoutMs into a single abortSignal
+		let timeoutId: ReturnType<typeof setTimeout> | undefined
 		if (options?.abortSignal && options?.timeoutMs && options.timeoutMs > 0) {
 			// When both are provided, create a merged signal that aborts when either fires
 			const controller = new AbortController()
-			const timeoutId = setTimeout(() => controller.abort(), options.timeoutMs)
-
-			options.abortSignal.addEventListener(
-				"abort",
-				() => {
-					clearTimeout(timeoutId)
-					controller.abort()
-				},
-				{ once: true },
-			)
+			if (options.abortSignal.aborted) {
+				controller.abort()
+			} else {
+				timeoutId = setTimeout(() => controller.abort(), options.timeoutMs)
+				options.abortSignal.addEventListener(
+					"abort",
+					() => {
+						clearTimeout(timeoutId)
+						controller.abort()
+					},
+					{ once: true },
+				)
+			}
 
 			generateOptions.abortSignal = controller.signal
 		} else if (options?.abortSignal) {
@@ -229,8 +233,13 @@ export abstract class OpenAICompatibleHandler extends BaseProvider implements Si
 			generateOptions.abortSignal = AbortSignal.timeout(options.timeoutMs)
 		}
 
-		const { text } = await generateText(generateOptions)
-
-		return text
+		try {
+			const { text } = await generateText(generateOptions)
+			return text
+		} finally {
+			if (timeoutId !== undefined) {
+				clearTimeout(timeoutId)
+			}
+		}
 	}
 }

src/api/providers/openai.ts

@@ -315,7 +315,7 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
 			if (options?.abortSignal) {
 				createOptions.signal = options.abortSignal
 			}
-			if (options?.timeoutMs) {
+			if (options?.timeoutMs !== undefined) {
 				createOptions.timeout = options.timeoutMs
 			}
 

src/api/providers/vscode-lm.ts

@@ -595,19 +595,18 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan
 				}
 			}
 			return result
+		} catch (error: any) {
+			if (error instanceof Error) {
+				throw new Error(`VSCode LM completion error: ${error.message}`)
+			}
+			throw error
 		} finally {
 			if (timeoutTimeout) {
 				clearTimeout(timeoutTimeout)
 			}
 			tokenSource.dispose()
 		}
 	}
-	catch(error: any) {
-		if (error instanceof Error) {
-			throw new Error(`VSCode LM completion error: ${error.message}`)
-		}
-		throw error
-	}
 }
 
 // Static blacklist of VS Code Language Model IDs that should be excluded from the model list e.g. because they will never work