feat(api): implement abort/timeout support for native-ollama provider

Author: easonliang28

src/api/providers/__tests__/native-ollama.spec.ts

@@ -8,11 +8,14 @@ import { getOllamaModels } from "../fetchers/ollama"
 
 // Mock the ollama package
 const mockChat = vitest.fn()
+const mockAbort = vitest.fn()
 vitest.mock("ollama", () => {
 	return {
-		Ollama: vitest.fn().mockImplementation(function () {
+		Ollama: vitest.fn().mockImplementation(function (options?: any) {
 			return {
 				chat: mockChat,
+				abort: mockAbort,
+				_host: options?.host ?? "http://localhost:11434",
 			}
 		}),
 		Message: vitest.fn(),
@@ -366,6 +369,78 @@ describe("NativeOllamaHandler", () => {
 			const result = await handler.completePrompt("Test prompt")
 			expect(result).toBe("Response")
 		})
+
+		it("should call client.abort() when timeoutMs is reached", async () => {
+			const testTimeout = 5000
+			let capturedFn: (() => void) | undefined
+
+			vitest.spyOn(global, "setTimeout").mockImplementation((fn: any, ms?: number) => {
+				if (ms === testTimeout) {
+					capturedFn = fn as () => void
+					return 0 as any
+				}
+				return 0 as any
+			})
+
+			mockChat.mockResolvedValue({
+				message: { content: "Response" },
+			})
+
+			await handler.completePrompt("Test prompt", { timeoutMs: testTimeout })
+
+			expect(capturedFn).toBeDefined()
+			if (capturedFn) capturedFn()
+			expect(mockAbort).toHaveBeenCalledTimes(1)
+		})
+
+		it("should call client.abort() when abortSignal is aborted", async () => {
+			const controller = new AbortController()
+			mockChat.mockResolvedValue({
+				message: { content: "Response" },
+			})
+
+			const promise = handler.completePrompt("Test prompt", { abortSignal: controller.signal })
+			controller.abort()
+			await promise
+
+			expect(mockAbort).toHaveBeenCalledTimes(1)
+		})
+
+		it("should call client.abort() immediately when abortSignal is already aborted", async () => {
+			const controller = new AbortController()
+			controller.abort()
+
+			mockChat.mockResolvedValue({
+				message: { content: "Response" },
+			})
+
+			await handler.completePrompt("Test prompt", { abortSignal: controller.signal })
+
+			expect(mockAbort).toHaveBeenCalledTimes(1)
+		})
+
+		it("should clear timeoutId in finally block on success", async () => {
+			let capturedDelay: number | undefined
+
+			vitest.spyOn(global, "setTimeout").mockImplementation((fn: any, ms?: number) => {
+				if (ms === 5000) {
+					capturedDelay = ms
+					return 1 as any // Return truthy value so timeoutId is set
+				}
+				return 0 as any
+			})
+
+			vitest.spyOn(global, "clearTimeout").mockImplementation(() => {})
+
+			mockChat.mockResolvedValue({
+				message: { content: "Response" },
+			})
+
+			await handler.completePrompt("Test prompt", { timeoutMs: 5000 })
+
+			// setTimeout should have been called with the correct delay
+			expect(capturedDelay).toBe(5000)
+		})
 	})
 
 	describe("error handling", () => {

src/api/providers/__tests__/openai-native.spec.ts

@@ -279,36 +279,36 @@ describe("OpenAiNativeHandler", () => {
 			expect(result).toBe("response")
 		})
 
-		it("completePrompt should pass timeoutMs through to client", async () => {
-			mockResponsesCreate.mockResolvedValue({
-				output: [
-					{
-						type: "message",
-						content: [{ type: "output_text", text: "response" }],
-					},
-				],
-			})
-
-			await handler.completePrompt("Test prompt", { timeoutMs: 5000 })
-			// Implementation creates an AbortSignal when timeoutMs is provided.
-			expect(mockResponsesCreate.mock.calls[0][1].signal).toBeInstanceOf(AbortSignal)
-		})
-
-		it("completePrompt should merge signal and timeoutMs together", async () => {
-			const controller = new AbortController()
-			mockResponsesCreate.mockResolvedValue({
-				output: [
-					{
-						type: "message",
-						content: [{ type: "output_text", text: "response" }],
-					},
-				],
-			})
-
-			await handler.completePrompt("Test prompt", { abortSignal: controller.signal, timeoutMs: 10000 })
-			// Implementation uses AbortSignal.any() to merge signals.
-			expect(mockResponsesCreate.mock.calls[0][1].signal).toBeInstanceOf(AbortSignal)
-		})
+						it("completePrompt should pass timeoutMs through to client", async () => {
+							mockResponsesCreate.mockResolvedValue({
+								output: [
+									{
+										type: "message",
+										content: [{ type: "output_text", text: "response" }],
+									},
+								],
+							})
+				
+							await handler.completePrompt("Test prompt", { timeoutMs: 5000 })
+							// Implementation passes a signal to the client (uses baseSignal when no abortSignal provided).
+							expect(mockResponsesCreate.mock.calls[0][1].signal).toBeInstanceOf(AbortSignal)
+						})
+				
+						it("completePrompt should merge signal and timeoutMs together", async () => {
+							const controller = new AbortController()
+							mockResponsesCreate.mockResolvedValue({
+								output: [
+									{
+										type: "message",
+										content: [{ type: "output_text", text: "response" }],
+									},
+								],
+							})
+				
+							await handler.completePrompt("Test prompt", { abortSignal: controller.signal, timeoutMs: 10000 })
+							// Implementation uses AbortSignal.any() to merge signals.
+							expect(mockResponsesCreate.mock.calls[0][1].signal).toBeInstanceOf(AbortSignal)
+						})
 	})
 
 	describe("getModel", () => {

src/api/providers/native-ollama.ts

@@ -347,13 +347,42 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio
 		}
 	}
 
-	async completePrompt(prompt: string, _options?: CompletePromptOptions): Promise<string> {
-		// Ollama native client doesn't support abort signals at all — accept param but ignore
+	async completePrompt(prompt: string, options?: CompletePromptOptions): Promise<string> {
+		// Ollama native client doesn't support external AbortSignal directly.
+		// For per-request cancellation, create a dedicated client instance when abortSignal is provided.
+		const hasAbortSignal = options?.abortSignal !== undefined
+		let localClient: Ollama | undefined
+		let timeoutId: ReturnType<typeof setTimeout> | undefined
+
 		try {
-			const client = this.ensureClient()
+			// Use local client if abortSignal is provided (per-request isolation)
+			const client = hasAbortSignal
+				? (localClient ??= new Ollama({ host: this.options.ollamaBaseUrl }))
+				: this.ensureClient()
 			const { id: modelId } = await this.fetchModel()
 			const useR1Format = modelId.toLowerCase().includes("deepseek-r1")
 
+			// Handle timeoutMs if provided
+			if (options?.timeoutMs !== undefined && options.timeoutMs > 0) {
+				timeoutId = setTimeout(() => client.abort(), options.timeoutMs)
+			}
+
+			// Propagate abortSignal into the local controller via client.abort()
+			if (options?.abortSignal) {
+				if (options.abortSignal.aborted) {
+					client.abort()
+				} else {
+					options.abortSignal.addEventListener(
+						"abort",
+						() => {
+							client.abort()
+							clearTimeout(timeoutId)
+						},
+						{ once: true },
+					)
+				}
+			}
+
 			// Build options object conditionally
 			const chatOptions: OllamaChatOptions = {
 				temperature: this.options.modelTemperature ?? (useR1Format ? DEEP_SEEK_DEFAULT_TEMPERATURE : 0),
@@ -377,6 +406,10 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio
 				throw new Error(`Ollama completion error: ${error.message}`)
 			}
 			throw error
+		} finally {
+			if (timeoutId) {
+				clearTimeout(timeoutId)
+			}
 		}
 	}
 }