@iowarp/clio-coder 0.4.9 → 0.5.0

package/src/engine/apis/ollama-native.ts

@@ -14,15 +14,12 @@ import type {
 	ToolCall,
 	Usage,
 } from "@earendil-works/pi-ai";
-import { createAssistantMessageEventStream } from "@earendil-works/pi-ai";
-import {
-	type ChatRequest,
-	type ChatResponse,
-	Ollama,
-	type Message as OllamaMessage,
-	type Options as OllamaOptions,
-	type Tool as OllamaTool,
-	type ToolCall as OllamaToolCall,
+import type {
+	ChatRequest,
+	Message as OllamaMessage,
+	Options as OllamaOptions,
+	Tool as OllamaTool,
+	ToolCall as OllamaToolCall,
 } from "ollama";
 import { residencyTargetKey } from "../../core/residency-target-key.js";
 import {
@@ -30,18 +27,36 @@ import {
 	resolveModelRuntimeCapabilitiesForModel,
 	type ThinkingLevel,
 } from "../../domains/providers/index.js";
+import { ollamaModelIds } from "../../domains/providers/runtimes/common/ollama-model-ids.js";
 import type { LocalModelQuirks, SamplingProfile } from "../../domains/providers/types/local-model-quirks.js";
 import { calculateEngineCost } from "../ai.js";
+import { resolvedRequestContext } from "../context.js";
 import { createGemmaChannelFilter, usesGemmaChannelMarkers } from "../gemma-channel-filter.js";
 import { createSentinelStripper } from "../strip-tokenizer-sentinels.js";
+import { createDegradedInferenceStream } from "./degraded-inference.js";
+import { ollamaJson, streamOllamaChat } from "./ollama-http.js";
 import { remainingContextMaxTokens } from "./output-budget.js";
-import { type ResidencyAdapter, reconcileResidency, residencyManagedFor } from "./residency.js";
-import type { ResidentModelInfo } from "./resident-models.js";
-import { mergeSamplingOverride } from "./sampling-overrides.js";
+import {
+	EXIT_RELEASE_MS,
+	forgetReleasedModel,
+	isClioLoaded,
+	markClioLoaded,
+	type ReleaseScope,
+	type ResidencyAdapter,
+	reconcileResidency,
+	registerExitRelease,
+	registerWorkerLoadAdopter,
+	reportClioModelLoad,
+	residencyManagedFor,
+} from "./residency.js";
+import { type ResidentModelInfo, residentMatchesKeep } from "./resident-models.js";
+import { pickSamplingProfile, samplingParamsFromProfile } from "./sampling-overrides.js";
 import type { EngineApiProvider } from "./types.js";
 
 const REASONING_CHARS_PER_TOKEN = 4;
 const ownedModelsByTarget = new Map<string, Set<string>>();
+/** How to reach each target that holds an owned model, for the release on exit. */
+const ownedEndpointsByTarget = new Map<string, { baseUrl: string; headers: Record<string, string> }>();
 
 interface ClioRuntimeMetadata {
 	clioCoder?: {
@@ -52,6 +67,7 @@ interface ClioRuntimeMetadata {
 		gateway?: boolean;
 		family?: string;
 		quirks?: LocalModelQuirks;
+		ollama?: { numCtx?: number };
 	};
 }
 
@@ -72,46 +88,45 @@ function ollamaTargetId(model: Model<"ollama-native">): string {
  * and nothing else ever reclaims. Best-effort: a failure never blocks the
  * turn.
  */
-async function reconcileOllamaResidency(model: Model<"ollama-native">, headers: Record<string, string>): Promise<void> {
+async function reconcileOllamaResidency(
+	model: Model<"ollama-native">,
+	headers: Record<string, string>,
+	signal?: AbortSignal,
+): Promise<boolean> {
 	const baseUrl = model.baseUrl;
-	if (!baseUrl) return;
+	if (!baseUrl) return false;
 	const metadata = (model as Model<"ollama-native"> & ClioRuntimeMetadata).clioCoder;
+	if (!residencyManagedFor(metadata?.lifecycle)) return false;
+	let resident: ResidentModelInfo[] = [];
 	const adapter: ResidencyAdapter = {
-		targetKey: residencyTargetKey("ollama-native", baseUrl),
+		targetKey: residencyTargetKey("ollama", baseUrl),
 		targetId: ollamaTargetId(model),
-		runtimeId: "ollama-native",
+		runtimeId: "ollama",
 		keepModelId: model.id,
-		managed: residencyManagedFor(metadata?.lifecycle),
+		managed: true,
 		strategy: "scheduler",
-		listResident: () => listResidentOllamaModels(ollamaEvictClient(baseUrl, headers)),
-		unload: (id) => unloadOllamaModel(baseUrl, id, headers),
+		...(signal ? { signal } : {}),
+		listResident: async () => (resident = await listResidentOllamaModels(baseUrl, headers, signal)),
+		unload: (id) =>
+			unloadOllamaModel(baseUrl, resident.find((entry) => entry.modelId === id) ?? { modelId: id }, headers, signal),
 	};
 	try {
-		await reconcileResidency(adapter);
+		const plan = await reconcileResidency(adapter);
+		const keep = resident.find((entry) => residentMatchesKeep(entry, model.id));
+		const ids = ollamaModelIds(model.id, ...(keep ? [keep.modelId, ...(keep.aliasIds ?? [])] : []));
+		return (
+			plan.decision === "reconcile" &&
+			(!plan.keepResident || ids.some((id) => ownedModelsByTarget.get(adapter.targetKey)?.has(id)))
+		);
 	} catch {
-		// Reconciliation is best-effort; a failure must never block the turn.
+		signal?.throwIfAborted();
+		return false;
 	}
 }
 
-function pickSamplingProfile(
-	quirks: LocalModelQuirks | undefined,
-	thinkingActive: boolean,
-): SamplingProfile | undefined {
-	const sampling = quirks?.sampling;
-	const profile = sampling ? (thinkingActive ? (sampling.thinking ?? sampling.instruct) : sampling.instruct) : undefined;
-	return mergeSamplingOverride(profile);
-}
-
 function applyOllamaSamplingProfile(opts: Partial<OllamaOptions>, profile: SamplingProfile): void {
 	if (profile.temperature !== undefined && opts.temperature === undefined) opts.temperature = profile.temperature;
-	if (profile.topP !== undefined && opts.top_p === undefined) opts.top_p = profile.topP;
-	if (profile.topK !== undefined && opts.top_k === undefined) opts.top_k = profile.topK;
-	if (profile.repeatPenalty !== undefined && opts.repeat_penalty === undefined)
-		opts.repeat_penalty = profile.repeatPenalty;
-	if (profile.presencePenalty !== undefined && opts.presence_penalty === undefined)
-		opts.presence_penalty = profile.presencePenalty;
-	if (profile.frequencyPenalty !== undefined && opts.frequency_penalty === undefined)
-		opts.frequency_penalty = profile.frequencyPenalty;
+	Object.assign(opts, { ...samplingParamsFromProfile(profile, "ollama"), ...opts });
 }
 
 function isOllamaEffort(value: string | undefined): value is "low" | "medium" | "high" {
@@ -214,15 +229,14 @@ function buildRequest(
 	context: Context,
 	options: StreamOptions | undefined,
 	thinkingLevel: ThinkingLevel,
+	pin: boolean,
 ): ChatRequest & { stream: true } {
 	const req: ChatRequest & { stream: true } = {
 		model: model.id,
 		messages: buildMessages(context),
 		stream: true,
-		// Pin the active model resident. The pre-turn reconciler releases
-		// Clio-pinned stragglers this leaves behind after a model switch;
-		// operator-loaded and configured models stay resident.
-		keep_alive: -1,
+		// Only pin loads we can own and release. Otherwise use the server's policy.
+		...(pin ? { keep_alive: -1 } : {}),
 	};
 	if (context.tools && context.tools.length > 0) req.tools = context.tools.map(toolToOllama);
 	const opts: Partial<OllamaOptions> = {};
@@ -234,6 +248,11 @@ function buildRequest(
 	if (samplingProfile) applyOllamaSamplingProfile(opts, samplingProfile);
 	if (options?.temperature !== undefined) opts.temperature = options.temperature;
 	opts.num_predict = remainingContextMaxTokens(model, context, options);
+	// Sent only when the operator configured it: a `num_ctx` that differs from
+	// the loaded one makes Ollama reload the model, which on a shared server
+	// evicts whatever window another client opened it at.
+	const numCtx = (model as Model<"ollama-native"> & ClioRuntimeMetadata).clioCoder?.ollama?.numCtx;
+	if (numCtx !== undefined) opts.num_ctx = numCtx;
 	if (Object.keys(opts).length > 0) req.options = opts;
 	return req;
 }
@@ -251,35 +270,24 @@ export interface EvictResidentResponse {
 	readonly models: ReadonlyArray<EvictResidentEntry>;
 }
 
-export interface EvictGenerateRequest {
-	readonly model: string;
-	readonly prompt: string;
-	readonly keep_alive: number;
-	readonly stream: false;
-}
-
-export interface OllamaEvictClient {
-	ps(): Promise<EvictResidentResponse>;
-	generate(req: EvictGenerateRequest): Promise<unknown>;
-}
-
-function ollamaEvictClient(baseUrl: string, headers?: Record<string, string>): OllamaEvictClient {
-	return new Ollama({ host: baseUrl, ...(headers ? { headers } : {}) });
-}
-
 /**
  * Map an Ollama `/api/ps` response to the runtime-agnostic resident shape,
  * preserving the GPU/total footprint when the server reports it.
  */
-async function listResidentOllamaModels(client: OllamaEvictClient): Promise<ResidentModelInfo[]> {
-	const resident = await client.ps();
+async function listResidentOllamaModels(
+	baseUrl: string,
+	headers?: Record<string, string>,
+	signal?: AbortSignal,
+): Promise<ResidentModelInfo[]> {
+	const resident = await ollamaJson<EvictResidentResponse>(baseUrl, "ps", undefined, { headers, signal });
 	return resident.models.map((entry) => {
 		const primary = entry.model || entry.name;
+		if (typeof primary !== "string" || !primary.trim()) throw new Error("Ollama returned an invalid resident model");
 		// Ollama reports both `model` and `name`; keep the other one as an alias so
 		// a keep target that matches either field is never evicted.
-		const aliases = [entry.model, entry.name].filter(
-			(id): id is string => typeof id === "string" && id.length > 0 && id !== primary,
-		);
+		const aliases = ollamaModelIds(
+			...[entry.model, entry.name].filter((id): id is string => typeof id === "string" && id.length > 0),
+		).filter((id) => id !== primary);
 		const info: ResidentModelInfo = { modelId: primary };
 		if (aliases.length > 0) info.aliasIds = aliases;
 		if (typeof entry.size_vram === "number") info.sizeVramBytes = entry.size_vram;
@@ -293,13 +301,92 @@ async function listResidentOllamaModels(client: OllamaEvictClient): Promise<Resi
  * `keep_alive: -1`, so eviction fires `keep_alive: 0` against it to let its
  * weights release.
  */
-async function unloadOllamaModel(baseUrl: string, modelId: string, headers?: Record<string, string>): Promise<void> {
-	const owned = ownedModelsByTarget.get(residencyTargetKey("ollama-native", baseUrl));
-	if (!owned?.has(modelId)) return;
-	await ollamaEvictClient(baseUrl, headers).generate({ model: modelId, prompt: "", keep_alive: 0, stream: false });
-	owned.delete(modelId);
+async function unloadOllamaModel(
+	baseUrl: string,
+	entry: ResidentModelInfo,
+	headers?: Record<string, string>,
+	signal?: AbortSignal,
+): Promise<void> {
+	const targetKey = residencyTargetKey("ollama", baseUrl);
+	const owned = ownedModelsByTarget.get(targetKey);
+	const ids = [entry.modelId, ...(entry.aliasIds ?? [])];
+	if (!ids.some((id) => owned?.has(id))) throw new Error("Ollama model ownership is unknown");
+	await ollamaJson(
+		baseUrl,
+		"generate",
+		{ model: entry.modelId, prompt: "", keep_alive: 0, stream: false },
+		{ headers, signal },
+	);
+	for (const id of ids) {
+		owned?.delete(id);
+		forgetReleasedModel(targetKey, id);
+	}
+}
+
+/**
+ * Release on process exit every Ollama model this process loaded and pinned
+ * with `keep_alive: -1`, so a finished run does not hold its weights forever
+ * (#379). A model counts as this process's own only when both registries agree:
+ * the residency reconciler saw it absent from the server before Clio's request
+ * and marked it Clio-loaded, and a chat for it then streamed, which is what
+ * pinned it. A model that was already resident when Clio first touched it, or
+ * one whose load never produced a response, is left alone, so an operator's
+ * model is never unloaded (#313). The resident list is read first, and only
+ * models still resident are released, so a model Ollama already dropped is
+ * not loaded again just to be unloaded. A `scope` narrows the release further,
+ * which is how a one-shot probe releases only what it loaded.
+ *
+ * Best-effort and bounded: every request shares one deadline, and a failure
+ * or timeout leaves the model pinned and never throws. A crashed process
+ * releases nothing.
+ */
+export async function releaseClioLoadedOllamaModels(
+	options: { timeoutMs?: number; scope?: ReleaseScope } = {},
+): Promise<void> {
+	const controller = new AbortController();
+	const timer = setTimeout(() => controller.abort(), options.timeoutMs ?? EXIT_RELEASE_MS);
+	timer.unref();
+	const releaseTarget = async (targetKey: string, owned: Set<string>): Promise<void> => {
+		const endpoint = ownedEndpointsByTarget.get(targetKey);
+		if (!endpoint || owned.size === 0) return;
+		const resident = await listResidentOllamaModels(endpoint.baseUrl, endpoint.headers, controller.signal);
+		for (const entry of resident) {
+			const ids = [entry.modelId, ...(entry.aliasIds ?? [])];
+			if (!ids.some((id) => owned.has(id)) || !isClioLoaded(targetKey, entry)) continue;
+			if (options.scope && !options.scope(targetKey, ids)) continue;
+			await unloadOllamaModel(endpoint.baseUrl, entry, endpoint.headers, controller.signal);
+		}
+	};
+	const releases = [...ownedModelsByTarget].map(([targetKey, owned]) =>
+		releaseTarget(targetKey, owned).catch(() => {
+			// Best-effort: an unreachable server keeps the model pinned.
+		}),
+	);
+	try {
+		await Promise.race([
+			Promise.all(releases),
+			new Promise<void>((resolve) => controller.signal.addEventListener("abort", () => resolve(), { once: true })),
+		]);
+	} finally {
+		clearTimeout(timer);
+		controller.abort();
+	}
 }
 
+registerExitRelease((scope) => releaseClioLoadedOllamaModels(scope ? { scope } : {}));
+
+// A dispatched worker's load becomes this process's to release (#379). The
+// endpoint comes from this process's own settings for the admitted target.
+registerWorkerLoadAdopter("ollama", (targetKey, endpoint, modelIds) => {
+	let owned = ownedModelsByTarget.get(targetKey);
+	if (!owned) {
+		owned = new Set<string>();
+		ownedModelsByTarget.set(targetKey, owned);
+	}
+	for (const id of modelIds) owned.add(id);
+	ownedEndpointsByTarget.set(targetKey, { baseUrl: endpoint.baseUrl, headers: endpoint.headers });
+});
+
 function mapStopReason(reason: string | undefined, hadToolCall: boolean): AssistantMessage["stopReason"] {
 	if (hadToolCall) return "toolUse";
 	if (reason === "length") return "length";
@@ -322,7 +409,7 @@ function emitToolCall(raw: OllamaToolCall, output: AssistantMessage, stream: Ass
 		type: "toolCall",
 		id: randomUUID(),
 		name: raw.function?.name ?? "",
-		arguments: args,
+		arguments: args as ToolCall["arguments"],
 	};
 	output.content.push(toolCall);
 	const idx = output.content.length - 1;
@@ -342,7 +429,6 @@ function runStream(
 	options: StreamOptions | undefined,
 	thinkingLevel: ThinkingLevel,
 ): AssistantMessageEventStream {
-	const stream: AssistantMessageEventStream = createAssistantMessageEventStream();
 	const output: AssistantMessage = {
 		role: "assistant",
 		content: [],
@@ -360,25 +446,36 @@ function runStream(
 		stopReason: "stop",
 		timestamp: Date.now(),
 	};
-	// Fresh instance per call so instance-level abort() scopes to this stream only.
-	const headers: Record<string, string> = {};
-	if (model.headers) Object.assign(headers, model.headers);
-	if (options?.headers) Object.assign(headers, options.headers);
-	const client = new Ollama({ host: model.baseUrl, headers });
+	const requestHeaders = new Headers(model.headers);
+	for (const [name, value] of Object.entries(options?.headers ?? {})) {
+		if (value === null) requestHeaders.delete(name);
+		else if (value !== undefined) requestHeaders.set(name, value);
+	}
+	if (options?.apiKey && !requestHeaders.has("authorization")) {
+		requestHeaders.set("authorization", `Bearer ${options.apiKey}`);
+	}
+	const headers = Object.fromEntries(requestHeaders);
 	const signal = options?.signal;
-	let aborted = signal?.aborted === true;
-	const onAbort = () => {
-		aborted = true;
-		client.abort();
-	};
-	if (signal && !signal.aborted) signal.addEventListener("abort", onAbort, { once: true });
+	const baseUrl = model.baseUrl;
+	// Events go straight into the watched stream, so this function's own catch
+	// still decides how a failed turn ends.
+	const stream = createDegradedInferenceStream({
+		targetId: ollamaTargetId(model),
+		runtimeId: "ollama",
+		model: model.id,
+		...(signal ? { signal } : {}),
+		...(baseUrl ? { listResident: () => listResidentOllamaModels(baseUrl, headers, signal) } : {}),
+	});
 	(async () => {
 		try {
-			if (aborted) throw new Error("Request was aborted");
-			await reconcileOllamaResidency(model, headers);
-			if (aborted) throw new Error("Request was aborted");
-			const iterator = await client.chat(buildRequest(model, context, options, thinkingLevel));
-			stream.push({ type: "start", partial: output });
+			signal?.throwIfAborted();
+			const pin = await reconcileOllamaResidency(model, headers, signal);
+			signal?.throwIfAborted();
+			const iterator = streamOllamaChat(model.baseUrl, buildRequest(model, context, options, thinkingLevel, pin), {
+				headers,
+				signal,
+				fetch: options?.fetch,
+			});
 			let active: TextContent | null = null;
 			let activeIdx = -1;
 			let activeThinking: ThinkingContent | null = null;
@@ -471,17 +568,40 @@ function runStream(
 				}
 			};
 			let recordedOwnership = false;
-			for await (const chunk of iterator) {
-				const response = chunk as ChatResponse;
-				if (!recordedOwnership && model.baseUrl) {
-					const targetKey = residencyTargetKey("ollama-native", model.baseUrl);
+			let started = false;
+			for await (const response of iterator) {
+				if (!started) {
+					stream.push({ type: "start", partial: output });
+					started = true;
+				}
+				if (pin && !recordedOwnership && model.baseUrl) {
+					const targetKey = residencyTargetKey("ollama", model.baseUrl);
 					let owned = ownedModelsByTarget.get(targetKey);
 					if (!owned) {
 						owned = new Set<string>();
 						ownedModelsByTarget.set(targetKey, owned);
 					}
-					owned.add(response.model || model.id);
+					const loadedId = response.model || model.id;
+					const firstRecord = !owned.has(loadedId);
+					for (const id of ollamaModelIds(model.id, loadedId)) {
+						owned.add(id);
+						markClioLoaded(targetKey, id);
+					}
+					ownedEndpointsByTarget.set(targetKey, { baseUrl: model.baseUrl, headers });
 					recordedOwnership = true;
+					// In a dispatched worker this hands the load to the orchestrator,
+					// which owns its release (#379). Elsewhere no sink is installed.
+					if (firstRecord) {
+						reportClioModelLoad(
+							{
+								runtimeId: "ollama",
+								targetId: ollamaTargetId(model),
+								modelId: model.id,
+								aliasIds: loadedId === model.id ? [] : [loadedId],
+							},
+							targetKey,
+						);
+					}
 				}
 				const msg = response.message;
 				if (msg?.thinking && msg.thinking.length > 0) {
@@ -514,17 +634,15 @@ function runStream(
 					doneReason = response.done_reason;
 				}
 			}
-			if (aborted) throw new Error("Request was aborted");
+			signal?.throwIfAborted();
 			output.stopReason = mapStopReason(doneReason, hadToolCall);
 			stream.push({ type: "done", reason: asDoneReason(output.stopReason), message: output });
 			stream.end();
 		} catch (err) {
-			output.stopReason = aborted ? "aborted" : "error";
+			output.stopReason = signal?.aborted ? "aborted" : "error";
 			output.errorMessage = err instanceof Error ? err.message : String(err);
 			stream.push({ type: "error", reason: output.stopReason, error: output });
 			stream.end();
-		} finally {
-			if (signal) signal.removeEventListener("abort", onAbort);
 		}
 	})();
 	return stream;
@@ -538,7 +656,8 @@ function stripReasoning(options: SimpleStreamOptions | undefined): StreamOptions
 
 export const ollamaNativeApiProvider: EngineApiProvider<"ollama-native"> = {
 	api: "ollama-native",
-	stream: (model, context, options) => runStream(model, context, options, fallbackThinkingLevel(model)),
+	stream: (model, context, options) =>
+		runStream(model, resolvedRequestContext(context), options, fallbackThinkingLevel(model)),
 	streamSimple: (model, context, options?: SimpleStreamOptions) =>
-		runStream(model, context, stripReasoning(options), thinkingLevelFromSimple(options)),
+		runStream(model, resolvedRequestContext(context), stripReasoning(options), thinkingLevelFromSimple(options)),
 };