lam-agent-vicinae/src/lib/ollama.ts

import { CLOUD_HOSTNAMES, DEFAULT_HOST } from "./constants";
import { formatBytes } from "./utils";
import type { ApiMessage, ModelInfo, ToolCall, ToolDefinition } from "./types";

/**
 * Thin client for the Ollama HTTP API.
 *
 * Everything here is plain `fetch`, deliberately: the extension runs in Node
 * inside Vicinae's extension runtime, so there is no reason to pull in the
 * `ollama` npm package and its own transport for four endpoints.
 */

/** An error with enough context for the UI to render something actionable rather than a raw status code. */
export class OllamaError extends Error {
  readonly status?: number;
  readonly hint?: string;

  constructor(message: string, options: { status?: number; hint?: string } = {}) {
    super(message);
    this.name = "OllamaError";
    this.status = options.status;
    this.hint = options.hint;
  }
}

/** Accepts `https://ollama.com/`, `http://localhost:11434`, or an empty preference, and yields a base URL with no trailing slash. */
export function normalizeHost(host: string | undefined): string {
  const trimmed = (host ?? "").trim();
  return (trimmed || DEFAULT_HOST).replace(/\/+$/, "");
}

function authHeaders(apiKey: string): Record<string, string> {
  const key = apiKey.trim();
  return key ? { Authorization: `Bearer ${key}` } : {};
}

/**
 * True when the host is a server whose models can be pulled and deleted.
 *
 * Decided from the hostname rather than the scheme or the port: a local server
 * reached over https is still local, and Ollama Cloud is the one host that
 * cannot be managed this way. A host that will not parse is treated as local,
 * because every request against it fails on its own terms anyway, and the caller
 * hides nothing it might have needed.
 */
export function isSelfHosted(host: string | undefined): boolean {
  try {
    return !CLOUD_HOSTNAMES.has(new URL(normalizeHost(host)).hostname);
  } catch {
    return true;
  }
}

async function safeText(response: Response): Promise<string> {
  try {
    return await response.text();
  } catch {
    return "";
  }
}

function describeHttpError(status: number, bodyText: string, host: string): OllamaError {
  const detail = bodyText.trim().slice(0, 300);

  switch (status) {
    case 401:
    case 403:
      return new OllamaError("Ollama rejected the API key.", {
        status,
        hint: "Check the Ollama API Key preference. Create a new key at https://ollama.com/settings/keys.",
      });
    case 404:
      return new OllamaError(`Ollama returned 404 from ${host}.`, {
        status,
        hint: "The model name may be wrong, or this host may not expose that endpoint.",
      });
    case 429:
      return new OllamaError("Ollama is rate limiting this key.", {
        status,
        hint: "Wait a moment and retry, or switch to a lighter model.",
      });
    default:
      return new OllamaError(`Ollama request failed with HTTP ${status}.`, { status, hint: detail });
  }
}

/**
 * Splits newline-delimited JSON into parsed objects.
 *
 * Exported for tests. Ollama streams NDJSON, and a network chunk can land
 * anywhere: mid-line, or between the bytes of a multi-byte UTF-8 character.
 * Both are handled here by buffering until a newline arrives and letting
 * TextDecoder hold incomplete sequences back.
 */
export async function* parseNdjsonStream<T>(stream: ReadableStream<Uint8Array>): AsyncGenerator<T> {
  const reader = stream.getReader();
  const decoder = new TextDecoder();
  let buffer = "";

  try {
    for (;;) {
      const { done, value } = await reader.read();
      if (done) break;

      buffer += decoder.decode(value, { stream: true });

      let newlineIndex = buffer.indexOf("\n");
      while (newlineIndex !== -1) {
        const line = buffer.slice(0, newlineIndex).trim();
        buffer = buffer.slice(newlineIndex + 1);
        if (line) yield JSON.parse(line) as T;
        newlineIndex = buffer.indexOf("\n");
      }
    }

    // Flush whatever the decoder is still holding, then the final unterminated line.
    buffer += decoder.decode();
    const tail = buffer.trim();
    if (tail) yield JSON.parse(tail) as T;
  } finally {
    try {
      reader.releaseLock();
    } catch {
      // The stream was already torn down, typically because the request was aborted.
    }
  }
}

