test(providers): cover provider completion options

Author: easonliang28

src/api/providers/__tests__/fake-ai.spec.ts

@@ -0,0 +1,79 @@
+import type { Anthropic } from "@anthropic-ai/sdk"
+
+import type { ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../../index"
+import { FakeAIHandler } from "../fake-ai"
+
+const modelInfo = {
+	contextWindow: 8192,
+	maxTokens: 4096,
+	supportsImages: false,
+	supportsPromptCache: false,
+}
+
+describe("FakeAIHandler", () => {
+	it("should delegate completePrompt with options to the cached FakeAI instance", async () => {
+		const completePrompt = vitest.fn().mockResolvedValue("delegated response")
+		const fakeAi: {
+			id: string
+			createMessage: () => AsyncGenerator<never, void, unknown>
+			getModel: () => { id: string; info: typeof modelInfo }
+			countTokens: ReturnType<typeof vitest.fn>
+			completePrompt: typeof completePrompt
+			removeFromCache?: () => void
+		} = {
+			id: "fake-ai-completePrompt-delegation",
+			createMessage: async function* () {},
+			getModel: () => ({ id: "fake-model", info: modelInfo }),
+			countTokens: vitest.fn().mockResolvedValue(0),
+			completePrompt,
+		}
+		const controller = new AbortController()
+		const options: CompletePromptOptions = { abortSignal: controller.signal, timeoutMs: 1234 }
+
+		const handler = new FakeAIHandler({ fakeAi })
+		const result = await handler.completePrompt("Test prompt", options)
+
+		expect(result).toBe("delegated response")
+		expect(completePrompt).toHaveBeenCalledWith("Test prompt", options)
+		fakeAi.removeFromCache?.()
+	})
+
+	it("should delegate createMessage, getModel, and countTokens to FakeAI", async () => {
+		const metadata = { taskId: "task-1" } as ApiHandlerCreateMessageMetadata
+		const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hello" }]
+		const content: Anthropic.Messages.ContentBlockParam[] = [{ type: "text", text: "Hello" }]
+		const createMessage = vitest.fn(async function* () {
+			yield { type: "text" as const, text: "Hello" }
+		})
+		const getModel = vitest.fn(() => ({ id: "fake-model", info: modelInfo }))
+		const countTokens = vitest.fn().mockResolvedValue(7)
+		const fakeAi: {
+			id: string
+			createMessage: typeof createMessage
+			getModel: typeof getModel
+			countTokens: typeof countTokens
+			completePrompt: ReturnType<typeof vitest.fn>
+			removeFromCache?: () => void
+		} = {
+			id: "fake-ai-handler-delegation",
+			createMessage,
+			getModel,
+			countTokens,
+			completePrompt: vitest.fn().mockResolvedValue("complete"),
+		}
+
+		const handler = new FakeAIHandler({ fakeAi })
+		const chunks = []
+		for await (const chunk of handler.createMessage("System", messages, metadata)) {
+			chunks.push(chunk)
+		}
+
+		expect(chunks).toEqual([{ type: "text", text: "Hello" }])
+		expect(createMessage).toHaveBeenCalledWith("System", messages, metadata)
+		expect(handler.getModel()).toEqual({ id: "fake-model", info: modelInfo })
+		expect(getModel).toHaveBeenCalledTimes(1)
+		await expect(handler.countTokens(content)).resolves.toBe(7)
+		expect(countTokens).toHaveBeenCalledWith(content)
+		fakeAi.removeFromCache?.()
+	})
+})

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

@@ -555,6 +555,63 @@ describe("GeminiHandler", () => {
 		})
 	})
 
