diff --git a/package.json b/package.json index 8431467918..0a914fb148 100644 --- a/package.json +++ b/package.json @@ -43,7 +43,8 @@ "rimraf": "6.0.1", "tsx": "4.22.4", "turbo": "2.10.0", - "typescript": "5.9.3" + "typescript": "5.9.3", + "vitest": "4.1.9" }, "lint-staged": { "*.{js,jsx,ts,tsx,json,css,md}": [ diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 183a5e02d1..a8972fd426 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -74,6 +74,9 @@ importers: typescript: specifier: 5.9.3 version: 5.9.3 + vitest: + specifier: 4.1.9 + version: 4.1.9(@opentelemetry/api@1.9.0)(@types/node@22.20.1)(@vitest/coverage-v8@4.1.9)(@vitest/ui@4.1.9)(jsdom@26.1.0)(vite@8.1.0(@types/node@22.20.1)(esbuild@0.28.1)(jiti@2.7.0)(tsx@4.22.4)(yaml@2.9.0)) apps/cli: dependencies: @@ -5037,6 +5040,7 @@ packages: eslint@9.39.4: resolution: {integrity: sha512-XoMjdBOwe/esVgEvLmNsD3IRHkm7fbKIUGvrleloJXUZgDHig2IPWNniv+GwjyJXzuNqVjlr5+4yVUZjycJwfQ==} engines: {node: ^18.18.0 || ^20.9.0 || >=21.1.0} + deprecated: This version is no longer supported. Please see https://eslint.org/version-support for other options. hasBin: true peerDependencies: jiti: '*' diff --git a/src/api/providers/__tests__/native-ollama.spec.ts b/src/api/providers/__tests__/native-ollama.spec.ts index 8fcbf4a0a1..ffcba936c3 100644 --- a/src/api/providers/__tests__/native-ollama.spec.ts +++ b/src/api/providers/__tests__/native-ollama.spec.ts @@ -2,32 +2,66 @@ import { Anthropic } from "@anthropic-ai/sdk" -import { NativeOllamaHandler } from "../native-ollama" +import { NativeOllamaHandler, raceWithAbortSignal } 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) +// Bound await helper: mutation runs time out at 5000ms per mutant, so a test +// that would hang under a mutant must fail fast instead. Capture the native +// setTimeout at module load so a lingering setTimeout spy cannot swallow it. +const nativeSetTimeout = globalThis.setTimeout +const withDeadline = async (promise: Promise, ms = 500): Promise => + Promise.race([ + promise, + new Promise((_resolve, reject) => { + nativeSetTimeout(() => reject(new Error(`withDeadline: unsettled after ${ms}ms`)), ms) + }), + ]) + describe("NativeOllamaHandler", () => { let handler: NativeOllamaHandler @@ -85,6 +119,67 @@ describe("NativeOllamaHandler", () => { expect(results[2]).toEqual({ type: "usage", inputTokens: 10, outputTokens: 2 }) }) + it("should classify reasoning chunks by the think tag", async () => { + mockChat.mockImplementation(async function* () { + yield { message: { content: "secrethi" } } + }) + const stream = handler.createMessage("System", [{ role: "user" as const, content: "Test" }]) + const parts: { type: string; text: string }[] = [] + for await (const part of stream) { + parts.push(part as { type: string; text: string }) + } + const reasoning = parts + .filter((p) => p.type === "reasoning") + .map((p) => p.text) + .join("") + const text = parts + .filter((p) => p.type === "text") + .map((p) => p.text) + .join("") + expect(reasoning).toBe("secret") + expect(text).toBe("hi") + }) + + // The matcher only opens a top-level tag at the start of the stream + // (TagMatcher position 0), so the thought tag needs its own stream to + // exercise the second entry of the tag list. + it("should classify reasoning chunks by the thought tag when it opens the stream", async () => { + mockChat.mockImplementation(async function* () { + yield { message: { content: "plandone" } } + }) + const stream = handler.createMessage("System", [{ role: "user" as const, content: "Test" }]) + const parts: { type: string; text: string }[] = [] + for await (const part of stream) { + parts.push(part as { type: string; text: string }) + } + const reasoning = parts + .filter((p) => p.type === "reasoning") + .map((p) => p.text) + .join("") + const text = parts + .filter((p) => p.type === "text") + .map((p) => p.text) + .join("") + expect(reasoning).toBe("plan") + expect(text).toBe("done") + }) + + it("should send the system prompt as the first message of the chat request", async () => { + mockChat.mockImplementation(async function* () { + yield { message: { content: "ok" } } + }) + const stream = handler.createMessage("You are terse", [{ role: "user" as const, content: "Hi" }]) + const consume = (async () => { + for await (const _ of stream) { + // drain + } + })() + await withDeadline(consume, 500) + // documented narrow cast: mockChat is an untyped vi.fn() + const callArgs = mockChat.mock.calls[0][0] as { messages: { role: string; content: string }[] } + expect(callArgs.messages[0]).toEqual({ role: "system", content: "You are terse" }) + }) + it("should map tool_result array content to a concatenated string, flushing base64 images", async () => { mockChat.mockImplementation(async function* () { yield { message: { content: "ok" } } @@ -696,6 +791,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 +856,266 @@ 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 wrap non-AbortError completion failures", async () => { + mockChat.mockRejectedValue(new Error("boom")) + const settledError = expect(handler.completePrompt("Test prompt")).rejects.toThrow( + "Ollama completion error: boom", + ) + await withDeadline(settledError, 500) + }) + + it("should detach the abort listener when completePrompt settles normally", async () => { + let abortCount = 0 + OllamaMock.mockImplementation(function (options?: { host?: string }) { + return { + chat: mockChat, + abort: () => { + abortCount++ + }, + _host: options?.host ?? "http://localhost:11434", + } + }) + mockChat.mockResolvedValue({ message: { content: "Response" } }) + const controller = new AbortController() + await handler.completePrompt("Test prompt", { abortSignal: controller.signal }) + expect(abortCount).toBe(0) + // The listener is detached in the finally block, so a late abort must + // not reach the per-request client. + controller.abort() + expect(abortCount).toBe(0) + }) }) it("should send think parameter in completePrompt when reasoningEffort is set", async () => { @@ -902,6 +1262,28 @@ describe("NativeOllamaHandler", () => { } }).rejects.toThrow("something else") }) + it("should rethrow stream AbortError without logging it", async () => { + const consoleError = vi.spyOn(console, "error").mockImplementation(() => {}) + mockChat.mockImplementation(async function* () { + yield* [] + throw new DOMException("This operation was aborted", "AbortError") + }) + const stream = handler.createMessage("System", [{ role: "user" as const, content: "Test" }]) + const consume = (async () => { + for await (const _ of stream) { + // drain + } + })() + try { + const settledStream = expect(consume).rejects.toMatchObject({ name: "AbortError" }) + await withDeadline(settledStream, 500) + // Assert before mockRestore(): restoring the spy resets its call + // history, which would make the assertion below a no-op. + expect(consoleError).not.toHaveBeenCalled() + } finally { + consoleError.mockRestore() + } + }) }) describe("getModel", () => { @@ -1602,5 +1984,803 @@ 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?.() + const settledTimeoutRace = expect(promise).rejects.toMatchObject({ name: "AbortError" }) + await withDeadline(settledTimeoutRace, 500) + 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", + }, + }), + ) + }) + + it("should default the client host to localhost when no base URL is configured", async () => { + const bareHandler = new NativeOllamaHandler({ + apiModelId: "llama2", + ollamaModelId: "llama2", + }) + mockChat.mockResolvedValue({ message: { content: "Response" } }) + await bareHandler.completePrompt("Test prompt") + expect(OllamaMock.mock.calls[0][0]).toEqual(expect.objectContaining({ host: "http://localhost:11434" })) + }) + }) + + describe("raceWithAbortSignal", () => { + it("should resolve with the pending result when no signal is given", async () => { + const result = await raceWithAbortSignal(Promise.resolve("value"), undefined) + expect(result).toBe("value") + }) + + it("should reject with AbortError when the signal is already aborted", async () => { + const controller = new AbortController() + controller.abort() + const pending = new Promise(() => {}) + const settledPre = expect(raceWithAbortSignal(pending, controller.signal)).rejects.toMatchObject({ + name: "AbortError", + }) + await withDeadline(settledPre, 500) + }) + + it("should reject with AbortError when the signal aborts before the promise settles", async () => { + const controller = new AbortController() + const pending = new Promise(() => {}) + const raced = raceWithAbortSignal(pending, controller.signal) + controller.abort() + const settledAbort = expect(raced).rejects.toMatchObject({ name: "AbortError" }) + await withDeadline(settledAbort, 500) + }) + + it("should resolve when the promise settles before the signal aborts", async () => { + const controller = new AbortController() + const raced = raceWithAbortSignal(Promise.resolve("done"), controller.signal) + const result = await withDeadline(raced, 500) + expect(result).toBe("done") + controller.abort() // a late abort must not disturb the settled result + }) + }) + + 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. + const settledHeldPost = expect(heldPost).rejects.toMatchObject({ name: "AbortError" }) + await withDeadline(settledHeldPost, 500) + + // 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* () {}, + }) + const settledConsume = expect(consume).rejects.toMatchObject({ name: "AbortError" }) + await withDeadline(settledConsume, 500) + }) + + 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 withDeadline(released, 500) + }) + + 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. + const settledTimeoutPost = expect(heldPost).rejects.toMatchObject({ name: "AbortError" }) + await withDeadline(settledTimeoutPost, 500) + + // Mirror the real SDK: the in-flight POST rejects, and the + // completion surfaces the AbortError unmodified. + rejectChat?.(new DOMException("This operation was aborted", "AbortError")) + const settledTimeoutPromise = expect(promise).rejects.toMatchObject({ name: "AbortError" }) + await withDeadline(settledTimeoutPromise, 500) + }) + + 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() + const settledMergedPost = expect(post).rejects.toMatchObject({ name: "AbortError" }) + await withDeadline(settledMergedPost, 500) + }) + + it("should abort the injected transport when the SDK fetch call carries no init signal", async () => { + const fetchSpy = heldOpenFetch() + vi.stubGlobal("fetch", fetchSpy) + mockChat.mockImplementation(async function* () { + yield { message: { content: "ok" } } + }) + const stream = handler.createMessage("System", [{ role: "user" as const, content: "Test" }]) + for (let i = 0; i < 50 && mockChat.mock.calls.length === 0; i++) { + await Promise.resolve() + } + // Consume the stream to completion; the finally block aborts the + // per-request controller, which must release the held transport. + const consume = (async () => { + for await (const _ of stream) { + // drain + } + })() + await withDeadline(consume, 500) + const clientOptions = OllamaMock.mock.calls[0][0] as ConstructorOptions + const heldPost = clientOptions.fetch!("http://localhost:11434/api/chat") + const settledPost = expect(heldPost).rejects.toMatchObject({ name: "AbortError" }) + await withDeadline(settledPost, 500) + }) + + it("should abort the per-request transport when completePrompt's abortSignal is already aborted", async () => { + const fetchSpy = heldOpenFetch() + vi.stubGlobal("fetch", fetchSpy) + + // Count the per-request client's abort() via a local spy factory, + // mirroring the capturedInstanceAbort pattern (scoped to the + // completePrompt describe, so not reachable from here). + let instanceAbortCount = 0 + OllamaMock.mockImplementation(function (options?: { host?: string }) { + return { + chat: mockChat, + abort: () => { + instanceAbortCount++ + }, + _host: options?.host ?? "http://localhost:11434", + } + }) + + const controller = new AbortController() + controller.abort() + const promise = handler.completePrompt("Test prompt", { abortSignal: controller.signal }) + const settledPromise = expect(promise).rejects.toMatchObject({ name: "AbortError" }) + await withDeadline(settledPromise, 500) + expect(mockChat).not.toHaveBeenCalled() + expect(instanceAbortCount).toBe(1) + const clientOptions = OllamaMock.mock.calls[0][0] as ConstructorOptions + const heldPost = clientOptions.fetch!("http://localhost:11434/api/chat") + const settledPost = expect(heldPost).rejects.toMatchObject({ name: "AbortError" }) + await withDeadline(settledPost, 500) + }) + + it("should abort the per-request transport when completePrompt's abortSignal fires during discovery", async () => { + const fetchSpy = heldOpenFetch() + vi.stubGlobal("fetch", fetchSpy) + 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 }) + for (let i = 0; i < 50 && OllamaMock.mock.calls.length === 0; i++) { + await Promise.resolve() + } + controller.abort() + const settledPromise = expect(promise).rejects.toMatchObject({ name: "AbortError" }) + await withDeadline(settledPromise, 500) + resolveFetch?.() + expect(mockChat).not.toHaveBeenCalled() + const clientOptions = OllamaMock.mock.calls[0][0] as ConstructorOptions + const heldPost = clientOptions.fetch!("http://localhost:11434/api/chat") + const settledPost = expect(heldPost).rejects.toMatchObject({ name: "AbortError" }) + await withDeadline(settledPost, 500) + }) + + it("should abort the per-request client only once when the external abort fires before the timeout", async () => { + let abortCount = 0 + OllamaMock.mockImplementation(function (options?: { host?: string }) { + return { + chat: mockChat, + abort: () => { + abortCount++ + }, + _host: options?.host ?? "http://localhost:11434", + } + }) + let resolveFetch: (() => void) | undefined + + // Hold model discovery open so the external abort lands while the + // per-request client is still in flight. + mockGetOllamaModels.mockImplementation( + () => + new Promise((resolve) => { + resolveFetch = () => resolve({}) + }), + ) + + const controller = new AbortController() + const promise = handler.completePrompt("Test prompt", { timeoutMs: 100, abortSignal: controller.signal }) + for (let i = 0; i < 50 && OllamaMock.mock.calls.length === 0; i++) { + await Promise.resolve() + } + controller.abort() + const settledPromise = expect(promise).rejects.toMatchObject({ name: "AbortError" }) + await withDeadline(settledPromise, 500) + // The external-abort path must be the only abort: the cleared timer must + // not fire a second client.abort() afterwards. + expect(abortCount).toBe(1) + await new Promise((resolve) => setTimeout(resolve, 150)) + expect(abortCount).toBe(1) + resolveFetch?.() + }) + }) + + 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 }), + ) + + const settledPreAbort = expect(async () => { + for await (const _ of stream) { + // consume stream + } + }).rejects.toMatchObject({ name: "AbortError" }) + await withDeadline(settledPreAbort, 500) + + // The pre-abort path must reject before allocating the per-request client. + expect(OllamaMock).not.toHaveBeenCalled() + expect(mockChat).not.toHaveBeenCalled() + }) + + 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() + + const settledInFlight = expect(consume).rejects.toMatchObject({ name: "AbortError" }) + await withDeadline(settledInFlight, 500) + }) + + it("should detach the abort listener when the stream settles normally", async () => { + let abortCount = 0 + OllamaMock.mockImplementation(function (options?: { host?: string }) { + return { + chat: mockChat, + abort: () => { + abortCount++ + }, + _host: options?.host ?? "http://localhost:11434", + } + }) + mockChat.mockImplementation(async function* () { + yield { message: { content: "ok" } } + }) + 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) { + // drain + } + })() + await withDeadline(consume, 500) + // The listener must be detached in the finally block, so a late abort + // after settlement must not reach the per-request client. + controller.abort() + expect(abortCount).toBe(0) + }) + + 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() + const settledDiscovery = expect(consume).rejects.toMatchObject({ name: "AbortError" }) + await withDeadline(settledDiscovery, 500) + 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 + } + })() + + const settledDispose = expect(consume).rejects.toMatchObject({ name: "AbortError" }) + await withDeadline(settledDispose, 500) + 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" }]) + + const settledPassthrough = expect(async () => { + for await (const _ of stream) { + // consume stream + } + }).rejects.toMatchObject({ name: "AbortError" }) + await withDeadline(settledPassthrough, 500) + }) }) }) diff --git a/src/api/providers/fetchers/__tests__/ollama.test.ts b/src/api/providers/fetchers/__tests__/ollama.test.ts index 9c0b547e88..b2db91eff9 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, isIpv4LoopbackHost, isSecureOllamaEndpoint, parseOllamaModel } from "../ollama" import ollamaModelsData from "./fixtures/ollama-model-details.json" // Mock axios @@ -116,7 +116,175 @@ 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 + expect(isSecureOllamaEndpoint("http://127.256.0.1:11434")).toBe(false) + expect(isSecureOllamaEndpoint("http://one.two.three.four:11434")).toBe(false) + }) + + 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("isIpv4LoopbackHost", () => { + it("should accept only four-octet 127/8 literals", () => { + expect(isIpv4LoopbackHost("127.0.0.1")).toBe(true) + expect(isIpv4LoopbackHost("127.255.255.255")).toBe(true) + expect(isIpv4LoopbackHost("0.0.0.1")).toBe(false) + }) + + it("should reject hosts with the wrong octet count", () => { + expect(isIpv4LoopbackHost("127.0.0")).toBe(false) + expect(isIpv4LoopbackHost("127.0.0.0.1")).toBe(false) + }) + + it("should reject non-numeric octets even when they parse as numbers", () => { + // Number("0x4") === 4 and Number("1e2") === 100, so only the anchored + // octet regex rejects these inputs. + expect(isIpv4LoopbackHost("127.0.0.0x4")).toBe(false) + expect(isIpv4LoopbackHost("127.1e2.0.1")).toBe(false) + }) + + it("should reject out-of-range octets", () => { + expect(isIpv4LoopbackHost("127.999.0.1")).toBe(false) + expect(isIpv4LoopbackHost("0127.0.0.1")).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 +504,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..2b4888c31f 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]" || 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). + */ +export 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 <= 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..e6d7dc73c1 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,36 @@ 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. + */ +export function raceWithAbortSignal(pending: Promise, signal: AbortSignal | undefined): Promise { + if (!signal) { + return pending + } + if (signal.aborted) { + return Promise.reject(createAbortError()) + } + return new Promise((resolve, reject) => { + signal.addEventListener("abort", () => reject(createAbortError())) + pending.then(resolve, reject) + }) +} + 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 +256,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 +263,41 @@ 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 - } - - // Add API key if provided (for Ollama cloud or authenticated instances) - if (this.options.ollamaApiKey) { - clientOptions.headers = { - Authorization: `Bearer ${this.options.ollamaApiKey}`, - } - } + /** + * 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 + } - this.client = new Ollama(clientOptions) - } catch (error: any) { - throw new Error(`Error creating Ollama client: ${error.message}`) + // 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}`, } } - return this.client + + // 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. The client always + // receives a per-request transport that merges the request signal into + // the SDK's fetch calls while preserving the SDK's own stream signal when + // present (AbortSignal.any honors either side). + const baseFetch = globalThis.fetch + clientOptions.fetch = (url, init) => { + const signal = init?.signal ? AbortSignal.any([init.signal, requestSignal]) : requestSignal + return baseFetch(url, { ...init, signal }) + } + + return new Ollama(clientOptions) } /** @@ -398,25 +442,50 @@ 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, - ) + // The external abort signal drives both the per-request client and the + // per-request transport for this request, since the Ollama SDK exposes + // cancellation only through the client-level abort() and the fetch signal. + const externalAbortSignal = metadata?.abortSignal + if (externalAbortSignal?.aborted) { + // The request is already cancelled: reject before allocating the + // per-request client or transport so nothing is left to clean up. + throw createAbortError() + } + + const requestController = new AbortController() + const client = this._createOllamaClient(requestController.signal) + let onExternalAbort: (() => void) | undefined try { + if (externalAbortSignal) { + onExternalAbort = () => { + requestController.abort() + client.abort() + } + externalAbortSignal.addEventListener("abort", onExternalAbort) + } + + // 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) + // Stryker disable next-line MethodExpression,StringLiteral: DEEP_SEEK_DEFAULT_TEMPERATURE (0.0) is numerically equal to the 0 fallback, so the R1 check cannot change the emitted temperature + 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 +506,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 +591,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 +623,19 @@ 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. + // Stryker disable next-line LogicalOperator: onExternalAbort is only assigned when externalAbortSignal is present, so || is equivalent + if (onExternalAbort && externalAbortSignal) { + 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 +653,57 @@ 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 + const abortSignal = options?.abortSignal + 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(() => { + 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 (abortSignal) { + if (abortSignal.aborted) { + requestController.abort() + client.abort() + throw createAbortError() + } else { + onAbort = () => { + requestController.abort() + client.abort() + // The per-request timeout is cleared by the finally block when this prompt + // settles; an external abort always settles it promptly (the discovery race + // or the injected transport rejects), so no timer clear is needed here. + } + 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) + const useR1Format = modelId.toLowerCase().includes("deepseek-r1") // Reuse the shared request-option builder so single-shot @@ -571,10 +721,18 @@ 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 { + clearTimeout(timeoutId) + // Stryker disable next-line LogicalOperator: onAbort is only assigned when abortSignal is present, so || is equivalent + if (onAbort && abortSignal) { + 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": {