interface OllamaChatEvent {
  message?: {
    role?: string;
    content?: string;
    thinking?: string;
    tool_calls?: ToolCall[];
  };
  done?: boolean;
  done_reason?: string;
  error?: string;
}

export interface ChatStreamChunk {
  content?: string;
  thinking?: string;
  toolCalls?: ToolCall[];
  done: boolean;
  doneReason?: string;
}

export interface StreamChatOptions {
  host: string;
  apiKey: string;
  model: string;
  messages: ApiMessage[];
  tools?: ToolDefinition[];
  /** Omit, or pass "off", to leave reasoning disabled. */
  think?: string | boolean;
  signal?: AbortSignal;
}

/**
 * Streams a chat completion.
 *
 * Because this is an async generator, the request is not issued until the
 * caller starts iterating, which lets the UI attach its loading state first.
 */
export async function* streamChat(options: StreamChatOptions): AsyncGenerator<ChatStreamChunk> {
  const host = normalizeHost(options.host);

  const body: Record<string, unknown> = {
    model: options.model,
    messages: options.messages,
    stream: true,
  };
  if (options.tools && options.tools.length > 0) body.tools = options.tools;
  if (options.think && options.think !== "off") body.think = options.think;

  const response = await fetch(`${host}/api/chat`, {
    method: "POST",
    headers: { "Content-Type": "application/json", ...authHeaders(options.apiKey) },
    body: JSON.stringify(body),
    signal: options.signal,
  });

  if (!response.ok) {
    throw describeHttpError(response.status, await safeText(response), host);
  }
  if (!response.body) {
    throw new OllamaError("Ollama returned an empty response body.");
  }

  for await (const event of parseNdjsonStream<OllamaChatEvent>(response.body)) {
    if (event.error) {
      throw new OllamaError(event.error);
    }

    yield {
      // Ollama sends empty strings for absent fields; normalise to undefined so
      // callers can use a simple truthiness check.
      content: event.message?.content || undefined,
      thinking: event.message?.thinking || undefined,
      toolCalls: event.message?.tool_calls,
      done: Boolean(event.done),
      doneReason: event.done_reason,
    };
  }
}

export interface ListModelsOptions {
  host: string;
  apiKey: string;
  signal?: AbortSignal;
}

/**
 * Lists the models this host can reach.
 *
 * On Ollama Cloud this reflects the account's entitlements. On a local server it
 * reflects what is installed, which mixes real weights with `:cloud` entries the
 * daemon proxies to ollama.com, so the two are told apart by `remoteHost` rather
 * than by anything in the name.
 */
export async function listModels(options: ListModelsOptions): Promise<ModelInfo[]> {
  const host = normalizeHost(options.host);

  const response = await fetch(`${host}/api/tags`, {
    headers: authHeaders(options.apiKey),
    signal: options.signal,
  });

  if (!response.ok) {
    throw describeHttpError(response.status, await safeText(response), host);
  }

  const payload = (await response.json()) as {
    models?: Array<{
      name?: string;
      model?: string;
      size?: number;
      remote_host?: string;
      capabilities?: string[];
      details?: { parameter_size?: string };
    }>;
  };

  return (payload.models ?? [])
    .map((model) => ({
      name: model.name ?? model.model ?? "",
      parameterSize: model.details?.parameter_size,
      size: typeof model.size === "number" ? model.size : undefined,
      // Every one of these is optional because older servers simply do not send
      // it, and a field that is missing must not be guessed at.
      capabilities: Array.isArray(model.capabilities) ? model.capabilities : undefined,
      remoteHost: typeof model.remote_host === "string" && model.remote_host ? model.remote_host : undefined,
    }))
    .filter((model) => model.name.length > 0)
    .sort((a, b) => a.name.localeCompare(b.name));
}