+	describe("completePrompt request options", () => {
+		it("should pass timeout and baseUrl through httpOptions", async () => {
+			const handlerWithBaseUrl = new GeminiHandler({
+				apiKey: "test-key",
+				apiModelId: GEMINI_MODEL_NAME,
+				geminiApiKey: "test-key",
+				googleGeminiBaseUrl: "https://gemini.example.test",
+			})
+			handlerWithBaseUrl["client"] = handler["client"] as any
+			;(handler["client"].models.generateContent as any).mockResolvedValue({ text: "Response" })
+
+			const result = await handlerWithBaseUrl.completePrompt("Test prompt", { timeoutMs: 1234 })
+
+			expect(result).toBe("Response")
+			expect(handler["client"].models.generateContent).toHaveBeenCalledWith(
+				expect.objectContaining({
+					config: expect.objectContaining({
+						httpOptions: {
+							timeout: 1234,
+							baseUrl: "https://gemini.example.test",
+						},
+					}),
+				}),
+			)
+		})
+
+		it("should pass abortSignal on config instead of httpOptions", async () => {
+			const controller = new AbortController()
+			;(handler["client"].models.generateContent as any).mockResolvedValue({ text: "Response" })
+
+			await handler.completePrompt("Test prompt", { abortSignal: controller.signal })
+
+			expect(handler["client"].models.generateContent).toHaveBeenCalledWith(
+				expect.objectContaining({
+					config: expect.objectContaining({
+						abortSignal: controller.signal,
+						httpOptions: undefined,
+					}),
+				}),
+			)
+		})
+
+		it("should omit httpOptions when timeoutMs and baseUrl are not provided", async () => {
+			;(handler["client"].models.generateContent as any).mockResolvedValue({ text: "Response" })
+
+			await handler.completePrompt("Test prompt")
+
+			expect(handler["client"].models.generateContent).toHaveBeenCalledWith(
+				expect.objectContaining({
+					config: expect.objectContaining({
+						httpOptions: undefined,
+					}),
+				}),
+			)
+		})
+	})
+
 	describe("error telemetry", () => {
 		const mockMessages: Anthropic.Messages.MessageParam[] = [
 			{

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

@@ -769,16 +769,17 @@ describe("NativeOllamaHandler", () => {
 			)
 		})
 
-		it("should accept options param but ignore it (no signal support)", async () => {
+		it("should use a request-local client when abortSignal is provided", async () => {
 			mockChat.mockResolvedValue({
 				message: { content: "Response" },
 			})
 
 			const controller = new AbortController()
 			await handler.completePrompt("Test prompt", { abortSignal: controller.signal })
 
-			// Verify that the call does NOT include any signal-related options
-			// Ollama implementation only passes the payload, not a second options argument
+			expect(OllamaMock).toHaveBeenCalledTimes(1)
+			expect(OllamaMock).toHaveBeenCalledWith({ host: "http://localhost:11434" })
+			// Ollama implementation only passes the payload, not a second options argument.
 			expect(mockChat).toHaveBeenCalledWith(
 				expect.objectContaining({
 					model: "llama2",
@@ -787,7 +788,6 @@ describe("NativeOllamaHandler", () => {
 					options: { temperature: 0 },
 				}),
 			)
-			// Verify no second argument was passed (no signal/options forwarded)
 			expect(mockChat).toHaveBeenCalledTimes(1)
 			expect(mockChat.mock.calls[0]).toHaveLength(1)
 		})
@@ -899,6 +899,47 @@ describe("NativeOllamaHandler", () => {
 			expect(capturedInstanceAbort!).toHaveBeenCalledTimes(1)
 		})
 
+		it("should not create a request-local client or timer for non-positive timeoutMs", async () => {
+			const setTimeoutSpy = vitest.spyOn(global, "setTimeout")
+			mockChat.mockResolvedValue({
+				message: { content: "Response" },
+			})
+
+			await handler.completePrompt("Test prompt", { timeoutMs: 0 })
+
+			expect(setTimeoutSpy).not.toHaveBeenCalled()
+			expect(OllamaMock).toHaveBeenCalledTimes(1)
+			expect(OllamaMock).toHaveBeenCalledWith({ host: "http://localhost:11434" })
+		})
+
+		it("should remove abort listener and clear timeout when abortSignal fires", async () => {
+			const controller = new AbortController()
+			const timeoutHandle = 1 as any
+			const clearTimeoutSpy = vitest.spyOn(global, "clearTimeout").mockImplementation(() => {})
+			vitest.spyOn(global, "setTimeout").mockImplementation(() => timeoutHandle)
+			const removeEventListenerSpy = vitest.spyOn(controller.signal, "removeEventListener")
+
+			let resolveChat: (value: { message: { content: string } }) => void = () => {}
+			mockChat.mockImplementation(
+				() =>
+					new Promise((resolve) => {
+						resolveChat = resolve
+					}),
+			)
+
+			const promise = handler.completePrompt("Test prompt", { abortSignal: controller.signal, timeoutMs: 5000 })
+			for (let i = 0; i < 10 && mockChat.mock.calls.length === 0; i++) {
+				await Promise.resolve()
+			}
+
+			controller.abort()
+			resolveChat({ message: { content: "Response" } })
+
+			await expect(promise).resolves.toBe("Response")
+			expect(clearTimeoutSpy).toHaveBeenCalledWith(timeoutHandle)
+			expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", expect.any(Function))
+		})
+
 		it("should clear timeoutId in finally block on success", async () => {
 			let capturedDelay: number | undefined
 			const timeoutHandle = 1 as any