@iowarp/clio-coder 0.4.9 → 0.5.0

package/src/engine/apis/openai-completions.ts

@@ -1,4 +1,5 @@
 import { TextDecoder } from "node:util";
+import type { TranscriptContext } from "@earendil-works/pi-ai";
 import {
 	type Api,
 	type AssistantMessage,
@@ -29,22 +30,31 @@ import {
 	reasoningClassForMechanism,
 	resolveModelRuntimeCapabilitiesForModel,
 } from "../../domains/providers/model-runtime-capabilities.js";
-import { lmStudioReasoningEffort } from "../../domains/providers/runtimes/common/lmstudio-http.js";
+import {
+	invalidateLmStudioCatalog,
+	lmStudioReasoningEffort,
+} from "../../domains/providers/runtimes/common/lmstudio-http.js";
 import type { ThinkingLevel } from "../../domains/providers/types/capability-flags.js";
-import type { LocalModelQuirks, SamplingProfile } from "../../domains/providers/types/local-model-quirks.js";
+import type { LocalModelQuirks } from "../../domains/providers/types/local-model-quirks.js";
 import type { LiteLLMTargetSettings, LmStudioTargetSettings } from "../../domains/providers/types/target-descriptor.js";
+import { normalizeContext, resolvedRequestContext } from "../context.js";
 import { filterGemmaChannelStream, usesGemmaChannelMarkers } from "../gemma-channel-filter.js";
 import { HarmonyResponseParser } from "../harmony-response.js";
 import { captureErrorBody, restoreTruncatedErrorBody } from "../provider-error-body.js";
 import { createSentinelStripper, stripTokenizerSentinels } from "../strip-tokenizer-sentinels.js";
-import { ensureLlamaCppResidency } from "./llamacpp-residency.js";
-import { ensureLmStudioResidency } from "./lmstudio.js";
+import { createDegradedInferenceStream, type WatchDegradedInferenceOptions } from "./degraded-inference.js";
+import { ensureLlamaCppResidency, listLlamaCppResidentModels } from "./llamacpp-residency.js";
+import { ensureLmStudioResidency, listLmStudioResidentModels } from "./lmstudio.js";
 import { remainingContextMaxTokens } from "./output-budget.js";
 import { residencyManagedFor } from "./residency.js";
-import { mergeSamplingOverride } from "./sampling-overrides.js";
+import { pickSamplingProfile, samplingParamsFromProfile } from "./sampling-overrides.js";
 import type { EngineApiProvider } from "./types.js";
 
 declare module "@earendil-works/pi-ai" {
+	interface Usage {
+		/** Whether the wire response explicitly supplied a cache-read count. */
+		cacheReadReported?: boolean;
+	}
 	interface AssistantMessage {
 		/** Direct observation of model-id presence in an OpenAI-compatible response. */
 		responseModelIdObservation?: ResponseModelIdObservation;
@@ -91,15 +101,6 @@ function clioQuirks(model: Model<"openai-completions">): LocalModelQuirks | unde
 	return (model as Model<"openai-completions"> & ClioRuntimeMetadata).clioCoder?.quirks;
 }
 
-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 isPlainRecord(value: unknown): value is Record<string, unknown> {
 	return value !== null && typeof value === "object" && !Array.isArray(value);
 }
@@ -108,6 +109,7 @@ interface ResponseModelIdCapture {
 	reportedModelId: string | null;
 	observed: boolean;
 	modelIdDone: boolean;
+	cacheReadReported: boolean;
 	backendTimings: BackendCompletionTimings | null;
 	backendTimingsSource: BackendTimingsSource | null;
 	gatewayRouting: GatewayRoutingObservation | null;
@@ -156,10 +158,6 @@ function backendTimingsSourceForModel(model: Model<Api>): BackendTimingsSource |
 	return null;
 }
 
-function captureCanStopEarly(capture: ResponseModelIdCapture): boolean {
-	return capture.modelIdDone && capture.backendTimingsSource === null;
-}
-
 function observeResponseMetadataLine(line: string, capture: ResponseModelIdCapture): void {
 	const normalized = line.endsWith("\r") ? line.slice(0, -1) : line;
 	if (!normalized.startsWith("data:")) return;
@@ -168,6 +166,14 @@ function observeResponseMetadataLine(line: string, capture: ResponseModelIdCaptu
 	try {
 		const payload = JSON.parse(data) as unknown;
 		if (!isPlainRecord(payload)) return;
+		const usage = isPlainRecord(payload.usage) ? payload.usage : {};
+		const details = isPlainRecord(usage.prompt_tokens_details) ? usage.prompt_tokens_details : {};
+		if (
+			nonnegativeFiniteNumber(details.cached_tokens) !== null ||
+			nonnegativeFiniteNumber(usage.cache_read_input_tokens) !== null
+		) {
+			capture.cacheReadReported = true;
+		}
 		if (!capture.modelIdDone) {
 			const model = payload.model;
 			if (typeof model === "string" && model.trim().length > 0) {
@@ -197,18 +203,13 @@ function observeResponseModelIdBytes(
 		const line = capture.buffer.slice(0, newline);
 		capture.buffer = capture.buffer.slice(newline + 1);
 		observeResponseMetadataLine(line, capture);
-		if (captureCanStopEarly(capture)) {
-			capture.buffer = "";
-			capture.decoder = null;
-			return;
-		}
 		newline = capture.buffer.indexOf("\n");
 	}
 	if (flush && capture.buffer.length > 0) {
 		observeResponseMetadataLine(capture.buffer, capture);
 		capture.buffer = "";
 	}
-	if (captureCanStopEarly(capture) || flush) capture.decoder = null;
+	if (flush) capture.decoder = null;
 }
 
 function captureResponseModelId(response: Response, capture: ResponseModelIdCapture, model: Model<Api>): Response {
@@ -248,6 +249,7 @@ function withResponseModelIdCapture<TOptions extends StreamOptions>(
 		reportedModelId: null,
 		observed: false,
 		modelIdDone: false,
+		cacheReadReported: false,
 		backendTimings: null,
 		backendTimingsSource: backendTimingsSourceForModel(model),
 		gatewayRouting: null,
@@ -272,6 +274,7 @@ function withResponseModelIdCapture<TOptions extends StreamOptions>(
 						: { state: "reported", reportedModelId: capture.reportedModelId }
 					: { state: "not-observed" };
 				if (event.type === "done") {
+					event.message.usage.cacheReadReported = capture.cacheReadReported;
 					event.message.responseModelIdObservation = observation;
 					if (capture.gatewayRouting !== null) event.message.gatewayRouting = capture.gatewayRouting;
 					if (capture.backendTimings !== null) event.message.backendTimings = capture.backendTimings;
@@ -279,6 +282,7 @@ function withResponseModelIdCapture<TOptions extends StreamOptions>(
 					if (event.error.errorMessage !== undefined && capture.errorBody !== null) {
 						event.error.errorMessage = restoreTruncatedErrorBody(event.error.errorMessage, await capture.errorBody);
 					}
+					event.error.usage.cacheReadReported = capture.cacheReadReported;
 					event.error.responseModelIdObservation = observation;
 					if (capture.gatewayRouting !== null) event.error.gatewayRouting = capture.gatewayRouting;
 					if (capture.backendTimings !== null) event.error.backendTimings = capture.backendTimings;
@@ -287,8 +291,7 @@ function withResponseModelIdCapture<TOptions extends StreamOptions>(
 			}
 			annotated.end();
 		} catch (err) {
-			const message = err instanceof Error ? err.message : String(err);
-			throw new Error(message);
+			failStream(annotated, model, err);
 		}
 	})();
 	return annotated;
@@ -302,46 +305,38 @@ function withLiteLLMRouteFailureAdvice(
 	const metadata = runtimeMetadata(model);
 	if (metadata?.runtimeId !== "litellm") return source;
 	const advised = createAssistantMessageEventStream();
+	let hasPartialResponse = false;
 	(async () => {
 		try {
 			for await (const event of source) {
+				if (event.type === "text_delta" || event.type === "thinking_delta" || event.type === "toolcall_delta") {
+					hasPartialResponse = true;
+				}
 				if (event.type === "error" && event.reason !== "aborted" && event.error.stopReason !== "aborted") {
-					event.error.errorMessage = liteLLMRouteFailureMessage(event.error.errorMessage, metadata.targetId, model.id);
+					event.error.errorMessage = liteLLMRouteFailureMessage(
+						event.error.errorMessage,
+						metadata.targetId,
+						model.id,
+						hasPartialResponse || event.error.content.length > 0,
+					);
 				}
 				advised.push(event as AssistantMessageEvent);
 			}
 			advised.end();
 		} catch (err) {
-			if (err instanceof Error && err.name === "AbortError") throw err;
 			const message = err instanceof Error ? err.message : String(err);
-			throw new Error(liteLLMRouteFailureMessage(message, metadata.targetId, model.id));
+			failStream(
+				advised,
+				model,
+				err instanceof Error && err.name === "AbortError"
+					? err
+					: new Error(liteLLMRouteFailureMessage(message, metadata.targetId, model.id, hasPartialResponse)),
+			);
 		}
 	})();
 	return advised;
 }
 
-function samplingParamsFromProfile(profile: SamplingProfile): Record<string, unknown> {
-	return {
-		...(profile.topP !== undefined ? { top_p: profile.topP } : {}),
-		...(profile.topK !== undefined ? { top_k: profile.topK } : {}),
-		...(profile.minP !== undefined ? { min_p: profile.minP } : {}),
-		...(profile.repeatPenalty !== undefined ? { repeat_penalty: profile.repeatPenalty } : {}),
-		...(profile.presencePenalty !== undefined ? { presence_penalty: profile.presencePenalty } : {}),
-		...(profile.frequencyPenalty !== undefined ? { frequency_penalty: profile.frequencyPenalty } : {}),
-	};
-}
-
-function piThinkingBudgets(quirks: LocalModelQuirks | undefined): ThinkingBudgets | undefined {
-	const budgets = quirks?.thinking?.budgetByLevel;
-	if (!budgets) return undefined;
-	return {
-		...(budgets.minimal !== undefined ? { minimal: budgets.minimal } : {}),
-		...(budgets.low !== undefined ? { low: budgets.low } : {}),
-		...(budgets.medium !== undefined ? { medium: budgets.medium } : {}),
-		...(budgets.high !== undefined ? { high: budgets.high } : {}),
-	};
-}
-
 type AnyOnPayload = (payload: unknown, model: Model<Api>) => unknown | undefined | Promise<unknown | undefined>;
 type StreamOptionsWithThinkingBudgets = StreamOptions & { thinkingBudgets?: ThinkingBudgets };
 
@@ -460,23 +455,14 @@ function applyThinkingPayload(
 	resolved: ResolvedModelRuntimeCapabilities,
 	model: Model<Api>,
 ): Record<string, unknown> {
-	if (applied.mechanism === "none") {
-		const next = stripThinkingRequestFields(payload);
+	if (applied.mechanism === "none" || applied.mechanism === "always-on") {
+		const next = applied.mechanism === "none" ? stripThinkingRequestFields(payload) : { ...payload };
 		if (resolved.request.chatTemplateKwargs && !chatTemplateKwargsUnsupported(model)) {
 			const existing = isPlainRecord(next.chat_template_kwargs) ? next.chat_template_kwargs : {};
 			next.chat_template_kwargs = { ...existing, ...resolved.request.chatTemplateKwargs };
 		}
 		return next;
 	}
-	if (applied.mechanism === "always-on") {
-		if (resolved.request.chatTemplateKwargs && !chatTemplateKwargsUnsupported(model)) {
-			const next: Record<string, unknown> = { ...payload };
-			const existing = isPlainRecord(next.chat_template_kwargs) ? next.chat_template_kwargs : {};
-			next.chat_template_kwargs = { ...existing, ...resolved.request.chatTemplateKwargs };
-			return next;
-		}
-		return payload;
-	}
 	const next: Record<string, unknown> = { ...payload };
 	if (
 		resolved.request.reasoningEffort &&
@@ -580,11 +566,8 @@ function composeThinkingOnPayload(
 		if (!isPlainRecord(payload)) {
 			return base ? await base(payload, model) : undefined;
 		}
-		const next = applyLmStudioPayload(
-			applyThinkingPayload(applyLlamaCppPromptCachePayload(payload, model, retention), resolved.thinking, resolved, model),
-			model,
-			resolved,
-		);
+		const cached = applyLlamaCppPromptCachePayload(payload, model, retention);
+		const next = applyLmStudioPayload(applyThinkingPayload(cached, resolved.thinking, resolved, model), model, resolved);
 		if (base) {
 			const fromBase = await base(next, model);
 			if (fromBase !== undefined) return fromBase;
@@ -603,24 +586,16 @@ function withSamplingOverrides<TOptions extends StreamOptions>(
 	const profile = pickSamplingProfile(quirks, applied.thinkingActive);
 	const promptCache = shouldApplyLlamaCppPromptCache(model);
 	const lmstudio = isLmStudioModel(model);
-	const vllmThinkingBudgets = resolved.runtimeId === "vllm" ? piThinkingBudgets(quirks) : undefined;
-	if (
-		!promptCache &&
-		!lmstudio &&
-		!profile &&
-		!vllmThinkingBudgets &&
-		applied.mechanism !== "effort-levels" &&
-		applied.mechanism !== "budget-tokens" &&
-		applied.mechanism !== "on-off" &&
-		applied.mechanism !== "none"
-	) {
+	const vllmThinkingBudgets = resolved.runtimeId === "vllm" ? quirks?.thinking?.budgetByLevel : undefined;
+	const needsPayloadControls = promptCache || lmstudio || applied.mechanism !== "always-on";
+	if (!profile && !vllmThinkingBudgets && !needsPayloadControls) {
 		return options;
 	}
 	const merged: Record<string, unknown> = { ...(options ?? {}) };
 	if (profile?.temperature !== undefined && merged.temperature === undefined) merged.temperature = profile.temperature;
 	if (profile) {
 		merged.samplingParams = {
-			...samplingParamsFromProfile(profile),
+			...samplingParamsFromProfile(profile, resolved.runtimeId),
 			...options?.samplingParams,
 		};
 	}
@@ -630,14 +605,7 @@ function withSamplingOverrides<TOptions extends StreamOptions>(
 			...(options as StreamOptionsWithThinkingBudgets | undefined)?.thinkingBudgets,
 		};
 	}
-	if (
-		promptCache ||
-		lmstudio ||
-		applied.mechanism === "effort-levels" ||
-		applied.mechanism === "budget-tokens" ||
-		applied.mechanism === "on-off" ||
-		applied.mechanism === "none"
-	) {
+	if (needsPayloadControls) {
 		merged.onPayload = composeThinkingOnPayload(resolved, options?.onPayload, options?.cacheRetention);
 	}
 	return merged as TOptions;
@@ -645,6 +613,7 @@ function withSamplingOverrides<TOptions extends StreamOptions>(
 
 function stripNeverReasoningFromStream(
 	source: AssistantMessageEventStream,
+	model: Model<Api>,
 	resolved: ResolvedModelRuntimeCapabilities,
 ): AssistantMessageEventStream {
 	if (!stripsThinking(resolved)) return source;
@@ -665,8 +634,7 @@ function stripNeverReasoningFromStream(
 			}
 			stripped.end();
 		} catch (err) {
-			const message = err instanceof Error ? err.message : String(err);
-			throw new Error(message);
+			failStream(stripped, model, err);
 		}
 	})();
 	return stripped;
@@ -749,7 +717,7 @@ function runtimeMetadata(model: Model<Api>): NonNullable<ClioRuntimeMetadata["cl
 	return (model as Model<Api> & ClioRuntimeMetadata).clioCoder;
 }
 
-function emptyErrorMessage(model: Model<"openai-completions">, message: string): AssistantMessage {
+function emptyErrorMessage(model: Model<Api>, message: string): AssistantMessage {
 	return {
 		role: "assistant",
 		content: [],
@@ -770,6 +738,13 @@ function emptyErrorMessage(model: Model<"openai-completions">, message: string):
 	};
 }
 
+function failStream(stream: AssistantMessageEventStream, model: Model<Api>, cause: unknown): void {
+	const error = emptyErrorMessage(model, cause instanceof Error ? cause.message : String(cause));
+	error.stopReason = cause instanceof Error && cause.name === "AbortError" ? "aborted" : "error";
+	stream.push({ type: "error", reason: error.stopReason, error });
+	stream.end(error);
+}
+
 function malformedToolArgsMessage(
 	model: Model<"openai-completions">,
 	toolName: string,
@@ -844,7 +819,10 @@ function applyOpenAICompatReasoningEstimate(message: AssistantMessage): void {
 	}
 }
 
-function withReasoningTokenEstimate(source: AssistantMessageEventStream): AssistantMessageEventStream {
+function withReasoningTokenEstimate(
+	source: AssistantMessageEventStream,
+	model: Model<Api>,
+): AssistantMessageEventStream {
 	const annotated = createAssistantMessageEventStream();
 	(async () => {
 		try {
@@ -858,8 +836,7 @@ function withReasoningTokenEstimate(source: AssistantMessageEventStream): Assist
 			}
 			annotated.end();
 		} catch (err) {
-			const message = err instanceof Error ? err.message : String(err);
-			throw new Error(message);
+			failStream(annotated, model, err);
 		}
 	})();
 	return annotated;
@@ -875,6 +852,7 @@ function withReasoningTokenEstimate(source: AssistantMessageEventStream): Assist
  */
 function stripSentinelsFromStream(
 	source: AssistantMessageEventStream,
+	model: Model<Api>,
 	resolved: ResolvedModelRuntimeCapabilities,
 ): AssistantMessageEventStream {
 	const sanitized = createAssistantMessageEventStream();
@@ -954,8 +932,7 @@ function stripSentinelsFromStream(
 			}
 			sanitized.end();
 		} catch (err) {
-			const message = err instanceof Error ? err.message : String(err);
-			throw new Error(message);
+			failStream(sanitized, model, err);
 		}
 	})();
 	return sanitized;
@@ -1017,7 +994,22 @@ function isManagedLlamaCppModel(model: Model<"openai-completions">): boolean {
 	return model.provider === "llamacpp" && metadata?.runtimeId === "llamacpp" && typeof model.baseUrl === "string";
 }
 
-async function ensureResidencyForModel(model: Model<"openai-completions">): Promise<void> {
+type LocalRequestOptions = Pick<StreamOptions, "apiKey" | "headers" | "signal">;
+
+function localRequestHeaders(model: Model<Api>, options: LocalRequestOptions): Record<string, string> {
+	const headers = new Headers(model.headers);
+	for (const [name, value] of Object.entries(options.headers ?? {})) {
+		if (value === null) headers.delete(name);
+		else if (value !== undefined) headers.set(name, value);
+	}
+	if (options.apiKey && !headers.has("authorization")) headers.set("authorization", `Bearer ${options.apiKey}`);
+	return Object.fromEntries(headers);
+}
+
+async function ensureResidencyForModel(
+	model: Model<"openai-completions">,
+	options: LocalRequestOptions,
+): Promise<void> {
 	const metadata = runtimeMetadata(model);
 	if (model.provider !== "llamacpp" || metadata?.runtimeId !== "llamacpp") return;
 	if (typeof model.baseUrl !== "string" || model.baseUrl.length === 0) return;
@@ -1027,133 +1019,143 @@ async function ensureResidencyForModel(model: Model<"openai-completions">): Prom
 		runtimeId: metadata.runtimeId,
 		keepModelId: model.id,
 		managed: residencyManagedFor(metadata.lifecycle),
+		headers: localRequestHeaders(model, options),
+		...(options.signal ? { signal: options.signal } : {}),
 	});
 }
 
 async function ensureLocalResidency(
 	model: Model<"openai-completions">,
-	options: { apiKey?: string; signal?: AbortSignal },
+	options: LocalRequestOptions,
 ): Promise<Model<"openai-completions">> {
 	if (isLmStudioModel(model)) {
-		const wireModelId = await ensureLmStudioResidency(model, options);
+		const requestModel = { ...model, headers: localRequestHeaders(model, options) };
+		const wireModelId = await ensureLmStudioResidency(requestModel, options);
+		// Residency may discover a smaller loaded window after turn preflight.
+		if (model.contextWindow <= 0 || requestModel.contextWindow < model.contextWindow) {
+			model.contextWindow = requestModel.contextWindow;
+		}
 		return wireModelId === model.id ? model : { ...model, id: wireModelId };
 	}
-	await ensureResidencyForModel(model);
+	await ensureResidencyForModel(model, options);
 	return model;
 }
 
-function withLocalResidency(
+function degradedWatchOptions(
+	model: Model<"openai-completions">,
+	options: LocalRequestOptions,
+): WatchDegradedInferenceOptions {
+	const metadata = runtimeMetadata(model);
+	const baseUrl = model.baseUrl;
+	const listResident = isLmStudioModel(model)
+		? () => listLmStudioResidentModels({ ...model, headers: localRequestHeaders(model, options) }, options)
+		: () =>
+				listLlamaCppResidentModels(baseUrl, fetch, {
+					headers: localRequestHeaders(model, options),
+					...(options.signal ? { signal: options.signal } : {}),
+				});
+	return {
+		targetId: metadata?.targetId ?? model.provider,
+		runtimeId: metadata?.runtimeId ?? model.provider,
+		model: model.id,
+		...(options.signal !== undefined ? { signal: options.signal } : {}),
+		listResident,
+	};
+}
+
+// Local servers answer a CPU spill with a crawl rather than an error, so the
+// streams reaching them run under the degraded-inference watchdog. The source
+// is iterated here, inside this catch, so a throw mid-stream still ends the
+// turn with an error event. Exported for the error-propagation contract.
+export function withLocalResidency(
 	model: Model<"openai-completions">,
-	options: { apiKey?: string; signal?: AbortSignal },
+	options: LocalRequestOptions,
 	sourceFactory: (requestModel: Model<"openai-completions">) => AssistantMessageEventStream,
 ): AssistantMessageEventStream {
 	if (!isManagedLlamaCppModel(model) && !isLmStudioModel(model)) return sourceFactory(model);
-	const stream = createAssistantMessageEventStream();
+	const stream = createDegradedInferenceStream(degradedWatchOptions(model, options));
 	(async () => {
 		try {
+			options.signal?.throwIfAborted();
 			const requestModel = await ensureLocalResidency(model, options);
+			options.signal?.throwIfAborted();
 			for await (const event of sourceFactory(requestModel)) {
+				if (event.type === "error" && isLmStudioModel(model)) {
+					invalidateLmStudioCatalog({
+						id: runtimeMetadata(model)?.targetId ?? model.provider,
+						runtime: "lmstudio",
+						url: model.baseUrl,
+					});
+				}
 				stream.push(event as AssistantMessageEvent);
 			}
 			stream.end();
 		} catch (err) {
 			const message = err instanceof Error ? err.message : String(err);
 			const error = emptyErrorMessage(model, message);
-			stream.push({ type: "error", reason: "error", error });
+			error.stopReason = options.signal?.aborted ? "aborted" : "error";
+			stream.push({ type: "error", reason: error.stopReason, error });
 			stream.end(error);
 		}
 	})();
 	return stream;
 }
 
+function streamCompletions<TOptions extends StreamOptions>(
+	model: Model<"openai-completions">,
+	context: Context,
+	options: TOptions | undefined,
+	level: ThinkingLevel,
+	start: (
+		model: Model<"openai-completions">,
+		context: TranscriptContext,
+		options: TOptions,
+	) => AssistantMessageEventStream,
+): AssistantMessageEventStream {
+	const resolved = resolvedCapabilitiesForModel(model, level);
+	const effectiveContext = stripsThinking(resolved) ? stripThinkingFromContext(context) : context;
+	// Pi's credential check reads request headers, while its HTTP client also
+	// reads model headers. Forward both so header-authenticated targets pass both.
+	const transportOptions = model.headers
+		? ({
+				...options,
+				headers: Object.fromEntries(
+					[...Object.entries(model.headers), ...Object.entries(options?.headers ?? {})].map(([name, value]) => [
+						name.toLowerCase(),
+						value,
+					]),
+				),
+			} as TOptions)
+		: options;
+	const requestOptions = withLiteLLMRequestOptions(model, transportOptions ?? ({} as TOptions));
+	const source = withResponseModelIdCapture(model, requestOptions, (capturedOptions) =>
+		withLocalResidency(model, options ?? {}, (requestModel) => {
+			return start(
+				requestModel,
+				normalizeContext(preserveInterruptedReasoning(effectiveContext, requestModel, resolved)),
+				withRemainingContextBudget(
+					requestModel,
+					effectiveContext,
+					withSamplingOverrides(requestModel, capturedOptions, resolved),
+				),
+			);
+		}),
+	);
+	const advised = withLiteLLMRouteFailureAdvice(model, source);
+	const channels = filterGemmaChannelStream(advised, usesGemmaChannelMarkers(resolved.modelId), (stream, error) =>
+		failStream(stream, model, error),
+	);
+	const thinking = stripNeverReasoningFromStream(channels, model, resolved);
+	const sanitized = stripSentinelsFromStream(thinking, model, resolved);
+	return guardMalformedToolCalls(withReasoningTokenEstimate(sanitized, model), model, resolvedRequestContext(context));
+}
+
 export const openAICompletionsApiProvider: EngineApiProvider<"openai-completions", OpenAICompletionsOptions> = {
 	api: "openai-completions",
 	stream: (model, context, options) => {
-		// Bare `stream` callers don't communicate thinking state; fall back to
-		// the model's reasoning capability so the catalog still applies. LM Studio
-		// is the exception because its HTTP server defaults to thinking off.
-		const defaultLevel = isLmStudioModel(model) ? "off" : model.reasoning === true ? "medium" : "off";
-		const resolved = resolvedCapabilitiesForModel(model, defaultLevel);
-		const effectiveContext = stripsThinking(resolved) ? stripThinkingFromContext(context) : context;
-		const withSamplers = withSamplingOverrides(model, options, resolved);
-		const withGatewayOptions = withLiteLLMRequestOptions(
-			model,
-			withRemainingContextBudget(model, effectiveContext, withSamplers),
-		);
-		return guardMalformedToolCalls(
-			withReasoningTokenEstimate(
-				stripSentinelsFromStream(
-					stripNeverReasoningFromStream(
-						filterGemmaChannelStream(
-							withLiteLLMRouteFailureAdvice(
-								model,
-								withResponseModelIdCapture(model, withGatewayOptions, (capturedOptions) =>
-									withLocalResidency(
-										model,
-										{
-											...(options?.apiKey !== undefined ? { apiKey: options.apiKey } : {}),
-											...(options?.signal !== undefined ? { signal: options.signal } : {}),
-										},
-										(requestModel) =>
-											piOpenAICompletions.stream(
-												requestModel,
-												preserveInterruptedReasoning(effectiveContext, requestModel, resolved),
-												capturedOptions,
-											),
-									),
-								),
-							),
-							usesGemmaChannelMarkers(resolved.modelId),
-						),
-						resolved,
-					),
-					resolved,
-				),
-			),
-			model,
-			context,
-		);
-	},
-	streamSimple: (model, context, options?: SimpleStreamOptions) => {
-		const resolved = resolvedCapabilitiesForModel(model, thinkingLevelFromSimple(options));
-		const effectiveContext = stripsThinking(resolved) ? stripThinkingFromContext(context) : context;
-		const withSamplers = withSamplingOverrides(model, options, resolved);
-		const withGatewayOptions = withLiteLLMRequestOptions(
-			model,
-			withRemainingContextBudget(model, effectiveContext, withSamplers),
-		);
-		return guardMalformedToolCalls(
-			withReasoningTokenEstimate(
-				stripSentinelsFromStream(
-					stripNeverReasoningFromStream(
-						filterGemmaChannelStream(
-							withLiteLLMRouteFailureAdvice(
-								model,
-								withResponseModelIdCapture(model, withGatewayOptions, (capturedOptions) =>
-									withLocalResidency(
-										model,
-										{
-											...(options?.apiKey !== undefined ? { apiKey: options.apiKey } : {}),
-											...(options?.signal !== undefined ? { signal: options.signal } : {}),
-										},
-										(requestModel) =>
-											piOpenAICompletions.streamSimple(
-												requestModel,
-												preserveInterruptedReasoning(effectiveContext, requestModel, resolved),
-												capturedOptions,
-											),
-									),
-								),
-							),
-							usesGemmaChannelMarkers(resolved.modelId),
-						),
-						resolved,
-					),
-					resolved,
-				),
-			),
-			model,
-			context,
-		);
+		const level = isLmStudioModel(model) ? "off" : model.reasoning ? "medium" : "off";
+		return streamCompletions(model, context, options, level, piOpenAICompletions.stream);
 	},
+	streamSimple: (model, context, options?: SimpleStreamOptions) =>
+		streamCompletions(model, context, options, thinkingLevelFromSimple(options), piOpenAICompletions.streamSimple),
 };

package/src/engine/apis/output-budget.ts

@@ -1,6 +1,7 @@
 import type { Api, Context, Model, StreamOptions } from "@earendil-works/pi-ai";
 import { CLIO_MIN_MAX_OUTPUT_TOKENS } from "../../core/context-floor.js";
 import { ceilChars, estimateAgentMessageTokens, toolSchemaChars } from "../../domains/session/context-accounting.js";
+import { resolvedRequestContext } from "../context.js";
 
 const CONTEXT_BUDGET_SAFETY_TOKENS = 1024;
 /**
@@ -66,7 +67,8 @@ function clampOutputToRemainingContext(ceiling: number, contextWindow: number, i
 	return Math.min(ceiling, available);
 }
 
-export function estimateInputTokensFromContext(context: Context): number {
+export function estimateInputTokensFromContext(input: Context): number {
+	const context = resolvedRequestContext(input);
 	const system = context.systemPrompt ? ceilChars(context.systemPrompt.length) : 0;
 	const messages = context.messages.reduce((sum, msg) => sum + estimateAgentMessageTokens(msg), 0);
 	const tools = (context.tools ?? []).reduce((sum, tool) => sum + ceilChars(toolSchemaChars(tool)), 0);

package/src/engine/apis/residency-lock.ts

@@ -23,9 +23,10 @@ function lockTargetFor(targetKey: string): string {
 }
 
 /** Run `fn` while holding the per-target residency lock. */
-export async function withResidencyLock<T>(targetKey: string, fn: () => Promise<T>): Promise<T> {
+export async function withResidencyLock<T>(targetKey: string, fn: () => Promise<T>, signal?: AbortSignal): Promise<T> {
 	return withStateFileLock(lockTargetFor(targetKey), fn, {
 		timeoutMs: LOCK_WAIT_MS,
 		onAcquireFailure: "run-unlocked",
+		...(signal ? { signal } : {}),
 	});
 }