export interface DeleteModelOptions {
  host: string;
  apiKey: string;
  model: string;
  signal?: AbortSignal;
}

/**
 * A model the server has in memory right now, as reported by `GET /api/ps`.
 *
 * Distinct from {@link ModelInfo}, which is what is on disk. A model can be
 * installed without being loaded, and loaded models are the ones holding VRAM.
 */
export interface RunningModel {
  name: string;
  parameterSize?: string;
  /**
   * Bytes held in memory. Equal to the on-disk size when the model is fully
   * offloaded to the GPU, and less when part of it is running on the CPU.
   */
  sizeVram?: number;
  /** When the server will drop it on its own, as a Unix timestamp. */
  expiresAt?: number;
}

export interface ListRunningModelsOptions {
  host: string;
  apiKey: string;
  signal?: AbortSignal;
}

/**
 * The models holding memory on this server.
 *
 * Resolves to an empty list rather than throwing. Nothing here is worth failing
 * over: a server that does not serve `/api/ps` (older builds answer 404) simply
 * has nothing known to be loaded, and taking the whole model list down because
 * the *running* list is unavailable would be a bad trade.
 */
export async function listRunningModels(options: ListRunningModelsOptions): Promise<RunningModel[]> {
  const host = normalizeHost(options.host);

  try {
    const response = await fetch(`${host}/api/ps`, {
      headers: authHeaders(options.apiKey),
      signal: options.signal,
    });

    if (!response.ok) return [];

    const payload = (await response.json()) as {
      models?: Array<{
        name?: string;
        model?: string;
        size_vram?: number;
        expires_at?: string;
        details?: { parameter_size?: string };
      }>;
    };

    return (payload.models ?? [])
      .map((model) => ({
        name: model.name ?? model.model ?? "",
        parameterSize: model.details?.parameter_size,
        sizeVram: typeof model.size_vram === "number" ? model.size_vram : undefined,
        expiresAt: model.expires_at ? Date.parse(model.expires_at) : undefined,
      }))
      .filter((model) => model.name.length > 0);
  } catch {
    return [];
  }
}

export interface UnloadModelOptions {
  host: string;
  apiKey: string;
  model: string;
  signal?: AbortSignal;
}

/**
 * Drops a model out of memory, freeing the VRAM it was holding.
 *
 * `POST /api/generate` with `keep_alive: 0` and no prompt is how Ollama says
 * "release this now"; its own `ollama stop` does the same thing. Probed against
 * 0.35.1: the reply is `{"done":true,"done_reason":"unload"}` and `/api/ps`
 * comes back empty. The `DELETE` form of the same route answers 405, so there is
 * no delete-shaped way to do this.
 *
 * Nothing is deleted, so a mistake costs a reload. That is why this needs no
 * confirmation while `deleteModel` does.
 */
export async function unloadModel(options: UnloadModelOptions): Promise<void> {
  const host = normalizeHost(options.host);

  const response = await fetch(`${host}/api/generate`, {
    method: "POST",
    headers: { "Content-Type": "application/json", ...authHeaders(options.apiKey) },
    // No `prompt`, so this never generates anything, and `stream: false` because
    // there is exactly one reply to read.
    body: JSON.stringify({ model: options.model, keep_alive: 0, stream: false }),
    signal: options.signal,
  });

  if (!response.ok) {
    throw describeHttpError(response.status, await safeText(response), host);
  }
}

/**
 * Removes a model from a server the user runs.
 *
 * `DELETE`, not `POST`: Ollama 0.35.1 answers the POST form of this route with
 * `405 method not allowed`, which is the kind of change that is invisible until a
 * user tries it.
 *
 * There is no undo and no trash, so the caller confirms first. For the same
 * reason a 404 is not swallowed into a success: "no model called that" is worth
 * saying out loud, because it usually means something else already removed it, or
 * that the host is not what the user thinks it is.
 */
