@prometheus-ai/agent 0.5.4 → 0.5.8

package/src/utils/image-loading.ts

@@ -1,12 +1,32 @@
 import * as fs from "node:fs/promises";
-import type { ImageContent } from "@prometheus-ai/ai";
+import type { ImageContent, Model } from "@prometheus-ai/ai";
 import { formatBytes, readImageMetadata, SUPPORTED_IMAGE_MIME_TYPES } from "@prometheus-ai/utils";
 import { resolveReadPath } from "../tools/path-utils";
-import { formatDimensionNote, resizeImage } from "./image-resize";
+import { formatDimensionNote, type ImageResizeOptions, resizeImage } from "./image-resize";
 
 export const MAX_IMAGE_INPUT_BYTES = 20 * 1024 * 1024;
 export const SUPPORTED_INPUT_IMAGE_MIME_TYPES = SUPPORTED_IMAGE_MIME_TYPES;
 
+/**
+ * Ollama and its local-backend family decode image input through llama.cpp /
+ * `stb_image`, which is compiled without WebP support, so a WebP upload fails
+ * with an opaque HTTP 400. Detect those models so the resize pipeline encodes
+ * to PNG/JPEG instead — the automatic equivalent of `PROMETHEUS_NO_WEBP=1`.
+ */
+export function modelLacksWebpSupport(model: Pick<Model, "provider" | "api"> | undefined): boolean {
+	if (!model) return false;
+	return model.provider === "ollama" || model.provider === "ollama-cloud" || model.api === "ollama-chat";
+}
+
+/**
+ * `true` when `model` cannot decode WebP, otherwise `undefined` so the
+ * `PROMETHEUS_NO_WEBP` env fallback in {@link resizeImage} still applies. Feed straight
+ * into {@link ImageResizeOptions.excludeWebP}.
+ */
+export function webpExclusionForModel(model: Pick<Model, "provider" | "api"> | undefined): true | undefined {
+	return modelLacksWebpSupport(model) ? true : undefined;
+}
+
 export interface LoadImageInputOptions {
 	path: string;
 	cwd: string;
@@ -14,6 +34,8 @@ export interface LoadImageInputOptions {
 	maxBytes?: number;
 	resolvedPath?: string;
 	detectedMimeType?: string;
+	/** Force non-WebP output (e.g. for Ollama). Leave unset to honor `PROMETHEUS_NO_WEBP`. */
+	excludeWebP?: boolean;
 }
 
 export interface LoadedImageInput {
@@ -50,6 +72,41 @@ export async function ensureSupportedImageInput(image: ImageContent): Promise<Im
 	}
 }
 
