@oh-my-pi/pi-ai 16.3.11 → 16.3.13

package/src/utils/idle-iterator.ts

@@ -135,6 +135,16 @@ export interface IdleTimeoutIteratorOptions {
 	 * keepalive/no-op events from keeping a stalled tool call alive forever.
 	 */
 	isProgressItem?: (item: unknown) => boolean;
+	/**
+	 * Reports consumer-side local work in flight for the stream: the provider
+	 * transport is waiting on a server-requested local tool bridge (e.g. the
+	 * Cursor exec channel) before anything can flow upstream again. While it
+	 * returns true, an expired idle / first-item deadline slides forward
+	 * instead of aborting — the silence is ours, not a provider stall. The
+	 * watchdog re-arms with a full budget once the local work completes, so a
+	 * provider that stalls afterwards is still caught.
+	 */
+	hasPendingLocalWork?: () => boolean;
 	/**
 	 * Cancel iteration as soon as this signal aborts. Required for caller-driven
 	 * cancellation (ESC) when the underlying transport does not surface signal
@@ -157,7 +167,7 @@ export async function* iterateWithIdleTimeout<T>(
 	options: IdleTimeoutIteratorOptions,
 ): AsyncGenerator<T> {
 	const firstItemTimeoutMs = options.firstItemTimeoutMs ?? options.idleTimeoutMs;
-	const firstItemDeadlineMs =
+	let firstItemDeadlineMs =
 		firstItemTimeoutMs !== undefined && firstItemTimeoutMs > 0 ? Date.now() + firstItemTimeoutMs : undefined;
 	const abortSignal = options.abortSignal;
 	const iterator = iterable[Symbol.asyncIterator]();
@@ -197,6 +207,28 @@ export async function* iterateWithIdleTimeout<T>(
 	};
 	let lastProgressAt = Date.now();
 
+	const hasPendingLocalWork = (): boolean => {
+		if (!options.hasPendingLocalWork) return false;
+		try {
+			return options.hasPendingLocalWork();
+		} catch {
+			return false;
+		}
+	};
+	// Local work means the current gap is attributable to the consumer side,
+	// not the provider: slide the active deadline a full budget past now
+	// instead of aborting. Once the work completes the watchdog resumes from
+	// the last extension, so a provider that stalls afterwards is still caught.
+	const extendDeadlineForLocalWork = (): void => {
+		if (awaitingFirstItem) {
+			if (firstItemDeadlineMs !== undefined && firstItemTimeoutMs !== undefined) {
+				firstItemDeadlineMs = Date.now() + firstItemTimeoutMs;
+			}
+		} else {
+			lastProgressAt = Date.now();
+		}
+	};
+
 	const noTimeoutEnforced =
 		(firstItemTimeoutMs === undefined || firstItemTimeoutMs <= 0) &&
 		(options.idleTimeoutMs === undefined || options.idleTimeoutMs <= 0);
@@ -271,6 +303,12 @@ export async function* iterateWithIdleTimeout<T>(
 		timer = setTimeout(onTimerFire, Math.max(0, deadlineMs - Date.now()));
 	};
 
+	// The in-flight iterator.next() promise, persisted across loop iterations:
+	// a deadline extension for pending local work loops without consuming it,
+	// and issuing a second next() while one is outstanding would drop an item.
+	let pendingNext:
+		| Promise<{ kind: "next"; result: IteratorResult<T> } | { kind: "error"; error: unknown }>
+		| undefined;
 	try {
 		let raceCount = 0;
 		while (true) {
@@ -291,21 +329,29 @@ export async function* iterateWithIdleTimeout<T>(
 				if (firstItemDeadlineMs !== undefined) {
 					activeTimeoutMs = firstItemDeadlineMs - Date.now();
 					if (activeTimeoutMs <= 0) {
-						options.onFirstItemTimeout?.();
-						closeIterator();
-						throw new AIError.StreamTimeoutError(options.firstItemErrorMessage ?? options.errorMessage);
+						if (!hasPendingLocalWork()) {
+							options.onFirstItemTimeout?.();
+							closeIterator();
+							throw new AIError.StreamTimeoutError(options.firstItemErrorMessage ?? options.errorMessage);
+						}
+						extendDeadlineForLocalWork();
+						activeTimeoutMs = firstItemDeadlineMs! - Date.now();
 					}
 				}
 			} else if (options.idleTimeoutMs !== undefined && options.idleTimeoutMs > 0) {
 				activeTimeoutMs = options.idleTimeoutMs - (Date.now() - lastProgressAt);
 				if (activeTimeoutMs <= 0) {
-					options.onIdle?.();
-					closeIterator();
-					throw new AIError.StreamTimeoutError(options.errorMessage);
+					if (!hasPendingLocalWork()) {
+						options.onIdle?.();
+						closeIterator();
+						throw new AIError.StreamTimeoutError(options.errorMessage);
+					}
+					extendDeadlineForLocalWork();
+					activeTimeoutMs = options.idleTimeoutMs;
 				}
 			}
 
-			const nextResultPromise = withRacy(iterator.next());
+			pendingNext ??= withRacy(iterator.next());
 
 			const racers: Array<
 				Promise<
@@ -314,7 +360,7 @@ export async function* iterateWithIdleTimeout<T>(
 					| { kind: "timeout" }
 					| { kind: "abort" }
 				>
-			> = [nextResultPromise];
+			> = [pendingNext];
 
 			const enforceTimeout = !noTimeoutEnforced && activeTimeoutMs !== undefined && activeTimeoutMs > 0;
 			if (enforceTimeout) {
@@ -333,11 +379,21 @@ export async function* iterateWithIdleTimeout<T>(
 			let continuing = false;
 			try {
 				const outcome = await Promise.race(racers);
+				if (outcome.kind === "next" || outcome.kind === "error") {
+					pendingNext = undefined;
+				}
 				if (outcome.kind === "abort") {
 					closeIterator();
 					throw abortReason(abortSignal!);
 				}
 				if (outcome.kind === "timeout") {
+					if (hasPendingLocalWork()) {
+						// A local tool is still running; the provider cannot make
+						// progress until we hand its result back. Keep waiting.
+						extendDeadlineForLocalWork();
+						continuing = true;
+						continue;
+					}
 					if (!awaitingFirstItem) {
 						options.onIdle?.();
 					} else {