diff --git a/src/api/providers/__tests__/native-ollama.spec.ts b/src/api/providers/__tests__/native-ollama.spec.ts index 8fcbf4a0a1..c8590c14a7 100644 --- a/src/api/providers/__tests__/native-ollama.spec.ts +++ b/src/api/providers/__tests__/native-ollama.spec.ts @@ -6,25 +6,47 @@ import { NativeOllamaHandler } from "../native-ollama" import { ApiHandlerOptions } from "../../../shared/api" import { getOllamaModels } from "../fetchers/ollama" +import { makeCreateMessageMetadata } from "../../../test-utils/api" import { clearAllMocks } from "../../../test-utils/reset" // Mock the ollama package const mockChat = vitest.fn() + +// Use vi.hoisted to define mocks that can be referenced both inside and outside the hoisted vi.mock +const mockedOllama = vi.hoisted(() => ({ + OllamaMock: vitest.fn(), +})) + vitest.mock("ollama", () => { + const { OllamaMock } = mockedOllama return { - Ollama: vitest.fn().mockImplementation(function () { + Ollama: OllamaMock.mockImplementation(function (options?: { host?: string }) { + const instanceAbort = vitest.fn() return { chat: mockChat, + abort: instanceAbort, + _host: options?.host ?? "http://localhost:11434", + _instanceAbort: instanceAbort, } }), Message: vitest.fn(), } }) -// Mock the getOllamaModels function -vitest.mock("../fetchers/ollama", () => ({ - getOllamaModels: vitest.fn(), -})) +// Export OllamaMock for test access +const OllamaMock = mockedOllama.OllamaMock + +// Mock only the model-list fetch. The real isSecureOllamaEndpoint is kept so +// the credential-gating assertions exercise the production predicate. +vitest.mock("../fetchers/ollama", async (importOriginal) => { + // The factory's importOriginal is typed as unknown; recover the module's + // real type so the production predicate can be spread without `any`. + const actual = (await importOriginal()) as typeof import("../fetchers/ollama") + return { + ...actual, + getOllamaModels: vitest.fn(), + } +}) const mockGetOllamaModels = vitest.mocked(getOllamaModels) @@ -696,6 +718,11 @@ describe("NativeOllamaHandler", () => { }) describe("completePrompt", () => { + // Shared capture for the per-request client's abort spy. The timeoutMs + // and abortSignal tests below each override the OllamaMock constructor + // and re-assign it inside their own mockImplementation. + let capturedInstanceAbort: (() => void) | undefined + it("should complete a prompt without streaming", async () => { mockChat.mockResolvedValue({ message: { content: "This is the response" }, @@ -756,6 +783,237 @@ describe("NativeOllamaHandler", () => { }), ) }) + 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 }) + + expect(OllamaMock).toHaveBeenCalledTimes(1) + expect(OllamaMock).toHaveBeenCalledWith(expect.objectContaining({ host: "http://localhost:11434" })) + // Ollama implementation only passes the payload, not a second options argument. + expect(mockChat).toHaveBeenCalledWith( + expect.objectContaining({ + model: "llama2", + messages: [{ role: "user", content: "Test prompt" }], + stream: false, + options: { temperature: 0 }, + }), + ) + expect(mockChat).toHaveBeenCalledTimes(1) + expect(mockChat.mock.calls[0]).toHaveLength(1) + }) + + it("should not include signal-related options when not provided", async () => { + mockChat.mockResolvedValue({ + message: { content: "Response" }, + }) + + await handler.completePrompt("Test prompt") + + expect(mockChat).toHaveBeenCalledWith( + expect.objectContaining({ + model: "llama2", + messages: [{ role: "user", content: "Test prompt" }], + stream: false, + options: { temperature: 0 }, + }), + ) + }) + + it("should work without options (backward compatible)", async () => { + mockChat.mockResolvedValue({ + message: { content: "Response" }, + }) + + 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 + + // Capture the per-request client's abort spy so the assertion targets + // the actual abort() call rather than just the constructor call. + OllamaMock.mockImplementation(function (options?: { host?: string }) { + const instanceAbort = vitest.fn() + capturedInstanceAbort = instanceAbort + return { + chat: mockChat, + abort: instanceAbort, + _host: options?.host ?? "http://localhost:11434", + _instanceAbort: instanceAbort, + } + }) + + // The timer id is never consumed (the callback is captured and fired + // manually), so bridge the ambient setTimeout signature through unknown. + vitest.spyOn(global, "setTimeout").mockImplementation(((fn: () => void, ms?: number) => { + if (ms === testTimeout) { + capturedFn = fn + } + return 0 + }) as unknown as typeof setTimeout) + + mockChat.mockResolvedValue({ + message: { content: "Response" }, + }) + + await handler.completePrompt("Test prompt", { timeoutMs: testTimeout }) + + expect(capturedFn).toBeDefined() + if (capturedFn) capturedFn() + // The timeout callback should have invoked client.abort() on the request-local instance + expect(capturedInstanceAbort).toBeDefined() + expect(capturedInstanceAbort).toHaveBeenCalledTimes(1) + }) + + it("should call instance.abort() when abortSignal is aborted", async () => { + const controller = new AbortController() + + // Override the constructor to capture the instance abort spy + OllamaMock.mockImplementation(function (options?: { host?: string }) { + const instanceAbort = vitest.fn() + capturedInstanceAbort = instanceAbort + return { + chat: mockChat, + abort: instanceAbort, + _host: options?.host ?? "http://localhost:11434", + _instanceAbort: instanceAbort, + } + }) + + mockChat.mockResolvedValue({ + message: { content: "Response" }, + }) + + const promise = handler.completePrompt("Test prompt", { abortSignal: controller.signal }) + controller.abort() + await expect(promise).rejects.toThrow("This operation was aborted") + + expect(capturedInstanceAbort).toBeDefined() + expect(capturedInstanceAbort).toHaveBeenCalledTimes(1) + }) + + it("should call instance.abort() immediately when abortSignal is already aborted", async () => { + const controller = new AbortController() + controller.abort() + + // Override the constructor to capture the instance abort spy + OllamaMock.mockImplementation(function (options?: { host?: string }) { + const instanceAbort = vitest.fn() + capturedInstanceAbort = instanceAbort + return { + chat: mockChat, + abort: instanceAbort, + _host: options?.host ?? "http://localhost:11434", + _instanceAbort: instanceAbort, + } + }) + + mockChat.mockResolvedValue({ + message: { content: "Response" }, + }) + + await expect(handler.completePrompt("Test prompt", { abortSignal: controller.signal })).rejects.toThrow( + "This operation was aborted", + ) + + expect(capturedInstanceAbort).toBeDefined() + expect(capturedInstanceAbort).toHaveBeenCalledTimes(1) + }) + + it("should reject with AbortError without fetching models when abortSignal is already aborted", async () => { + const controller = new AbortController() + controller.abort() + + await expect( + handler.completePrompt("Test prompt", { abortSignal: controller.signal }), + ).rejects.toMatchObject({ + name: "AbortError", + }) + + // The pre-aborted check must short-circuit before any network call. + expect(mockGetOllamaModels).not.toHaveBeenCalled() + expect(mockChat).not.toHaveBeenCalled() + }) + + it("should not start a timeout 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(expect.objectContaining({ host: "http://localhost:11434" })) + }) + + it("should remove abort listener and clear timeout when abortSignal fires", async () => { + const controller = new AbortController() + // The timer id is never consumed directly (clearTimeout is mocked in these + // tests), so bridge the ambient setTimeout return type through unknown. + const timeoutHandle = 1 as unknown as ReturnType + const clearTimeoutSpy = vitest.spyOn(global, "clearTimeout").mockImplementation(() => {}) + vitest.spyOn(global, "setTimeout").mockImplementation(() => timeoutHandle) + const removeEventListenerSpy = vitest.spyOn(controller.signal, "removeEventListener") + const addEventListenerSpy = vitest.spyOn(controller.signal, "addEventListener") + + 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) + // The listener removed in the finally block must be the exact function + // that was registered, proving the same reference is detached. + expect(addEventListenerSpy).toHaveBeenCalledWith("abort", expect.any(Function), { once: true }) + const registeredHandler = addEventListenerSpy.mock.calls[0]?.[1] + expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", registeredHandler) + }) + + it("should clear timeoutId in finally block on success", async () => { + let capturedDelay: number | undefined + // The timer id is never consumed directly (clearTimeout is mocked in these + // tests), so bridge the ambient setTimeout return type through unknown. + const timeoutHandle = 1 as unknown as ReturnType + + vitest.spyOn(global, "setTimeout").mockImplementation(((fn: () => void, ms?: number) => { + if (ms === 5000) { + capturedDelay = ms + } + return timeoutHandle + }) as unknown as typeof setTimeout) + + const clearTimeoutSpy = 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) + expect(clearTimeoutSpy).toHaveBeenCalledWith(timeoutHandle) + }) }) it("should send think parameter in completePrompt when reasoningEffort is set", async () => { @@ -1602,5 +1860,590 @@ describe("NativeOllamaHandler", () => { expect(userImageMessages).toHaveLength(1) expect(userImageMessages[0].images).toEqual(["img-a"]) }) + + it("should reject with AbortError when the abort signal fires during model discovery", async () => { + let resolveFetch: (() => void) | undefined + + // Hold model discovery open so the abort lands before the chat request. + mockGetOllamaModels.mockImplementation( + () => + new Promise((resolve) => { + resolveFetch = () => resolve({}) + }), + ) + + const controller = new AbortController() + const promise = handler.completePrompt("Test prompt", { abortSignal: controller.signal }) + + // The discovery fetch is in flight; abort before it settles. + controller.abort() + resolveFetch?.() + + await expect(promise).rejects.toMatchObject({ name: "AbortError" }) + expect(mockGetOllamaModels).toHaveBeenCalledTimes(1) + expect(mockChat).not.toHaveBeenCalled() + }) + + it("should reject with AbortError when timeoutMs fires during model discovery", async () => { + const testTimeout = 5000 + let capturedFn: (() => void) | undefined + let resolveFetch: (() => void) | undefined + + // Hold model discovery open so the timeout lands before the chat request. + mockGetOllamaModels.mockImplementation( + () => + new Promise((resolve) => { + resolveFetch = () => resolve({}) + }), + ) + + // The timer id is never consumed (the callback is captured and fired + // manually), so bridge the ambient setTimeout signature through unknown. + vitest.spyOn(global, "setTimeout").mockImplementation(((fn: () => void, ms?: number) => { + if (ms === testTimeout) { + capturedFn = fn + } + return 0 + }) as unknown as typeof setTimeout) + + const promise = handler.completePrompt("Test prompt", { timeoutMs: testTimeout }) + expect(capturedFn).toBeDefined() + + // Fire the timeout while discovery is still in flight. The discovery + // race must reject on the timeout WITHOUT waiting for the held fetch + // to settle. + capturedFn?.() + await expect(promise).rejects.toMatchObject({ name: "AbortError" }) + expect(mockChat).not.toHaveBeenCalled() + + // Settle the held fetch after the assertion so no promise is left + // dangling. + resolveFetch?.() + }) + }) + + describe("per-request client creation", () => { + it("should create a new Ollama client for each completePrompt call (per-request pattern)", async () => { + mockChat.mockResolvedValue({ + message: { content: "Response" }, + }) + + const handler = new NativeOllamaHandler({ + apiModelId: "llama2", + ollamaModelId: "llama2", + ollamaBaseUrl: "http://localhost:11434", + }) + + // First call + await handler.completePrompt("Test prompt 1") + + // Second call - should create a new client each time (per-request pattern) + await handler.completePrompt("Test prompt 2") + + // Verify the Ollama constructor was called twice (per-request pattern, not singleton) + expect(OllamaMock).toHaveBeenCalledTimes(2) + }) + + it("should pass API key through constructor headers option", async () => { + mockChat.mockResolvedValue({ + message: { content: "Response" }, + }) + + const handler = new NativeOllamaHandler({ + apiModelId: "llama2", + ollamaModelId: "llama2", + ollamaBaseUrl: "http://localhost:11434", + ollamaApiKey: "test-api-key-123", + }) + + await handler.completePrompt("Test prompt") + + // Verify Ollama was constructed with headers containing the API key + expect(OllamaMock).toHaveBeenCalledWith( + expect.objectContaining({ + headers: { + Authorization: "Bearer test-api-key-123", + }, + }), + ) + }) + + it("should work without API key (no headers)", async () => { + mockChat.mockResolvedValue({ + message: { content: "Response" }, + }) + + const handler = new NativeOllamaHandler({ + apiModelId: "llama2", + ollamaModelId: "llama2", + ollamaBaseUrl: "http://localhost:11434", + }) + + await handler.completePrompt("Test prompt") + + // Verify Ollama was constructed without headers when no API key is provided + expect(OllamaMock).toHaveBeenCalledWith( + expect.objectContaining({ + host: "http://localhost:11434", + }), + ) + // headers should not be present when no API key + const callArgs = OllamaMock.mock.calls[0][0] + expect(callArgs.headers).toBeUndefined() + }) + + it("should use custom baseUrl in client options", async () => { + mockChat.mockResolvedValue({ + message: { content: "Response" }, + }) + + const handler = new NativeOllamaHandler({ + apiModelId: "llama2", + ollamaModelId: "llama2", + ollamaBaseUrl: "http://custom-ollama:11434", + }) + + await handler.completePrompt("Test prompt") + + expect(OllamaMock).toHaveBeenCalledWith( + expect.objectContaining({ + host: "http://custom-ollama:11434", + }), + ) + }) + + it("should not attach the API key for a remote HTTP endpoint", async () => { + mockChat.mockResolvedValue({ + message: { content: "Response" }, + }) + + const handler = new NativeOllamaHandler({ + apiModelId: "llama2", + ollamaModelId: "llama2", + ollamaBaseUrl: "http://ollama.example.com:11434", + ollamaApiKey: "test-api-key-123", + }) + + await handler.completePrompt("Test prompt") + + // Plaintext HTTP to a remote host would leak the credential (CWE-319), + // so the Authorization header must be omitted. + const callArgs = OllamaMock.mock.calls[0][0] + expect(callArgs.headers).toBeUndefined() + }) + + it("should attach the API key for an HTTPS endpoint", async () => { + mockChat.mockResolvedValue({ + message: { content: "Response" }, + }) + + const handler = new NativeOllamaHandler({ + apiModelId: "llama2", + ollamaModelId: "llama2", + ollamaBaseUrl: "https://ollama.example.com:11434", + ollamaApiKey: "test-api-key-123", + }) + + await handler.completePrompt("Test prompt") + + expect(OllamaMock).toHaveBeenCalledWith( + expect.objectContaining({ + host: "https://ollama.example.com:11434", + headers: { + Authorization: "Bearer test-api-key-123", + }, + }), + ) + }) + + it("should attach the API key for a loopback IP endpoint", async () => { + mockChat.mockResolvedValue({ + message: { content: "Response" }, + }) + + const handler = new NativeOllamaHandler({ + apiModelId: "llama2", + ollamaModelId: "llama2", + ollamaBaseUrl: "http://127.0.0.1:11434", + ollamaApiKey: "test-api-key-123", + }) + + await handler.completePrompt("Test prompt") + + expect(OllamaMock).toHaveBeenCalledWith( + expect.objectContaining({ + host: "http://127.0.0.1:11434", + headers: { + Authorization: "Bearer test-api-key-123", + }, + }), + ) + }) + }) + + describe("per-request abortable transport", () => { + // Earlier tests in this file install a persistent global setTimeout spy + // (vi.spyOn; vi.clearAllMocks does not restore spy implementations), so + // restore the native mocks before each test to keep this suite + // order-independent. + beforeEach(() => { + vi.restoreAllMocks() + }) + + // Drop stubbed globals (the injected fetch) after each test so a fake + // transport cannot leak into later tests. + afterEach(() => { + vi.unstubAllGlobals() + }) + + // The vi.fn() OllamaMock is untyped, so the constructor options read + // from mock.calls need one documented narrow cast to reach the injected + // per-request transport. + type ConstructorOptions = { + fetch?: (url: string, init?: RequestInit) => Promise + } + + // A held-open fetch: it rejects with an AbortError when the passed + // signal aborts and stays pending otherwise. + const heldOpenFetch = () => + vi.fn(async (_url: string, init?: RequestInit) => { + if (init?.signal?.aborted) { + throw new DOMException("This operation was aborted", "AbortError") + } + await new Promise((_resolve, reject) => { + init?.signal?.addEventListener( + "abort", + () => reject(new DOMException("This operation was aborted", "AbortError")), + { once: true }, + ) + }) + }) + + it("should make a held-open chat POST abortable via the per-request transport when the external signal aborts", async () => { + const fetchSpy = heldOpenFetch() + vi.stubGlobal("fetch", fetchSpy) + + let settleChat: ((value: object) => void) | undefined + mockChat.mockImplementation( + () => + new Promise((resolve) => { + settleChat = resolve + }), + ) + + const controller = new AbortController() + const stream = handler.createMessage( + "System", + [{ role: "user" as const, content: "Test" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + const consume = (async () => { + for await (const _ of stream) { + // consume stream + } + })() + + // Spin until the mocked SDK chat() has been invoked (the mocked SDK + // never uses the transport, so its chat() is held manually). + for (let i = 0; i < 50 && mockChat.mock.calls.length === 0; i++) { + await Promise.resolve() + } + expect(mockChat).toHaveBeenCalledTimes(1) + + const clientOptions = OllamaMock.mock.calls[0][0] as ConstructorOptions + expect(clientOptions.fetch).toBeTypeOf("function") + + // Non-stream phase: the SDK passes no signal, so the transport must + // still hand the per-request signal to fetch. + const heldPost = clientOptions.fetch!("http://localhost:11434/api/chat", {}) + controller.abort() + + // The held POST is aborted with the response never released. + await expect(heldPost).rejects.toMatchObject({ name: "AbortError" }) + + // Prove the dispose path: a stream that resolves after cancellation + // is disposed, not consumed. + settleChat?.({ + message: { content: "late" }, + abort: vi.fn(), + [Symbol.asyncIterator]: async function* () {}, + }) + await expect(consume).rejects.toMatchObject({ name: "AbortError" }) + }) + + it("should abort the per-request transport when the consumer stops iterating early", async () => { + const fetchSpy = heldOpenFetch() + vi.stubGlobal("fetch", fetchSpy) + mockChat.mockResolvedValue({ + message: { content: "Response" }, + abort: vi.fn(), + [Symbol.asyncIterator]: async function* () { + yield { message: { content: "first" } } + yield { message: { content: "second" } } + }, + }) + + // No metadata → no external abort bridge: finalization must still + // release the transport. + const stream = handler.createMessage("System", [{ role: "user" as const, content: "Test" }]) + const iterator = stream[Symbol.asyncIterator]() + await iterator.next() + + const clientOptions = OllamaMock.mock.calls[0][0] as ConstructorOptions + const heldPost = clientOptions.fetch!("http://localhost:11434/api/chat", {}) + // Finalization aborts the per-request controller, which must reject + // the held POST. Attach the expectation before the finalization + // triggers the rejection so it cannot surface as an unhandled + // rejection. + const released = expect(heldPost).rejects.toMatchObject({ name: "AbortError" }) + await iterator.return?.(undefined) + await released + }) + + it("should abort a held-open POST via the per-request transport when timeoutMs fires", async () => { + const fetchSpy = heldOpenFetch() + vi.stubGlobal("fetch", fetchSpy) + + let rejectChat: ((reason: unknown) => void) | undefined + mockChat.mockImplementation( + () => + new Promise((_resolve, reject) => { + rejectChat = reject + }), + ) + + const promise = handler.completePrompt("Test prompt", { timeoutMs: 100 }) + + // Spin until the mocked SDK chat() has been invoked. + for (let i = 0; i < 50 && mockChat.mock.calls.length === 0; i++) { + await Promise.resolve() + } + expect(mockChat).toHaveBeenCalledTimes(1) + + const clientOptions = OllamaMock.mock.calls[0][0] as ConstructorOptions + const heldPost = clientOptions.fetch!("http://localhost:11434/api/chat", {}) + + // The real 100ms timer aborts the per-request signal, which must + // reject the held POST. + await expect(heldPost).rejects.toMatchObject({ name: "AbortError" }) + + // Mirror the real SDK: the in-flight POST rejects, and the + // completion surfaces the AbortError unmodified. + rejectChat?.(new DOMException("This operation was aborted", "AbortError")) + await expect(promise).rejects.toMatchObject({ name: "AbortError" }) + }) + + it("should preserve the SDK stream signal when merging it with the per-request signal", async () => { + let receivedSignal: AbortSignal | null | undefined + const fetchSpy = vi.fn(async (_url: string, init?: RequestInit) => { + receivedSignal = init?.signal + if (init?.signal?.aborted) { + throw new DOMException("This operation was aborted", "AbortError") + } + await new Promise((_resolve, reject) => { + init?.signal?.addEventListener( + "abort", + () => reject(new DOMException("This operation was aborted", "AbortError")), + { once: true }, + ) + }) + }) + vi.stubGlobal("fetch", fetchSpy) + + mockChat.mockResolvedValue({ message: { content: "Response" } }) + + // Drive a plain completePrompt so the per-request client is + // constructed; the transport is attached even without an external + // signal or timeout. + await handler.completePrompt("Test prompt") + + const clientOptions = OllamaMock.mock.calls[0][0] as ConstructorOptions + expect(clientOptions.fetch).toBeTypeOf("function") + + const sdkController = new AbortController() + const post = clientOptions.fetch!("http://localhost:11434/api/chat", { signal: sdkController.signal }) + + // The transport must merge, not replace: the SDK's own signal is + // preserved as one side of the AbortSignal.any merge. + expect(receivedSignal).toBeInstanceOf(AbortSignal) + expect(receivedSignal).not.toBe(sdkController.signal) + expect(receivedSignal?.aborted).toBe(false) + + // The SDK-side internal controller must still cancel the POST (the + // preserved-signal side of the merge). + sdkController.abort() + await expect(post).rejects.toMatchObject({ name: "AbortError" }) + }) + }) + + describe("createMessage abort signal", () => { + it("should reject with AbortError when external abortSignal is already aborted", async () => { + const controller = new AbortController() + controller.abort() + + const stream = handler.createMessage( + "System", + [{ role: "user" as const, content: "Test" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + await expect(async () => { + for await (const _ of stream) { + // consume stream + } + }).rejects.toMatchObject({ name: "AbortError" }) + }) + + it("should abort the in-flight request when external abortSignal fires", async () => { + let rejectChat: ((reason: unknown) => void) | undefined + + // Wire the per-request client's abort() to reject the in-flight chat + // request, mirroring how the real Ollama SDK surfaces client.abort(). + OllamaMock.mockImplementation(function (options?: { host?: string }) { + return { + chat: mockChat, + abort: () => { + rejectChat?.(new DOMException("This operation was aborted", "AbortError")) + }, + _host: options?.host ?? "http://localhost:11434", + } + }) + mockChat.mockImplementation( + () => + new Promise((_, reject) => { + rejectChat = reject + }), + ) + + const controller = new AbortController() + const stream = handler.createMessage( + "System", + [{ role: "user" as const, content: "Test" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const consume = (async () => { + for await (const _ of stream) { + // consume stream + } + })() + + // Wait until the in-flight request has actually started. + for (let i = 0; i < 50 && mockChat.mock.calls.length === 0; i++) { + await Promise.resolve() + } + expect(mockChat).toHaveBeenCalledTimes(1) + + controller.abort() + + await expect(consume).rejects.toMatchObject({ name: "AbortError" }) + }) + + it("should reject with AbortError when the abort signal fires during model discovery", async () => { + let resolveFetch: (() => void) | undefined + let clientAborted = false + + // Hold model discovery open so the abort lands between fetchModel() + // starting and the chat request being issued. + mockGetOllamaModels.mockImplementation( + () => + new Promise((resolve) => { + resolveFetch = () => resolve({}) + }), + ) + + // Wire the per-request client's abort() so the bridge into the SDK + // client can be observed; the chat request must never start. + OllamaMock.mockImplementation(function (options?: { host?: string }) { + return { + chat: mockChat, + abort: () => { + clientAborted = true + }, + _host: options?.host ?? "http://localhost:11434", + } + }) + + const controller = new AbortController() + const stream = handler.createMessage( + "System", + [{ role: "user" as const, content: "Test" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const consume = (async () => { + for await (const _ of stream) { + // consume stream + } + })() + + // Abort while the model list is still being fetched. The race must + // reject the generator with AbortError before chat() is attempted. + controller.abort() + await expect(consume).rejects.toMatchObject({ name: "AbortError" }) + expect(mockChat).not.toHaveBeenCalled() + expect(mockGetOllamaModels).toHaveBeenCalledTimes(1) + expect(clientAborted).toBe(true) + + // Settle the held discovery fetch so no promise is left dangling. + resolveFetch?.() + }) + + it("should dispose a stream that resolves after cancellation instead of consuming it", async () => { + const controller = new AbortController() + let streamAborted = false + let consumed = false + + // The POST resolves exactly when the abort lands (response headers + // first): the resolved iterator must be disposed, not consumed. + mockChat.mockImplementation(() => { + controller.abort() + return Promise.resolve({ + message: { content: "late" }, + abort: () => { + streamAborted = true + }, + [Symbol.asyncIterator]: async function* () { + consumed = true + yield { message: { content: "late" } } + }, + }) + }) + + const stream = handler.createMessage( + "System", + [{ role: "user" as const, content: "Test" }], + makeCreateMessageMetadata({ abortSignal: controller.signal }), + ) + + const consume = (async () => { + for await (const _ of stream) { + // consume stream + } + })() + + await expect(consume).rejects.toMatchObject({ name: "AbortError" }) + expect(streamAborted).toBe(true) + expect(consumed).toBe(false) + }) + + it("should surface an AbortError thrown by the stream body without wrapping it", async () => { + // Mirror the real SDK: the chat() iterator is backed by the in-flight + // POST, so an aborted request rejects the first next() with the + // DOMException AbortError. The zero-item delegation satisfies + // require-yield without emitting a part before the rejection. + mockChat.mockImplementation(async function* () { + yield* [] + throw new DOMException("This operation was aborted", "AbortError") + }) + + const stream = handler.createMessage("System", [{ role: "user" as const, content: "Test" }]) + + await expect(async () => { + for await (const _ of stream) { + // consume stream + } + }).rejects.toMatchObject({ name: "AbortError" }) + }) }) }) diff --git a/src/api/providers/fetchers/__tests__/ollama.test.ts b/src/api/providers/fetchers/__tests__/ollama.test.ts index 9c0b547e88..dc781ae18d 100644 --- a/src/api/providers/fetchers/__tests__/ollama.test.ts +++ b/src/api/providers/fetchers/__tests__/ollama.test.ts @@ -1,6 +1,6 @@ import axios from "axios" -import { getOllamaModels, parseOllamaModel } from "../ollama" +import { getOllamaModels, isSecureOllamaEndpoint, parseOllamaModel } from "../ollama" import ollamaModelsData from "./fixtures/ollama-model-details.json" // Mock axios @@ -116,7 +116,148 @@ describe("Ollama Fetcher", () => { }) }) + describe("isSecureOllamaEndpoint", () => { + it("should treat loopback HTTP endpoints as safe", () => { + expect(isSecureOllamaEndpoint("http://localhost:11434")).toBe(true) + expect(isSecureOllamaEndpoint("http://127.0.0.1:11434")).toBe(true) + expect(isSecureOllamaEndpoint("http://[::1]:11434")).toBe(true) + }) + + it("should only treat strict 127.0.0.0/8 IPv4 literals as safe", () => { + // Regression: the old /^127\./ prefix also matched DNS names such as + // 127.example.com, which would leak the key over cleartext HTTP. + expect(isSecureOllamaEndpoint("http://127.example.com:11434")).toBe(false) + expect(isSecureOllamaEndpoint("http://127.99.99.99:11434")).toBe(true) // rest of 127.0.0.0/8 + expect(isSecureOllamaEndpoint("http://127.0.0.255:11434")).toBe(true) // range edge + expect(isSecureOllamaEndpoint("http://1270.0.0.1:11434")).toBe(false) // 1270 is not a valid octet + }) + + it("should treat any HTTPS endpoint as safe", () => { + expect(isSecureOllamaEndpoint("https://ollama.example.com")).toBe(true) + expect(isSecureOllamaEndpoint("https://10.0.0.5:11434")).toBe(true) + }) + + it("should treat remote HTTP endpoints as unsafe", () => { + expect(isSecureOllamaEndpoint("http://ollama.example.com:11434")).toBe(false) + expect(isSecureOllamaEndpoint("http://10.0.0.5:11434")).toBe(false) + }) + + it("should treat unparseable endpoints as unsafe", () => { + expect(isSecureOllamaEndpoint("not a url")).toBe(false) + }) + }) + describe("getOllamaModels", () => { + // Shared /api/tags and /api/show mock payloads. Each call returns fresh + // objects so tests cannot share mutable state. + const makeOllamaTagsPayload = (modelName: string) => ({ + models: [ + { + name: modelName, + model: modelName, + modified_at: "2025-06-03T09:23:22.610222878-04:00", + size: 14333928010, + digest: "6a5f0c01d2c96c687d79e32fdd25b87087feb376bf9838f854d10be8cf3c10a5", + details: { + family: "llama", + families: ["llama"], + format: "gguf", + parameter_size: "23.6B", + parent_model: "", + quantization_level: "Q4_K_M", + }, + }, + ], + }) + const makeOllamaShowPayload = () => ({ + license: "Mock License", + modelfile: "FROM /path/to/blob\nTEMPLATE {{ .Prompt }}", + parameters: "num_ctx 4096\nstop_token ", + template: "{{ .System }}USER: {{ .Prompt }}ASSISTANT:", + modified_at: "2025-06-03T09:23:22.610222878-04:00", + details: { + parent_model: "", + format: "gguf", + family: "llama", + families: ["llama"], + parameter_size: "23.6B", + quantization_level: "Q4_K_M", + }, + model_info: { + "ollama.context_length": 4096, + "some.other.info": "value", + }, + capabilities: ["completion", "tools"], // Has tools capability + }) + + it("should omit the Authorization header for a remote HTTP endpoint", async () => { + const baseUrl = "http://ollama.example.com:11434" + const apiKey = "test-api-key-123" + + mockedAxios.get.mockResolvedValueOnce({ data: { models: [] } }) + + const result = await getOllamaModels(baseUrl, apiKey) + + expect(mockedAxios.get).toHaveBeenCalledTimes(1) + expect(mockedAxios.get).toHaveBeenCalledWith(`${baseUrl}/api/tags`, { headers: {} }) + expect(result).toEqual({}) + }) + + it("should include the Authorization header for an HTTPS endpoint", async () => { + const baseUrl = "https://ollama.example.com:11434" + const apiKey = "test-api-key-123" + + mockedAxios.get.mockResolvedValueOnce({ data: { models: [] } }) + + const result = await getOllamaModels(baseUrl, apiKey) + + expect(mockedAxios.get).toHaveBeenCalledTimes(1) + expect(mockedAxios.get).toHaveBeenCalledWith(`${baseUrl}/api/tags`, { + headers: { Authorization: `Bearer ${apiKey}` }, + }) + expect(result).toEqual({}) + }) + + it.each(["http://localhost:11434", "http://[::1]:11434"])( + "should send the Authorization header and disable proxy routing for a loopback endpoint %s", + async (baseUrl) => { + const apiKey = "test-api-key-123" + const modelName = "devstral2to16:latest" + + mockedAxios.get.mockResolvedValueOnce({ data: makeOllamaTagsPayload(modelName) }) + mockedAxios.post.mockResolvedValueOnce({ data: makeOllamaShowPayload() }) + + const result = await getOllamaModels(baseUrl, apiKey) + + // A cleartext loopback request carrying the bearer token must not + // be routed through any HTTP proxy, otherwise the proxy would see + // the credential in cleartext (CWE-319). Both requests bypass the + // proxy. The IPv6 loopback [::1] is treated the same as the IPv4 + // loopback: the credential is attached and proxy routing is + // disabled. + const expectedHeaders = { Authorization: `Bearer ${apiKey}` } + expect(mockedAxios.get).toHaveBeenCalledTimes(1) + expect(mockedAxios.get).toHaveBeenCalledWith(`${baseUrl}/api/tags`, { + headers: expectedHeaders, + proxy: false, + }) + + expect(mockedAxios.post).toHaveBeenCalledTimes(1) + expect(mockedAxios.post).toHaveBeenCalledWith( + `${baseUrl}/api/show`, + { model: modelName }, + { + headers: expectedHeaders, + proxy: false, + }, + ) + + expect(typeof result).toBe("object") + expect(Object.keys(result).length).toBe(1) + expect(result[modelName]).toBeDefined() + }, + ) + it("should fetch model list from /api/tags and include models with tools capability", async () => { const baseUrl = "http://localhost:11434" const modelName = "devstral2to16:latest" @@ -336,61 +477,25 @@ describe("Ollama Fetcher", () => { const apiKey = "test-api-key-123" const modelName = "test-model:latest" - const mockApiTagsResponse = { - models: [ - { - name: modelName, - model: modelName, - modified_at: "2025-06-03T09:23:22.610222878-04:00", - size: 14333928010, - digest: "6a5f0c01d2c96c687d79e32fdd25b87087feb376bf9838f854d10be8cf3c10a5", - details: { - family: "llama", - families: ["llama"], - format: "gguf", - parameter_size: "23.6B", - parent_model: "", - quantization_level: "Q4_K_M", - }, - }, - ], - } - const mockApiShowResponse = { - license: "Mock License", - modelfile: "FROM /path/to/blob\nTEMPLATE {{ .Prompt }}", - parameters: "num_ctx 4096\nstop_token ", - template: "{{ .System }}USER: {{ .Prompt }}ASSISTANT:", - modified_at: "2025-06-03T09:23:22.610222878-04:00", - details: { - parent_model: "", - format: "gguf", - family: "llama", - families: ["llama"], - parameter_size: "23.6B", - quantization_level: "Q4_K_M", - }, - model_info: { - "ollama.context_length": 4096, - "some.other.info": "value", - }, - capabilities: ["completion", "tools"], // Has tools capability - } - - mockedAxios.get.mockResolvedValueOnce({ data: mockApiTagsResponse }) - mockedAxios.post.mockResolvedValueOnce({ data: mockApiShowResponse }) + mockedAxios.get.mockResolvedValueOnce({ data: makeOllamaTagsPayload(modelName) }) + mockedAxios.post.mockResolvedValueOnce({ data: makeOllamaShowPayload() }) const result = await getOllamaModels(baseUrl, apiKey) const expectedHeaders = { Authorization: `Bearer ${apiKey}` } - + // Cleartext loopback + bearer token: the fetcher also disables proxy + // routing on both requests (CWE-319). expect(mockedAxios.get).toHaveBeenCalledTimes(1) - expect(mockedAxios.get).toHaveBeenCalledWith(`${baseUrl}/api/tags`, { headers: expectedHeaders }) + expect(mockedAxios.get).toHaveBeenCalledWith(`${baseUrl}/api/tags`, { + headers: expectedHeaders, + proxy: false, + }) expect(mockedAxios.post).toHaveBeenCalledTimes(1) expect(mockedAxios.post).toHaveBeenCalledWith( `${baseUrl}/api/show`, { model: modelName }, - { headers: expectedHeaders }, + { headers: expectedHeaders, proxy: false }, ) expect(typeof result).toBe("object") diff --git a/src/api/providers/fetchers/ollama.ts b/src/api/providers/fetchers/ollama.ts index 9f88d75327..a4bd42b28c 100644 --- a/src/api/providers/fetchers/ollama.ts +++ b/src/api/providers/fetchers/ollama.ts @@ -62,6 +62,48 @@ export const parseOllamaModel = (rawModel: OllamaModelInfoResponse): ModelInfo | return modelInfo } +/** + * Determines whether an Ollama endpoint is safe to carry an API key. + * + * Ollama's default installation listens on loopback, where plaintext HTTP is + * the norm. API keys are secrets, though, and must not be sent in cleartext to + * a remote host (CWE-319). Only HTTPS or a loopback host is considered safe + * enough to attach the Authorization header. + * + * Shared by the axios fetcher (getOllamaModels) and the native provider + * (NativeOllamaHandler), so this strict check gates credentials on both paths. + */ +export function isSecureOllamaEndpoint(baseUrl: string): boolean { + if (!URL.canParse(baseUrl)) { + return false + } + const url = new URL(baseUrl) + if (url.protocol === "https:") { + return true + } + const host = url.hostname + return host === "localhost" || host === "::1" || host === "[::1]" || isIpv4LoopbackHost(host) +} + +/** + * Strict 127.0.0.0/8 IPv4-loopback literal check. `host` comes from + * `URL.hostname`, which is WHATWG-normalized (IPv4 literals carry no brackets + * and no leading zeros), so the anchored octet check is exact. A loose + * /^127\./ prefix would also match DNS names like `127.example.com`, letting + * the key leak over cleartext HTTP (CWE-319). + */ +function isIpv4LoopbackHost(host: string): boolean { + const octets = host.split(".") + if (octets.length !== 4) { + return false + } + if (!octets.every((octet) => /^\d{1,3}$/.test(octet))) { + return false + } + const values = octets.map((octet) => Number(octet)) + return values[0] === 127 && values.every((value) => value >= 0 && value <= 255) +} + export async function getOllamaModels( baseUrl = "http://localhost:11434", apiKey?: string, @@ -76,13 +118,24 @@ export async function getOllamaModels( return models } - // Prepare headers with optional API key + // Prepare headers with optional API key. The credential is only attached + // when the endpoint is HTTPS or loopback; sending it over plaintext HTTP + // to a remote host would leak it (CWE-319). + const credentialGated = Boolean(apiKey && isSecureOllamaEndpoint(baseUrl)) const headers: Record = {} - if (apiKey) { + if (credentialGated) { headers["Authorization"] = `Bearer ${apiKey}` } - const response = await axios.get(`${baseUrl}/api/tags`, { headers }) + // A loopback HTTP endpoint carrying the key must bypass any HTTP proxy, + // otherwise the proxy would see the bearer token in cleartext (CWE-319). + // HTTPS endpoints keep normal proxy behavior (traffic stays encrypted). + // Parsing is safe here: baseUrl is normalized ("" → default) above and + // the !URL.canParse(baseUrl) guard returned early. + const cleartextLoopback = credentialGated && new URL(baseUrl).protocol === "http:" + const proxyConfig: { proxy?: false } = cleartextLoopback ? { proxy: false } : {} + + const response = await axios.get(`${baseUrl}/api/tags`, { headers, ...proxyConfig }) const parsedResponse = OllamaModelsResponseSchema.safeParse(response.data) const modelInfoPromises = [] @@ -95,7 +148,7 @@ export async function getOllamaModels( { model: ollamaModel.model, }, - { headers }, + { headers, ...proxyConfig }, ) .then((ollamaModelInfo) => { const modelInfo = parseOllamaModel(ollamaModelInfo.data) diff --git a/src/api/providers/native-ollama.ts b/src/api/providers/native-ollama.ts index eb380d1eb2..b9a7993dd5 100644 --- a/src/api/providers/native-ollama.ts +++ b/src/api/providers/native-ollama.ts @@ -5,7 +5,8 @@ import { ModelInfo, openAiModelInfoSaneDefaults, DEEP_SEEK_DEFAULT_TEMPERATURE } import { ApiStream } from "../transform/stream" import { BaseProvider } from "./base-provider" import type { ApiHandlerOptions } from "../../shared/api" -import { getOllamaModels } from "./fetchers/ollama" +import { getOllamaModels, isSecureOllamaEndpoint } from "./fetchers/ollama" +import { mergeAbortSignalAndTimeout } from "./utils/abort-signal" import { TagMatcher } from "../../utils/tag-matcher" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" @@ -23,6 +24,48 @@ type ReasoningContentBlock = { type: "reasoning"; text: string } type ThinkingContentBlock = { type: "thinking"; thinking: string } type AssistantContentBlock = Anthropic.ContentBlock | ReasoningContentBlock | ThinkingContentBlock +/** + * Creates the canonical abort error shared by every cancellation path so + * callers can identify cancellations by `name === "AbortError"`. + */ +function createAbortError(): Error { + const abortError = new Error("This operation was aborted") + abortError.name = "AbortError" + return abortError +} + +/** + * Races `pending` against `signal`: if the signal aborts before the promise + * settles, the returned promise rejects with an AbortError instead of + * resolving with a stale result. The Ollama model-discovery fetch accepts no + * signal, so the underlying request is left to settle in the background + * rather than being cancelled. + */ +function raceWithAbortSignal(pending: Promise, signal: AbortSignal | undefined): Promise { + if (!signal) { + return pending + } + if (signal.aborted) { + return Promise.reject(createAbortError()) + } + return new Promise((resolve, reject) => { + const onAbort = () => { + reject(createAbortError()) + } + signal.addEventListener("abort", onAbort, { once: true }) + pending.then( + (value) => { + signal.removeEventListener("abort", onAbort) + resolve(value) + }, + (error) => { + signal.removeEventListener("abort", onAbort) + reject(error) + }, + ) + }) +} + function convertToOllamaMessages(anthropicMessages: Anthropic.Messages.MessageParam[]): Message[] { const ollamaMessages: Message[] = [] // Track tool use IDs to tool names so tool results can be sent with @@ -225,7 +268,6 @@ function convertToOllamaMessages(anthropicMessages: Anthropic.Messages.MessagePa export class NativeOllamaHandler extends BaseProvider implements SingleCompletionHandler { protected options: ApiHandlerOptions - private client: Ollama | undefined protected models: Record = {} constructor(options: ApiHandlerOptions) { @@ -233,27 +275,43 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio this.options = options } - private ensureClient(): Ollama { - if (!this.client) { - try { - const clientOptions: OllamaOptions = { - host: this.options.ollamaBaseUrl || "http://localhost:11434", - // Note: The ollama npm package handles timeouts internally - } + /** + * Creates a new Ollama client instance with the configured options. + * Uses constructor `headers` option for API key instead of mutating config. + * The per-request transport makes both the in-flight POST and the response + * body cancellable; see the inline rationale. + */ + private _createOllamaClient(requestSignal?: AbortSignal): Ollama { + const clientOptions: OllamaOptions = { + host: this.options.ollamaBaseUrl || "http://localhost:11434", + // Note: The ollama npm package handles timeouts internally + } - // Add API key if provided (for Ollama cloud or authenticated instances) - if (this.options.ollamaApiKey) { - clientOptions.headers = { - Authorization: `Bearer ${this.options.ollamaApiKey}`, - } - } + // Add API key if provided (for Ollama cloud or authenticated instances). + // Use constructor `headers` option instead of mutating (request as any).config. + // The credential is only attached when the endpoint is HTTPS or loopback; + // sending it over plaintext HTTP to a remote host would leak it (CWE-319). + if (this.options.ollamaApiKey && isSecureOllamaEndpoint(clientOptions.host)) { + clientOptions.headers = { + Authorization: `Bearer ${this.options.ollamaApiKey}`, + } + } - this.client = new Ollama(clientOptions) - } catch (error: any) { - throw new Error(`Error creating Ollama client: ${error.message}`) + // The Ollama SDK (0.6.x) only tracks streaming iterators after the POST + // resolves and passes no signal for stream:false requests, so client + // .abort() alone cannot cancel an in-flight POST. Inject a per-request + // transport that merges our cancellation signal into every fetch while + // preserving the SDK's own stream signal when present (AbortSignal.any + // honors either side). + if (requestSignal) { + const baseFetch = globalThis.fetch + clientOptions.fetch = (url, init) => { + const signal = init?.signal ? AbortSignal.any([init.signal, requestSignal]) : requestSignal + return baseFetch(url, { ...init, signal }) } } - return this.client + + return new Ollama(clientOptions) } /** @@ -398,25 +456,54 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio messages: Anthropic.Messages.MessageParam[], metadata?: ApiHandlerCreateMessageMetadata, ): ApiStream { - const client = this.ensureClient() - const { id: modelId } = await this.fetchModel() - const useR1Format = modelId.toLowerCase().includes("deepseek-r1") - - const ollamaMessages: Message[] = [ - { role: "system", content: systemPrompt }, - ...convertToOllamaMessages(messages), - ] - - const matcher = new TagMatcher( - ["think", "thought"], - (chunk) => - ({ - type: chunk.matched ? "reasoning" : "text", - text: chunk.data, - }) as const, - ) + // Per-request client (SDK doesn't support per-call signal). The + // per-request controller drives the abortable transport injected into + // the client, so an in-flight POST is cancellable as well. + const requestController = new AbortController() + const client = this._createOllamaClient(requestController.signal) + + // Bridge the external abort signal into the per-request client and the + // per-request transport. The Ollama SDK exposes cancellation only + // through the client-level abort() and the stream's own signal, so the + // external signal drives both of this same request instead of a + // separate internal controller. + const externalAbortSignal = metadata?.abortSignal + if (externalAbortSignal?.aborted) { + requestController.abort() + client.abort() + throw createAbortError() + } + let onExternalAbort: (() => void) | undefined try { + if (externalAbortSignal) { + onExternalAbort = () => { + requestController.abort() + client.abort() + } + externalAbortSignal.addEventListener("abort", onExternalAbort, { once: true }) + } + + // Race model discovery against the external signal: the SDK's model + // fetch accepts no signal, so a request aborted during discovery must + // reject instead of proceeding to chat() with a stale model. + const { id: modelId } = await raceWithAbortSignal(this.fetchModel(), externalAbortSignal) + const useR1Format = modelId.toLowerCase().includes("deepseek-r1") + + const ollamaMessages: Message[] = [ + { role: "system", content: systemPrompt }, + ...convertToOllamaMessages(messages), + ] + + const matcher = new TagMatcher( + ["think", "thought"], + (chunk) => + ({ + type: chunk.matched ? "reasoning" : "text", + text: chunk.data, + }) as const, + ) + // Build the shared chat options and conditional think parameter. // Conditionally enabling Ollama's native think parameter lets // reasoning models (qwen3, deepseek-r1, etc.) emit thinking via @@ -437,6 +524,14 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio ...(thinkParam !== undefined ? { think: thinkParam } : {}), }) + // The POST can resolve in the same instant the abort lands (response + // headers arrived first). In that case the iterator must be disposed + // instead of being consumed after cancellation. + if (externalAbortSignal?.aborted) { + stream.abort() + throw createAbortError() + } + let totalInputTokens = 0 let totalOutputTokens = 0 // Track tool calls across chunks (Ollama may send complete tool_calls in final chunk) @@ -514,10 +609,22 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio } } } catch (streamError: any) { + // A body read aborted by the per-request transport surfaces as an + // AbortError DOMException; keep its name unmodified (same contract + // as the outer catch) so callers can identify cancellations. + if (streamError instanceof Error && streamError.name === "AbortError") { + throw streamError + } console.error("Error processing Ollama stream:", streamError) throw new Error(`Ollama stream processing error: ${streamError.message || "Unknown error"}`) } } catch (error: any) { + // Let AbortError surface unmodified so callers can identify cancellations + // by `name === "AbortError"`; every other error keeps the historical wrap. + if (error instanceof Error && error.name === "AbortError") { + throw error + } + // Enhance error reporting const statusCode = error.status || error.statusCode const errorMessage = error.message || "Unknown error" @@ -534,6 +641,18 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio console.error(`Ollama API error (${statusCode || "unknown"}): ${errorMessage}`) throw error + } finally { + // Detach the external-signal listener so it cannot outlive this request, + // including when model discovery rejects before the chat request starts. + if (onExternalAbort) { + externalAbortSignal?.removeEventListener("abort", onExternalAbort) + } + // A consumer that stops iterating early (break/return) would otherwise + // leave the in-flight response body open, and when metadata.abortSignal + // is undefined no abort bridge exists at all. The per-request controller + // drives the injected fetch transport, so aborting it here releases the + // body; after normal completion the abort is a no-op. + requestController.abort() } } @@ -551,9 +670,67 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio } async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { + // Ollama native client cancellation calls client.abort(), so any request with + // cancellation behavior must use a dedicated client instance to avoid aborting + // unrelated requests sharing a provider-level client. The Ollama SDK has no + // per-call signal, so every request uses a per-request client. The + // per-request controller drives the abortable transport injected into the + // client, so the in-flight POST itself is cancellable. + const requestController = new AbortController() + const client = this._createOllamaClient(requestController.signal) + let timeoutId: ReturnType | undefined + let onAbort: (() => void) | undefined + // Set when the timeoutMs timer fires so a model-discovery fetch still in + // flight is rejected instead of proceeding to chat() past the deadline. + let timedOut = false + try { - const client = this.ensureClient() - const { id: modelId } = await this.fetchModel() + // Handle timeoutMs if provided (client already exists above, so the + // timer can abort both the per-request transport and the client) + if (options?.timeoutMs !== undefined && options.timeoutMs > 0) { + timeoutId = setTimeout(() => { + timedOut = true + requestController.abort() + client.abort() + }, options.timeoutMs) + } + + // Propagate abortSignal into the per-request client via client.abort() + // and the per-request transport. Set this up before fetchModel() so an + // already-aborted signal (or one that aborts during the model-list + // fetch) is honored before the request proceeds. + if (options?.abortSignal) { + if (options.abortSignal.aborted) { + requestController.abort() + client.abort() + throw createAbortError() + } else { + onAbort = () => { + requestController.abort() + client.abort() + if (timeoutId !== undefined) { + clearTimeout(timeoutId) + } + } + options.abortSignal.addEventListener("abort", onAbort, { once: true }) + } + } + + // Model discovery goes through Axios (getOllamaModels) and never touches + // the Ollama client, so client.abort() alone cannot reject it. Race + // discovery against the per-request signal (bridged from the external + // abort) and the per-request timeout so a stalled discovery rejects + // instead of hanging past the deadline. + const discoverySignal = mergeAbortSignalAndTimeout(requestController.signal, options?.timeoutMs) + const { id: modelId } = await raceWithAbortSignal(this.fetchModel(), discoverySignal) + + // Re-check after the model-list fetch: if the request aborted or timed + // out during fetchModel(), the listener/timer above already aborted the + // client, and the request must reject instead of proceeding to chat(). + if (timedOut || options?.abortSignal?.aborted) { + throw createAbortError() + } + const useR1Format = modelId.toLowerCase().includes("deepseek-r1") // Reuse the shared request-option builder so single-shot @@ -571,10 +748,19 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio return response.message?.content || "" } catch (error) { - if (error instanceof Error) { + // Let AbortError surface unmodified so callers can identify cancellations + // by `name === "AbortError"`; every other error keeps the historical wrap. + if (error instanceof Error && error.name !== "AbortError") { throw new Error(`Ollama completion error: ${error.message}`) } throw error + } finally { + if (timeoutId) { + clearTimeout(timeoutId) + } + if (onAbort && options?.abortSignal) { + options.abortSignal.removeEventListener("abort", onAbort) + } } } } diff --git a/src/eslint-suppressions.json b/src/eslint-suppressions.json index 0706dbe6fb..822ac24f08 100644 --- a/src/eslint-suppressions.json +++ b/src/eslint-suppressions.json @@ -386,7 +386,7 @@ }, "api/providers/native-ollama.ts": { "@typescript-eslint/no-explicit-any": { - "count": 3 + "count": 2 } }, "api/providers/openai-codex.ts": {