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/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)
+			}
 		}
 	}
 }