export async function deleteModel(options: DeleteModelOptions): Promise<void> {
  const host = normalizeHost(options.host);

  const response = await fetch(`${host}/api/delete`, {
    method: "DELETE",
    headers: { "Content-Type": "application/json", ...authHeaders(options.apiKey) },
    body: JSON.stringify({ model: options.model }),
    signal: options.signal,
  });

  if (response.status === 404) {
    throw new OllamaError(`${host} has no model called ${options.model}.`, {
      status: 404,
      hint: "It may have been removed already, or by another machine sharing this server.",
    });
  }

  if (!response.ok) {
    throw describeHttpError(response.status, await safeText(response), host);
  }
}

/**
 * One line of pull progress, exactly as Ollama reports it.
 *
 * `total` and `completed` are only present once a download starts, so the caller
 * has to render a status with no numbers in it for the first few ticks.
 */
export interface PullProgress {
  status: string;
  total?: number;
  completed?: number;
}

export interface PullModelOptions {
  host: string;
  apiKey: string;
  model: string;
  signal?: AbortSignal;
}

/**
 * Pulls a model onto a server the user runs, yielding progress as it arrives.
 *
 * Streaming rather than a blocking request, because a download is minutes long
 * and a request that returns nothing until it is finished leaves the window with
 * no way to tell a slow pull from a hung one.
 *
 * A bad model name comes back as **HTTP 200 with an `error` field in the body**
 * (verified against 0.35.1), so the status line is not enough to decide whether
 * this worked: the stream has to be read and checked, exactly as the chat stream
 * already does.
 */
export async function* pullModel(options: PullModelOptions): AsyncGenerator<PullProgress> {
  const host = normalizeHost(options.host);

  const response = await fetch(`${host}/api/pull`, {
    method: "POST",
    headers: { "Content-Type": "application/json", ...authHeaders(options.apiKey) },
    body: JSON.stringify({ model: options.model, stream: true }),
    signal: options.signal,
  });

  if (!response.ok) {
    throw describeHttpError(response.status, await safeText(response), host);
  }
  if (!response.body) {
    throw new OllamaError("Ollama returned an empty response body.");
  }

  for await (const event of parseNdjsonStream<{
    status?: string;
    total?: number;
    completed?: number;
    error?: string;
  }>(response.body)) {
    if (event.error) throw new OllamaError(event.error);
    if (!event.status) continue;

    yield { status: event.status, total: event.total, completed: event.completed };
  }
}

/**
 * One line describing a pull, for the toast that stays up while it runs.
 *
 * Lives here rather than in the view because it is a reading of Ollama's own
 * payload rather than a layout decision, and because the arithmetic is the part
 * of a progress bar that goes wrong silently: a total of zero, or a `completed`
 * that arrives before the `total` it is measured against, would otherwise print
 * `NaN%` or `Infinity%` at the user.
 */
export function describePullProgress(progress: PullProgress): string {
  const { status, total, completed } = progress;

  if (!total || total <= 0 || completed === undefined || !Number.isFinite(completed)) {
    return status;
  }

  const percent = Math.min(100, Math.max(0, Math.round((completed / total) * 100)));

  return `${status} · ${percent}% · ${formatBytes(completed)} / ${formatBytes(total)}`;
}

export interface ModelCapabilitiesOptions {
  host: string;
  apiKey: string;
  model: string;
  signal?: AbortSignal;
}

/**
 * The features a model reports for itself, straight from `POST /api/show`.
 *
 * `GET /api/tags` is not enough here: it returns names, sizes and quantization
 * levels and says nothing about what the model can actually do. The capability
 * list is what turns "image input" from a guess about the model's name into a
 * fact, so it is worth the extra request.
 */