+export interface NormalizeModelContextImagesOptions {
+	/** Model the images are bound for; used to derive encoder constraints (WebP exclusion for Ollama). */
+	model?: Model;
+	resize?: ImageResizeOptions;
+}
+
+/**
+ * Normalize image blocks before they enter agent/model context. This keeps
+ * provider request construction from having to resize an unbounded batch of
+ * large images on the streaming hot path. Images are processed sequentially on
+ * purpose: `resizeImage` may fan out multiple encoders for one image, so the
+ * outer image batch must stay bounded.
+ */
+export async function normalizeModelContextImages(
+	images: ImageContent[] | undefined,
+	options?: NormalizeModelContextImagesOptions,
+): Promise<ImageContent[] | undefined> {
+	if (!images || images.length === 0) return undefined;
+	const resize: ImageResizeOptions | undefined = modelLacksWebpSupport(options?.model)
+		? { ...options?.resize, excludeWebP: true }
+		: options?.resize;
+	const normalized: ImageContent[] = [];
+	for (const image of images) {
+		try {
+			const resized = await resizeImage(image, resize);
+			normalized.push({ type: "image", data: resized.data, mimeType: resized.mimeType });
+		} catch {
+			// Preserve existing caller behavior for decode/resize failures: keep the
+			// user's image block rather than dropping it from the turn.
+			normalized.push(image);
+		}
+	}
+	return normalized;
+}
+
 export async function loadImageInput(options: LoadImageInputOptions): Promise<LoadedImageInput | null> {
 	const maxBytes = options.maxBytes ?? MAX_IMAGE_INPUT_BYTES;
 	const resolvedPath = options.resolvedPath ?? resolveReadPath(options.path, options.cwd);
@@ -74,9 +131,13 @@ export async function loadImageInput(options: LoadImageInputOptions): Promise<Lo
 	let outputBytes = inputBuffer.byteLength;
 	let dimensionNote: string | undefined;
 
-	if (options.autoResize) {
+	const shouldReencodeWebP = options.excludeWebP === true && mimeType === "image/webp";
+	if (options.autoResize || shouldReencodeWebP) {
 		try {
-			const resized = await resizeImage({ type: "image", data: outputData, mimeType });
+			const resized = await resizeImage(
+				{ type: "image", data: outputData, mimeType },
+				{ excludeWebP: options.excludeWebP },
+			);
 			outputData = resized.data;
 			outputMimeType = resized.mimeType;
 			outputBytes = resized.buffer.byteLength;

package/src/utils/session-color.ts

@@ -1,4 +1,4 @@
-import { hslToHex, relativeLuminance } from "@prometheus-ai/utils";
+import { hexToHsv, hslToHex, relativeLuminance } from "@prometheus-ai/utils";
 
 /**
  * Derive a stable hue (0-359) from a string using djb2 hash.
@@ -25,23 +25,97 @@ function accentLuminanceCap(surfaceLuminance: number): number {
 	return Math.max(0, (surfaceLuminance + 0.05) / ACCENT_MIN_CONTRAST - 0.05);
 }
 
+/** Minimum angular distance in hue degrees from any theme color to avoid visual collision. */
+const MIN_HUE_DISTANCE = 10;
+/** Saturation threshold below which hue is meaningless (near-gray). */
+const MIN_SATURATION_FOR_HUE = 0.1;
+
+/** Angular distance between two hue values (0-360). */
+function hueDistance(a: number, b: number): number {
+	const d = Math.abs(a - b);
+	return Math.min(d, 360 - d);
+}
+
+/**
+ * Parse hue (0-360) from a hex color string.
+ * Returns undefined for near-gray colors where hue is not meaningful.
+ */
+function hexToHue(hex: string): number | undefined {
+	const hsv = hexToHsv(hex);
+	if (hsv.s < MIN_SATURATION_FOR_HUE) return undefined;
+	return hsv.h;
+}
+
+/**
+ * Find a hue at least {@link MIN_HUE_DISTANCE} from all occupied hues,
+ * clamped to [lo, hi] to prevent leaving the intended hue band.
+ * Returns `target` unchanged if no safe hue exists within bounds.
+ */
+function findSafeHue(target: number, occupied: number[], lo: number, hi: number): number {
+	if (occupied.length === 0) return target;
+	if (occupied.every(h => hueDistance(target, h) >= MIN_HUE_DISTANCE)) {
+		return target;
+	}
+	for (let d = 1; d <= hi - lo; d++) {
+		for (const dir of [1, -1]) {
+			const candidate = Math.max(lo, Math.min(hi, target + d * dir));
+			if (occupied.every(h => hueDistance(candidate, h) >= MIN_HUE_DISTANCE)) {
+				return candidate;
+			}
+		}
+	}
+	// fallback: keep the original target if no safe spot exists within the band
+	return target;
+}
+
+/** Hue range low and high for dark themes (warm: red → yellow → green). */
+const DARK_HUE_START = 0;
+const DARK_HUE_END = 120;
+/** Hue range low and high for light themes (cool: cyan → blue → purple). */
+const LIGHT_HUE_START = 180;
+const LIGHT_HUE_END = 300;
+
 /**
- * Derive a stable CSS hex accent color from a session name.
+ * Derive a stable CSS hex accent color from a session name and the active theme.
+ *
+ * Picks a hue from a **dark/light-specific range** so the accent feels natural
+ * for the theme type (warm on dark, cool on light). The session name hash
+ * determines the exact hue within the range. The result is checked against
+ * all theme color hues and shifted if it lands within {@link MIN_HUE_DISTANCE}
+ * of an existing theme hue, but is clamped to the hue band so it never
+ * drifts into an unrelated part of the spectrum.
  *
  * On dark themes (`surfaceLuminance` undefined) the accent is vivid (high
  * saturation, high lightness). On light themes the lightness is reduced until the
  * accent's perceived luminance clears {@link ACCENT_MIN_CONTRAST} against the
  * actual surface it renders on — so it stays legible on near-white *and* mid-light
- * backgrounds — while keeping the same per-session hue.
+ * backgrounds.
+ *
+ * @param name — session name for per-session uniqueness.
+ * @param themeColorHexes — all theme colors to check collision against.
+ * @param surfaceLuminance — undefined on dark themes; WCAG luminance of the
+ *   status-line background on light themes.
  */
-export function getSessionAccentHex(name: string, surfaceLuminance?: number): string {
-	const hue = nameToHue(name);
+export function getSessionAccentHex(name: string, themeColorHexes: string[], surfaceLuminance?: number): string {
+	// 1. Pick hue range based on theme mode
+	const hueStart = surfaceLuminance === undefined ? DARK_HUE_START : LIGHT_HUE_START;
+	const hueEnd = surfaceLuminance === undefined ? DARK_HUE_END : LIGHT_HUE_END;
+	const range = hueEnd - hueStart;
+
+	// 2. Session name picks within the range
+	let targetHue = hueStart + (nameToHue(name) % range);
+
+	// 3. Shift away if too close to any theme color — stays within [hueStart, hueEnd]
+	const themeHues = themeColorHexes.map(hexToHue).filter((h): h is number => h !== undefined);
+	targetHue = findSafeHue(targetHue, themeHues, hueStart, hueEnd);
+
+	// 4. Lightness/contrast — vivid on dark, bisected for AA on light
 	if (surfaceLuminance === undefined) {
-		return hslToHex(hue, ACCENT_SATURATION, ACCENT_DARK_LIGHTNESS);
+		return hslToHex(targetHue, ACCENT_SATURATION, ACCENT_DARK_LIGHTNESS);
 	}
 
 	const cap = accentLuminanceCap(surfaceLuminance);
-	const top = hslToHex(hue, ACCENT_SATURATION, ACCENT_DARK_LIGHTNESS);
+	const top = hslToHex(targetHue, ACCENT_SATURATION, ACCENT_DARK_LIGHTNESS);
 	if ((relativeLuminance(top) ?? 0) <= cap) return top;
 
 	// Bisect lightness: `lo` always yields luminance <= cap, `hi` always above it.
@@ -49,13 +123,13 @@ export function getSessionAccentHex(name: string, surfaceLuminance?: number): st
 	let hi = ACCENT_DARK_LIGHTNESS;
 	for (let i = 0; i < 20; i++) {
 		const mid = (lo + hi) / 2;
-		if ((relativeLuminance(hslToHex(hue, ACCENT_SATURATION, mid)) ?? 0) > cap) {
+		if ((relativeLuminance(hslToHex(targetHue, ACCENT_SATURATION, mid)) ?? 0) > cap) {
 			hi = mid;
 		} else {
 			lo = mid;
 		}
 	}
-	return hslToHex(hue, ACCENT_SATURATION, lo);
+	return hslToHex(targetHue, ACCENT_SATURATION, lo);
 }
 
 /**

package/src/utils/thinking-display.ts

@@ -0,0 +1,37 @@
+import type { AssistantMessage } from "@prometheus-ai/ai";
+
+type AssistantContentBlock = AssistantMessage["content"][number];
+type ThinkingBlock = Extract<AssistantContentBlock, { type: "thinking" }>;
+
+function isDotOnlyThinking(text: string): boolean {
+	let sawDot = false;
+	for (let i = 0; i < text.length; i++) {
+		const code = text.charCodeAt(i);
+		if (code === 0x2e || code === 0x2026) {
+			sawDot = true;
+			continue;
+		}
+		if (code === 0x20 || code === 0x09 || code === 0x0a || code === 0x0d) continue;
+		return false;
+	}
+	return sawDot;
+}
+
+/**
+ * Returns the operator-visible thinking text for a block.
+ *
+ * Some OpenAI-compatible reasoning gateways require a non-empty
+ * `reasoning_content` field on historical assistant tool-call turns even when
+ * the model did not emit any reasoning. The provider adapter uses a single dot
+ * as the wire-only placeholder those gateways accept; if that value is later
+ * replayed or echoed as a thinking block, it should not render as model thought.
+ */
+export function getVisibleThinkingText(block: ThinkingBlock): string {
+	const text = block.thinking.trim();
+	if (text.length === 0) return "";
+	return isDotOnlyThinking(text) ? "" : text;
+}
+
+export function hasVisibleThinking(block: ThinkingBlock): boolean {
+	return getVisibleThinkingText(block).length > 0;
+}

package/src/utils/title-generator.ts

@@ -6,14 +6,19 @@ import * as path from "node:path";
 import { type Api, type AssistantMessage, completeSimple, type Model, type Tool } from "@prometheus-ai/ai";
 import { logger, prompt } from "@prometheus-ai/utils";
 import type { ModelRegistry } from "../config/model-registry";
+
 import { resolveRoleSelection } from "../config/model-resolver";
 import type { Settings } from "../config/settings";
+import titleMarkerInstruction from "../prompts/system/title-marker-instruction.md" with { type: "text" };
 import titleSystemPrompt from "../prompts/system/title-system.md" with { type: "text" };
+import titleMarkerSystemPrompt from "../prompts/system/title-system-marker.md" with { type: "text" };
 import { ONLINE_TINY_TITLE_MODEL_KEY } from "../tiny/models";
 import { formatTitleUserMessage, isLowSignalTitleInput, normalizeGeneratedTitle } from "../tiny/text";
 import { tinyTitleClient } from "../tiny/title-client";
 
 const TITLE_SYSTEM_PROMPT = prompt.render(titleSystemPrompt);
+const TITLE_MARKER_SYSTEM_PROMPT = prompt.render(titleMarkerSystemPrompt);
+const TITLE_MARKER_INSTRUCTION = prompt.render(titleMarkerInstruction);
 
 const DEFAULT_TERMINAL_TITLE = "π";
 const TERMINAL_TITLE_CONTROL_CHARS = /[\u0000-\u001f\u007f-\u009f]/g;
@@ -32,7 +37,7 @@ const setTitleTool: Tool = {
 			title: {
 				type: "string",
 				description:
-					'A concise 3-6 word title for the session, or exactly "none" when the message carries no concrete task yet (greeting, small talk, vague).',
+					'The generated session title, or exactly "none" when the message carries no concrete task yet.',
 			},
 		},
 		required: ["title"],
@@ -40,11 +45,35 @@ const setTitleTool: Tool = {
 	},
 };
 
+/** Matches the title a tool-choice-less model wraps in `<title>...</title>`. */
+const TITLE_MARKER_RE = /<title>([\s\S]*?)<\/title>/i;
+
+/**
+ * Whether the model honors a forced `tool_choice` so the `set_title` tool can be
+ * required. Providers/models that reject forced tool calls (chat-completions
+ * hosts without `tool_choice` support, Claude Fable/Mythos) can't be made to
+ * emit a structured call, so the caller falls back to marker-wrapped text.
+ */
+function modelSupportsForcedToolChoice(model: Model<Api>): boolean {
+	// `compat` is a union across APIs and `supportsToolChoice` lives only on the
+	// OpenAI-completions variant, so read both flags through a structural view.
+	const compat = model.compat as { supportsToolChoice?: boolean; supportsForcedToolChoice?: boolean } | undefined;
+	if (!compat) return true;
+	// A forced tool call first requires sending `tool_choice` at all. Hosts that
+	// drop the parameter entirely (`supportsToolChoice: false`, e.g. direct
+	// DeepSeek reasoning) can never be forced even when they otherwise accept
+	// forced values, so this veto wins over `supportsForcedToolChoice`.
+	if (compat.supportsToolChoice === false) return false;
+	if (typeof compat.supportsForcedToolChoice === "boolean") return compat.supportsForcedToolChoice;
+	if (typeof compat.supportsToolChoice === "boolean") return compat.supportsToolChoice;
+	return true;
+}
+
 function getTitleModel(registry: ModelRegistry, settings: Settings, currentModel?: Model<Api>): Model<Api> | undefined {
 	const availableModels = registry.getAvailable();
 	if (availableModels.length === 0) return undefined;
 
-	const titleModel = resolveRoleSelection(["commit", "smol"], settings, availableModels, registry)?.model;
+	const titleModel = resolveRoleSelection(["title", "commit", "smol"], settings, availableModels, registry)?.model;
 	if (titleModel) return titleModel;
 
 	if (currentModel) return currentModel;
@@ -136,6 +165,7 @@ export async function raceFirstNonNull<T>(
  *   to produce request metadata (e.g. user_id for session attribution). Using a
  *   resolver instead of a pre-evaluated value ensures the metadata's account_uuid
  *   reflects the credential actually selected for this request.
+ * @param customSystemPrompt Optional title-specific system prompt override
  */
 export async function generateSessionTitle(
 	firstMessage: string,
@@ -144,6 +174,7 @@ export async function generateSessionTitle(
 	sessionId?: string,
 	currentModel?: Model<Api>,
 	metadataResolver?: (provider: string) => Record<string, unknown> | undefined,
+	customSystemPrompt?: string,
 ): Promise<string | null> {
 	// Defer titling for greetings / acknowledgements / empty input. The default
 	// tiny title model can't reliably decline trivial input, so this happens
@@ -154,13 +185,26 @@ export async function generateSessionTitle(
 		return null;
 	}
 
+	const titleSystemPrompt = customSystemPrompt?.trim() || undefined;
 	const tinyModel = settings.get("providers.tinyModel");
 	if (tinyModel === ONLINE_TINY_TITLE_MODEL_KEY) {
-		return generateTitleOnline(firstMessage, registry, settings, sessionId, currentModel, metadataResolver);
+		return generateTitleOnline(
+			firstMessage,
+			registry,
+			settings,
+			sessionId,
+			currentModel,
+			metadataResolver,
+			undefined,
+			titleSystemPrompt,
+		);
 	}
 
 	const onlineAbortController = new AbortController();
-	const localTitle = tinyTitleClient.generate(tinyModel, firstMessage).then(
+	const localTitlePromise = titleSystemPrompt
+		? tinyTitleClient.generate(tinyModel, firstMessage, { systemPrompt: titleSystemPrompt })
+		: tinyTitleClient.generate(tinyModel, firstMessage);
+	const localTitle = localTitlePromise.then(
 		title => title || null,
 		err => {
 			logger.warn("title-generator: local model error", {
@@ -180,6 +224,7 @@ export async function generateSessionTitle(
 			currentModel,
 			metadataResolver,
 			onlineAbortController.signal,
+			titleSystemPrompt,
 		);
 
 	return raceFirstNonNull(localTitle, startOnline, TITLE_LOCAL_FALLBACK_DELAY_MS, () => {
@@ -195,6 +240,7 @@ export async function generateTitleOnline(
 	currentModel?: Model<Api>,
 	metadataResolver?: (provider: string) => Record<string, unknown> | undefined,
 	signal?: AbortSignal,
+	customSystemPrompt?: string,
 ): Promise<string | null> {
 	const model = getTitleModel(registry, settings, currentModel);
 	if (!model) {
@@ -202,6 +248,17 @@ export async function generateTitleOnline(
 		return null;
 	}
 
+	const titleSystemPrompt = customSystemPrompt?.trim() || undefined;
+	// Some providers can't be forced to call a tool — chat-completions hosts
+	// without `tool_choice` support, Claude Fable/Mythos — so a required
+	// `set_title` call never arrives. For those, ask the model to wrap the title
+	// in `<title>...</title>` markers and parse it from text instead.
+	const useForcedTool = modelSupportsForcedToolChoice(model);
+	const systemPrompt = useForcedTool
+		? [titleSystemPrompt ?? TITLE_SYSTEM_PROMPT]
+		: titleSystemPrompt
+			? [titleSystemPrompt, TITLE_MARKER_INSTRUCTION]
+			: [TITLE_MARKER_SYSTEM_PROMPT];
 	const userMessage = formatTitleUserMessage(firstMessage);
 	const modelName = `${model.provider}/${model.id}`;
 	const modelContext = {
@@ -223,7 +280,7 @@ export async function generateTitleOnline(
 		// account_uuid rather than the snapshot-at-call-site value.
 		const metadata = metadataResolver?.(model.provider);
 
-		// Title generation is a 3-6 word task, but some reasoning backends ignore
+		// Title generation is a 3-7 word task, but some reasoning backends ignore
 		// disableReasoning. Keep the normal cheap budget for non-reasoning models
 		// while reserving enough output room for reasoning models to still emit
 		// the forced tool call after any unavoidable thinking tokens.
@@ -233,15 +290,15 @@ export async function generateTitleOnline(
 		const response = await completeSimple(
 			model,
 			{
-				systemPrompt: [TITLE_SYSTEM_PROMPT],
+				systemPrompt,
 				messages: [{ role: "user", content: userMessage, timestamp: Date.now() }],
-				tools: [setTitleTool],
+				tools: useForcedTool ? [setTitleTool] : undefined,
 			},
 			{
-				apiKey,
+				apiKey: registry.resolver(model, sessionId),
 				maxTokens,
 				disableReasoning: true,
-				toolChoice: { type: "tool", name: SET_TITLE_TOOL_NAME },
+				toolChoice: useForcedTool ? { type: "tool", name: SET_TITLE_TOOL_NAME } : undefined,
 				metadata,
 				signal,
 			},
@@ -299,7 +356,13 @@ function extractGeneratedTitle(contentBlocks: AssistantMessage["content"]): stri
 			textTitle += content.text;
 		}
 	}
-	return textTitle.trim();
+	// Tool-choice-less models are asked to wrap the title in <title>...</title>,
+	// but stay lenient: prefer the marker when the model closed it, otherwise
+	// accept a plain sentence after stripping any stray/unclosed tag fragment
+	// (e.g. output truncated before the closing tag).
+	const marker = TITLE_MARKER_RE.exec(textTitle);
+	if (marker) return marker[1].trim();
+	return textTitle.replace(/<\/?title>/gi, "").trim();
 }
 
 /**

package/src/utils/tool-choice.ts

@@ -31,3 +31,19 @@ export function buildNamedToolChoice(toolName: string, model?: Model<Api>): Tool
 
 	return undefined;
 }
+
+/**
+ * Whether the given tool choice can be satisfied by the active tool set for the
+ * upcoming turn. Non-named choices (`"none"`, `"required"`, etc.) do not name a
+ * specific tool and are therefore always active.
+ */
+export function isToolChoiceActive(toolChoice: ToolChoice | undefined, tools: readonly { name: string }[]): boolean {
+	if (!toolChoice || typeof toolChoice === "string") return true;
+	const name =
+		toolChoice.type === "tool"
+			? toolChoice.name
+			: "function" in toolChoice
+				? toolChoice.function.name
+				: toolChoice.name;
+	return tools.some(tool => tool.name === name);
+}

package/src/utils/tools-manager.ts

@@ -16,6 +16,16 @@ interface ToolConfig {
 	getAssetName: (version: string, plat: string, architecture: string) => string | null;
 }
 
+// ffmpeg static-binary asset names (eugeneware/ffmpeg-static direct binaries).
+// Maps node arch (arm64|x64) only; everything else is unsupported.
+export function ffmpegAssetName(_version: string, plat: string, architecture: string): string | null {
+	if (architecture !== "arm64" && architecture !== "x64") return null;
+	if (plat === "darwin") return `ffmpeg-darwin-${architecture}`;
+	if (plat === "linux") return `ffmpeg-linux-${architecture}`;
+	if (plat === "win32") return architecture === "x64" ? "ffmpeg-win32-x64" : null;
+	return null;
+}
+
 const TOOLS: Record<string, ToolConfig> = {
 	sd: {
 		name: "sd",
@@ -72,6 +82,14 @@ const TOOLS: Record<string, ToolConfig> = {
 			return null;
 		},
 	},
+	ffmpeg: {
+		name: "ffmpeg",
+		repo: "eugeneware/ffmpeg-static",
+		binaryName: "ffmpeg",
+		tagPrefix: "",
+		isDirectBinary: true,
+		getAssetName: ffmpegAssetName,
+	},
 };
 
 // CLI packages installed via uv/pip
@@ -89,7 +107,7 @@ const PYTHON_TOOLS: Record<string, PythonPackageToolConfig> = {
 	},
 };
 
-export type ToolName = "sd" | "sg" | "yt-dlp" | "trafilatura";
+export type ToolName = "sd" | "sg" | "yt-dlp" | "trafilatura" | "ffmpeg";
 
 // Get the path to a tool (system-wide or in our tools dir)
 export function getToolPath(tool: ToolName): string | null {

package/src/web/kagi.ts

@@ -6,7 +6,7 @@
  * through the shared {@link AuthStorage} broker (Bearer token), and responses
  * are categorized result buckets rather than the legacy flat object array.
  */
-import type { AuthStorage } from "@prometheus-ai/ai";
+import { type AuthStorage, type FetchImpl, withAuth } from "@prometheus-ai/ai";
 import { withHardTimeout } from "./search/providers/utils";
 
 const KAGI_SEARCH_URL = "https://kagi.com/api/v1/search";
@@ -156,6 +156,7 @@ export interface KagiSearchOptions {
 	recency?: "day" | "week" | "month" | "year";
 	sessionId?: string;
 	signal?: AbortSignal;
+	fetch?: FetchImpl;
 }
 
 export interface KagiSearchSource {
@@ -172,14 +173,6 @@ export interface KagiSearchResult {
 	answer?: string;
 }
 
-export async function findKagiApiKey(
-	authStorage: AuthStorage,
-	sessionId?: string,
-	signal?: AbortSignal,
-): Promise<string | null> {
-	return (await authStorage.getApiKey("kagi", sessionId, { signal })) ?? null;
-}
-
 /**
  * Compute a YYYY-MM-DD date string `recency` units before now, in UTC.
  * UTC keeps the recency window deterministic regardless of host timezone and
@@ -246,25 +239,34 @@ export async function searchWithKagi(
 	options: KagiSearchOptions = {},
 	authStorage: AuthStorage,
 ): Promise<KagiSearchResult> {
-	const apiKey = await findKagiApiKey(authStorage, options.sessionId, options.signal);
-	if (!apiKey) {
-		throw new KagiApiError("Kagi credentials not found. Set KAGI_API_KEY or login with 'prometheus /login kagi'.");
-	}
+	const fetchImpl = options.fetch ?? fetch;
+	const body = JSON.stringify(buildRequestBody(query, options));
+
+	const response = await withAuth(
+		authStorage.resolver("kagi", { sessionId: options.sessionId }),
+		async apiKey => {
+			const res = await fetchImpl(KAGI_SEARCH_URL, {
+				method: "POST",
+				headers: {
+					Authorization: `Bearer ${apiKey}`,
+					"Content-Type": "application/json",
+					Accept: "application/json",
+				},
+				body,
+				signal: withHardTimeout(options.signal),
+			});
+
+			if (!res.ok) {
+				throw parseKagiErrorResponse(res.status, await res.text());
+			}
 
-	const response = await fetch(KAGI_SEARCH_URL, {
-		method: "POST",
-		headers: {
-			Authorization: `Bearer ${apiKey}`,
-			"Content-Type": "application/json",
-			Accept: "application/json",
+			return res;
 		},
-		body: JSON.stringify(buildRequestBody(query, options)),
-		signal: withHardTimeout(options.signal),
-	});
-
-	if (!response.ok) {
-		throw parseKagiErrorResponse(response.status, await response.text());
-	}
+		{
+			signal: options.signal,
+			missingKeyMessage: "Kagi credentials not found. Set KAGI_API_KEY or login with 'prometheus /login kagi'.",
+		},
+	);
 
 	const payload = (await response.json()) as KagiSearchResponse;
 	if (payload.error && payload.error.length > 0) {

package/src/web/parallel.ts

@@ -1,4 +1,4 @@
-import { getEnvApiKey } from "@prometheus-ai/ai";
+import { type FetchImpl, getEnvApiKey } from "@prometheus-ai/ai";
 import type { AgentStorage } from "../session/agent-storage";
 import { findCredential, withHardTimeout } from "./search/providers/utils";
 
@@ -54,6 +54,7 @@ export interface ParallelSearchOptions {
 	mode?: "fast" | "research";
 	maxCharsPerResult?: number;
 	signal?: AbortSignal;
+	fetch?: FetchImpl;
 }
 
 export interface ParallelExtractOptions {
@@ -62,6 +63,7 @@ export interface ParallelExtractOptions {
 	excerpts?: boolean;
 	fullContent?: boolean;
 	signal?: AbortSignal;
+	fetch?: FetchImpl;
 }
 
 export class ParallelApiError extends Error {
@@ -295,7 +297,8 @@ export async function searchWithParallel(
 		);
 	}
 
-	const response = await fetch(PARALLEL_SEARCH_URL, {
+	const fetchImpl = options.fetch ?? fetch;
+	const response = await fetchImpl(PARALLEL_SEARCH_URL, {
 		method: "POST",
 		headers: getAuthHeaders(apiKey),
 		body: JSON.stringify({
@@ -328,7 +331,8 @@ export async function extractWithParallel(
 		);
 	}
 
-	const response = await fetch(PARALLEL_EXTRACT_URL, {
+	const fetchImpl = options.fetch ?? fetch;
+	const response = await fetchImpl(PARALLEL_EXTRACT_URL, {
 		method: "POST",
 		headers: getAuthHeaders(apiKey),
 		body: JSON.stringify({

package/src/web/scrapers/arxiv.ts

@@ -1,4 +1,3 @@
-import { parseHTML } from "linkedom";
 import type { RenderResult, SpecialHandler } from "./types";
 import { buildResult, loadPage } from "./types";
 import { convertWithMarkit, fetchBinary } from "./utils";
@@ -31,6 +30,7 @@ export const handleArxiv: SpecialHandler = async (
 		if (!result.ok) return null;
 
 		// Parse the Atom feed response
+		const { parseHTML } = await import("linkedom");
 		const doc = parseHTML(result.content).document;
 		const entry = doc.querySelector("entry");