pi-omlx-picker 0.3.0 → 0.3.1

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "pi-omlx-picker",
-	"version": "0.3.0",
+	"version": "0.3.1",
 	"type": "module",
 	"description": "Pi extension that discovers models from a local OMLX server and registers them as a native Pi provider.",
 	"license": "MIT",

package/src/stream-events.ts

@@ -27,19 +27,6 @@ export function isThinkingEvent(event: AssistantMessageEvent): boolean {
 	);
 }
 
-export function mergeAbortSignals(
-	parent: AbortSignal | undefined,
-	child: AbortSignal,
-): AbortSignal {
-	if (!parent) return child;
-	if (parent.aborted) return parent;
-	const controller = new AbortController();
-	const abort = () => controller.abort();
-	parent.addEventListener("abort", abort, { once: true });
-	child.addEventListener("abort", abort, { once: true });
-	return controller.signal;
-}
-
 export function eventPartial(
 	event: AssistantMessageEvent,
 	model: Model<Api>,

package/src/stream.ts

@@ -7,18 +7,31 @@ import {
 	type Model,
 	type SimpleStreamOptions,
 } from "@earendil-works/pi-ai";
-import { streamSimple as streamSimpleOpenAICompletions } from "@earendil-works/pi-ai/compat";
+// Resolve the concrete OpenAI Completions stream once, via the lazy API factory
+// re-exported from compat. Calling the compat `streamSimple` dispatcher from
+// inside this wrapper would re-resolve through the api-provider registry; and
+// because this extension registers itself as the openai-completions handler,
+// that routes dispatch -> wrapper -> dispatch -> ... and overflows the stack.
+// The lazy factory returns a closure over the concrete implementation that
+// loads the module on first call and calls it directly — no registry, no
+// re-dispatch. It is captured here at module load, BEFORE the wrapper is
+// registered, so it can never be the wrapper itself.
+//
+// Note: pi's extension loader (jiti) only aliases a fixed set of pi-ai
+// subpaths (root, /compat, /oauth). Importing `@earendil-works/pi-ai/api/...`
+// is not resolvable there, so the concrete module is reached through compat.
+import { openAICompletionsApi } from "@earendil-works/pi-ai/compat";
+import { PROVIDER_KEY } from "./auth-storage.ts";
 import { normalizeErrorEvent } from "./overflow.ts";
 import { isRepeatStop } from "./repeat-stop.ts";
-import {
-	isMeaningfulBodyEvent,
-	isThinkingEvent,
-	mergeAbortSignals,
-} from "./stream-events.ts";
+import { isMeaningfulBodyEvent, isThinkingEvent } from "./stream-events.ts";
 import { StreamWriter } from "./stream-writer.ts";
 
+const streamOpenAICompletionsImpl = openAICompletionsApi().streamSimple;
+
 const DEFAULT_FIRST_DELTA_TIMEOUT_MS = 120_000;
 const FIRST_DELTA_MAX_ATTEMPTS = 2;
+const MAX_REISSUES = 1;
 
 export type StreamTimeoutEvent = {
 	model: string;
@@ -45,6 +58,32 @@ export function resolveFirstDeltaTimeoutMs(): number {
 		: DEFAULT_FIRST_DELTA_TIMEOUT_MS;
 }
 
+// Merge the parent (caller) signal with our own timeout signal. We compose the
+// raw source signals directly via AbortSignal.any rather than chaining through
+// a freshly-created controller per call: chaining previously-merged signals
+// accumulates abort listeners across a long session and, when abort fires,
+// propagates through N recursive .abort() calls that overflow the stack.
+// AbortSignal.any keeps the merged signal detached from either source's
+// listener set, and returns an already-aborted signal if either input is
+// aborted (so a pre-aborted parent propagates immediately). Node >=22
+// guarantees AbortSignal.any is available (engines).
+function mergeTimeoutSignal(
+	parent: AbortSignal | undefined,
+	own: AbortSignal,
+): AbortSignal {
+	if (!parent) return own;
+	return AbortSignal.any([parent, own]) as AbortSignal;
+}
+
+// Always flush buffered thinking events before leaving runAttempt, including
+// the timed-out and thrown paths. Previously these were dropped silently.
+function flushThinking(
+	writer: StreamWriter,
+	events: AssistantMessageEvent[],
+): void {
+	for (const held of events) writer.push(held);
+}
+
 async function runAttempt(
 	writer: StreamWriter,
 	model: Model<Api>,
@@ -56,11 +95,14 @@ async function runAttempt(
 	allowReissue: boolean,
 ): Promise<AttemptResult> {
 	const controller = new AbortController();
-	const signal = mergeAbortSignals(options?.signal, controller.signal);
+	const signal = mergeTimeoutSignal(options?.signal, controller.signal);
 	let timedOut = false;
 	let firstMeaningfulEvent = false;
 
 	const timer = setTimeout(() => {
+		// clearTimeout below does not stop a callback already queued on the
+		// event loop; firstMeaningfulEvent is a load-bearing guard against the
+		// timer firing just after the first body event cleared it.
 		if (writer.closed || firstMeaningfulEvent) return;
 		timedOut = true;
 		onTimeout?.({
@@ -76,7 +118,9 @@ async function runAttempt(
 	let bufferedThinking: AssistantMessageEvent[] = [];
 
 	try {
-		const inner = streamSimpleOpenAICompletions(
+		// Direct call into the concrete implementation: no dispatch, so a model
+		// whose provider registered this wrapper cannot recurse back into it.
+		const inner = streamOpenAICompletionsImpl(
 			model as Model<"openai-completions">,
 			context,
 			{ ...options, signal },
@@ -96,24 +140,26 @@ async function runAttempt(
 				continue;
 			}
 			if (allowReissue && isRepeatStop(event, context)) {
+				// Tear down the inner stream's network connection immediately so it
+				// does not drain unpredictably while the caller reissues.
+				controller.abort();
 				bufferedThinking = [];
 				return "reissue";
 			}
-			for (const held of bufferedThinking) writer.push(held);
+			flushThinking(writer, bufferedThinking);
 			bufferedThinking = [];
 			writer.push(normalizeErrorEvent(event));
 			if (event.type === "done" || event.type === "error") break;
 		}
 	} catch (err) {
+		flushThinking(writer, bufferedThinking);
 		if (timedOut) return "timed-out";
 		throw err;
 	} finally {
 		clearTimeout(timer);
 	}
 
-	if (!timedOut) {
-		for (const held of bufferedThinking) writer.push(held);
-	}
+	flushThinking(writer, bufferedThinking);
 
 	return timedOut ? "timed-out" : "completed";
 }
@@ -125,13 +171,32 @@ export function streamOmlxOpenAICompletions(
 	firstDeltaTimeoutMs: number,
 	onTimeout: OnStreamTimeout | undefined,
 ): AssistantMessageEventStream {
+	// Pi dispatches stream handlers by api id, not provider. Registering this
+	// wrapper as the openai-completions streamSimple handler (the mechanism pi
+	// exposes) replaces the shared entry, so non-oMLX OpenAI-compatible models
+	// (groq, zai, glm, ...) also arrive here. Pass them straight through to the
+	// concrete implementation with no oMLX-specific timeout/reissue logic, so the
+	// extension cannot perturb unrelated providers.
+	if (model.provider !== PROVIDER_KEY) {
+		return streamOpenAICompletionsImpl(
+			model as Model<"openai-completions">,
+			context,
+			options,
+		);
+	}
+
 	const stream = createAssistantMessageEventStream();
 	const writer = new StreamWriter(stream, model);
 
 	(async () => {
 		try {
-			let reissued = false;
-			for (let attempt = 1; attempt <= FIRST_DELTA_MAX_ATTEMPTS; attempt++) {
+			// Separate budgets so a reissue never consumes a timeout attempt and
+			// the loop index is never mutated: timeoutAttemptsLeft governs how
+			// many timed-out retries remain, reissueLeft governs reissues.
+			let timeoutAttemptsLeft = FIRST_DELTA_MAX_ATTEMPTS;
+			let reissueLeft = MAX_REISSUES;
+			while (true) {
+				const attempt = FIRST_DELTA_MAX_ATTEMPTS - timeoutAttemptsLeft + 1;
 				const result = await runAttempt(
 					writer,
 					model,
@@ -140,20 +205,22 @@ export function streamOmlxOpenAICompletions(
 					firstDeltaTimeoutMs,
 					attempt,
 					onTimeout,
-					!reissued,
+					reissueLeft > 0,
 				);
 				if (writer.closed) return;
 
 				if (result === "reissue") {
-					reissued = true;
-					attempt--; // a re-issue doesn't consume a timeout attempt
+					reissueLeft--;
 					continue;
 				}
 				if (result === "completed") {
 					writer.end();
 					return;
 				}
-				if (attempt >= FIRST_DELTA_MAX_ATTEMPTS) {
+
+				// timed-out
+				timeoutAttemptsLeft--;
+				if (timeoutAttemptsLeft <= 0) {
 					writer.pushError(
 						firstDeltaTimeoutMessage(
 							firstDeltaTimeoutMs,