export async function modelCapabilities(options: ModelCapabilitiesOptions): Promise<string[]> {
  const host = normalizeHost(options.host);

  const response = await fetch(`${host}/api/show`, {
    method: "POST",
    headers: { "Content-Type": "application/json", ...authHeaders(options.apiKey) },
    body: JSON.stringify({ model: options.model }),
    signal: options.signal,
  });

  if (!response.ok) {
    throw describeHttpError(response.status, await safeText(response), host);
  }

  const payload = (await response.json()) as { capabilities?: string[] };
  return payload.capabilities ?? [];
}

/** Ollama reports image input as the `vision` capability. */
export function supportsVision(capabilities: string[]): boolean {
  return capabilities.includes("vision");
}

/**
 * Posts to a web tool, at the path this kind of host actually serves it from.
 *
 * The two kinds of host expose the same two operations at different paths, and
 * there is no path that works for both. Measured, not looked up:
 *
 * - `https://ollama.com/api/web_search` answers 401 without a key, so the route
 *   exists, and its `/api/experimental/` alias is a 404.
 * - A local 0.35.1 daemon answers 404 to that path, and serves
 *   `/api/experimental/web_search` with real results.
 *
 * This was a live bug: pointing the extension at a local server made every web
 * tool call a 404, and the model, handed a failure and no explanation, told the
 * user its web tools simply were not configured.
 *
 * The other path is tried when the first answers 404, because that is a server
 * built the other way round rather than a request being refused. A 404 body is
 * drained before the retry so the connection is not left hanging.
 */
async function postWebTool(options: {
  host: string;
  apiKey: string;
  kind: "search" | "fetch";
  body: Record<string, unknown>;
  signal?: AbortSignal;
}): Promise<Response> {
  const host = normalizeHost(options.host);

  const [preferred, fallback] = isSelfHosted(host)
    ? [`/api/experimental/web_${options.kind}`, `/api/web_${options.kind}`]
    : [`/api/web_${options.kind}`, `/api/experimental/web_${options.kind}`];

  const post = (path: string) =>
    fetch(`${host}${path}`, {
      method: "POST",
      headers: { "Content-Type": "application/json", ...authHeaders(options.apiKey) },
      body: JSON.stringify(options.body),
      signal: options.signal,
    });

  const response = await post(preferred);
  if (response.status !== 404) return response;

  await response.text().catch(() => "");
  return post(fallback);
}

export interface WebSearchResult {
  title: string;
  url: string;
  content: string;
}

export interface WebSearchOptions {
  host: string;
  apiKey: string;
  query: string;
  maxResults?: number;
  signal?: AbortSignal;
}

export async function webSearch(options: WebSearchOptions): Promise<WebSearchResult[]> {
  const body: Record<string, unknown> = { query: options.query };
  if (typeof options.maxResults === "number") body.max_results = options.maxResults;

  const response = await postWebTool({
    host: options.host,
    apiKey: options.apiKey,
    kind: "search",
    body,
    signal: options.signal,
  });

  if (!response.ok) {
    throw describeHttpError(response.status, await safeText(response), normalizeHost(options.host));
  }

  const payload = (await response.json()) as { results?: WebSearchResult[] };
  return payload.results ?? [];
}

export interface WebFetchResult {
  title: string;
  content: string;
  links: string[];
}

export interface WebFetchOptions {
  host: string;
  apiKey: string;
  url: string;
  signal?: AbortSignal;
}

export async function webFetch(options: WebFetchOptions): Promise<WebFetchResult> {
  const response = await postWebTool({
    host: options.host,
    apiKey: options.apiKey,
    kind: "fetch",
    body: { url: options.url },
    signal: options.signal,
  });

  if (!response.ok) {
    throw describeHttpError(response.status, await safeText(response), normalizeHost(options.host));
  }

  // `links` comes back as null from a local daemon rather than as an empty list,
  // which `??` covers. (Probed directly: a local `web_fetch` answered
  // `{"title":...,"content":...,"links":null}`.)
  const payload = (await response.json()) as Partial<WebFetchResult>;
  return {
    title: payload.title ?? options.url,
    content: payload.content ?? "",
    links: payload.links ?? [],
  };
}