@prometheus-ai/ai 0.5.4 → 0.5.8

package/src/usage.ts

@@ -4,8 +4,8 @@
  * Provides a normalized schema to represent multiple limit windows, model tiers,
  * and shared quotas across providers.
  */
-import * as z from "zod/v4";
-import type { Provider } from "./types";
+import { z } from "zod/v4";
+import type { FetchImpl, Provider } from "./types";
 export type UsageUnit = "percent" | "tokens" | "requests" | "usd" | "minutes" | "bytes" | "unknown";
 
 export type UsageStatus = "ok" | "warning" | "exhausted" | "unknown";
@@ -63,15 +63,78 @@ export interface UsageLimit {
 	notes?: string[];
 }
 
+/**
+ * Saved/banked rate-limit resets an account can redeem on demand.
+ *
+ * Surfaced by providers that let users defer a usage-window reset and spend it
+ * later (OpenAI Codex "saved rate limit resets"). The redeem itself is a
+ * separate, provider-specific action; this is the read-only count for display.
+ */
+export interface UsageResetCredits {
+	/** Number of resets available to redeem right now. */
+	availableCount: number;
+}
+
 /** Aggregated usage report for a provider. */
 export interface UsageReport {
 	provider: Provider;
 	fetchedAt: number;
 	limits: UsageLimit[];
+	/** Saved rate-limit resets the account can redeem, when the provider reports them. */
+	resetCredits?: UsageResetCredits;
 	metadata?: Record<string, unknown>;
 	raw?: unknown;
 }
 
+/**
+ * Resolve a limit's used fraction (0..1; >1 means overage) from whichever
+ * amount fields the provider populated. Precedence mirrors the usage UIs:
+ * explicit fraction > used/limit > percent-unit used > inverted remaining.
+ */
+export function resolveUsedFraction(limit: UsageLimit): number | undefined {
+	const amount = limit.amount;
+	if (amount.usedFraction !== undefined) return amount.usedFraction;
+	if (amount.used !== undefined && amount.limit !== undefined && amount.limit > 0) {
+		return amount.used / amount.limit;
+	}
+	if (amount.unit === "percent" && amount.used !== undefined) return amount.used / 100;
+	if (amount.remainingFraction !== undefined) return Math.max(0, 1 - amount.remainingFraction);
+	return undefined;
+}
+
+/**
+ * One recorded usage-limit snapshot: a single limit window of one account at
+ * a point in time. The usage cache itself is latest-snapshot-only; history
+ * rows are appended by the auth storage layer whenever a fresh report is
+ * fetched, so limit utilization stays inspectable over time.
+ */
+export interface UsageHistoryEntry {
+	/** Epoch ms the report was fetched. */
+	recordedAt: number;
+	provider: Provider;
+	/** Stable credential identity key (account/email/project derived). */
+	accountKey: string;
+	email?: string;
+	accountId?: string;
+	/** {@link UsageLimit.id} of the recorded window. */
+	limitId: string;
+	/** Human label of the limit. */
+	label: string;
+	windowLabel?: string;
+	/** Used fraction (0..1) when resolvable. */
+	usedFraction?: number;
+	status?: UsageStatus;
+	/** Epoch ms the window resets, when known. */
+	resetsAt?: number;
+}
+
+/** Filter for reading recorded usage history. */
+export interface UsageHistoryQuery {
+	provider?: string;
+	/** Inclusive lower bound on {@link UsageHistoryEntry.recordedAt} (epoch ms). */
+	sinceMs?: number;
+}
+
 // ─── Zod schemas (wire-shape validation for the broker `/v1/usage` endpoint) ─
 
 export const usageUnitSchema = z.enum(["percent", "tokens", "requests", "usd", "minutes", "bytes", "unknown"]);
@@ -114,10 +177,15 @@ export const usageLimitSchema = z.object({
 	notes: z.array(z.string()).optional(),
 });
 
+export const usageResetCreditsSchema = z.object({
+	availableCount: z.number(),
+});
+
 export const usageReportSchema = z.object({
 	provider: z.string(),
 	fetchedAt: z.number(),
 	limits: z.array(usageLimitSchema),
+	resetCredits: usageResetCreditsSchema.optional(),
 	metadata: z.record(z.string(), z.unknown()).optional(),
 	// `raw` is provider-specific and may be anything; the broker strips it before
 	// sending the report over the wire, so accept-but-ignore here.
@@ -154,7 +222,7 @@ export interface UsageFetchParams {
 
 /** Shared runtime utilities for fetchers. */
 export interface UsageFetchContext {
-	fetch: typeof fetch;
+	fetch: FetchImpl;
 	logger?: UsageLogger;
 	retryWait?: (delayMs: number, signal?: AbortSignal) => Promise<void>;
 }
@@ -168,13 +236,34 @@ export interface UsageProvider {
 	supports?(params: UsageFetchParams): boolean;
 }
 
+/** Request context used when ranking usage for a specific model. */
+export interface CredentialRankingContext {
+	/** Provider model id, when the caller is selecting a credential for one model. */
+	modelId?: string;
+}
+
 /** Strategy for usage-based credential ranking. Providers implement this to opt into smart credential selection. */
 export interface CredentialRankingStrategy {
 	/** Extract the primary (short) and secondary (long) window limits from a usage report. */
-	findWindowLimits(report: UsageReport): {
+	findWindowLimits(
+		report: UsageReport,
+		context?: CredentialRankingContext,
+	): {
 		primary?: UsageLimit;
 		secondary?: UsageLimit;
 	};
+	/**
+	 * Restrict limits to the ones relevant for the requested model before
+	 * credential-wide exhaustion checks and ranking. Providers with shared
+	 * account-wide quotas can omit this and use all limits.
+	 */
+	scopeLimits?(report: UsageReport, context?: CredentialRankingContext): UsageLimit[];
+	/**
+	 * Return a provider-local backoff scope for the requested model. Providers
+	 * with backend-specific quotas use this so one exhausted model family does
+	 * not block unrelated families on the same OAuth credential.
+	 */
+	blockScope?(context?: CredentialRankingContext): string | undefined;
 	/** Fallback window durations (ms) when limits don't specify durationMs. */
 	windowDefaults: {
 		primaryMs: number;

package/src/utils/abort.ts

@@ -49,3 +49,17 @@ export function createAbortSourceTracker(callerSignal?: AbortSignal): AbortSourc
 		},
 	};
 }
+
+/**
+ * Race a shared promise against a caller's AbortSignal without coupling the
+ * underlying work to that signal. The shared promise keeps running (and caches
+ * its result) even when an individual caller bails out.
+ */
+export function raceWithSignal<T>(promise: Promise<T>, signal: AbortSignal | undefined): Promise<T> {
+	if (!signal) return promise;
+	if (signal.aborted) return Promise.reject(signal.reason ?? new Error("Request was aborted"));
+	const { promise: aborted, reject } = Promise.withResolvers<never>();
+	const onAbort = () => reject(signal.reason ?? new Error("Request was aborted"));
+	signal.addEventListener("abort", onAbort, { once: true });
+	return Promise.race([promise, aborted]).finally(() => signal.removeEventListener("abort", onAbort));
+}

package/src/utils/event-stream.ts

@@ -5,6 +5,8 @@ export class EventStream<T, R = T> implements AsyncIterable<T> {
 	queue: T[] = [];
 	waiting: Array<{ resolve: (value: IteratorResult<T>) => void; reject: (err: unknown) => void }> = [];
 	done = false;
+	/** True once finalResultPromise has been resolved or rejected. */
+	resultSettled = false;
 	#failed = false;
 	#error: unknown = undefined;
 	finalResultPromise: Promise<R>;
@@ -30,6 +32,7 @@ export class EventStream<T, R = T> implements AsyncIterable<T> {
 
 		if (this.isComplete(event)) {
 			this.done = true;
+			this.resultSettled = true;
 			this.resolveFinalResult(this.extractResult(event));
 		}
 
@@ -54,7 +57,13 @@ export class EventStream<T, R = T> implements AsyncIterable<T> {
 	end(result?: R): void {
 		this.done = true;
 		if (result !== undefined) {
+			this.resultSettled = true;
 			this.resolveFinalResult(result);
+		} else if (!this.resultSettled) {
+			// end() without a terminal value must still settle result() —
+			// otherwise complete()/result() awaits hang forever.
+			this.resultSettled = true;
+			this.rejectFinalResult(new Error("Stream ended without a final result"));
 		}
 		// Notify all waiting consumers that we're done
 		while (this.waiting.length > 0) {
@@ -75,6 +84,7 @@ export class EventStream<T, R = T> implements AsyncIterable<T> {
 		this.done = true;
 		this.#failed = true;
 		this.#error = err;
+		this.resultSettled = true;
 		this.rejectFinalResult(err);
 		while (this.waiting.length > 0) {
 			const waiter = this.waiting.shift()!;
@@ -126,6 +136,7 @@ export class AssistantMessageEventStream extends EventStream<AssistantMessageEve
 		// Completion resolves the final result and still emits the terminal event.
 		if (this.isComplete(event)) {
 			this.done = true;
+			this.resultSettled = true;
 			this.resolveFinalResult(this.extractResult(event));
 		}
 
@@ -135,7 +146,13 @@ export class AssistantMessageEventStream extends EventStream<AssistantMessageEve
 	override end(result?: AssistantMessage): void {
 		this.done = true;
 		if (result !== undefined) {
+			this.resultSettled = true;
 			this.resolveFinalResult(result);
+		} else if (!this.resultSettled) {
+			// Mirror the base class: a result-less end() must not leave
+			// result() pending forever.
+			this.resultSettled = true;
+			this.rejectFinalResult(new Error("Stream ended without a final result"));
 		}
 		this.endWaiting();
 	}

package/src/utils/http-inspector.ts

@@ -1,5 +1,5 @@
 import * as path from "node:path";
-import { extractHttpStatusFromError, getLogsDir } from "@prometheus-ai/utils";
+import { extractHttpStatusFromError, getLogsDir, isBunTestRuntime } from "@prometheus-ai/utils";
 import { isCopilotTransientModelError } from "./retry.js";
 import { formatErrorMessageWithRetryAfter } from "./retry-after.js";
 
@@ -20,10 +20,6 @@ export type CapturedHttpErrorResponse = {
 	bodyJson?: unknown;
 };
 
-type ErrorWithStatus = {
-	status?: unknown;
-};
-
 const SENSITIVE_HEADERS = ["authorization", "x-api-key", "api-key", "cookie", "set-cookie", "proxy-authorization"];
 
 export async function appendRawHttpRequestDumpFor400(
@@ -31,7 +27,9 @@ export async function appendRawHttpRequestDumpFor400(
 	error: unknown,
 	dump: RawHttpRequestDump | undefined,
 ): Promise<string> {
-	if (!dump || extractHttpStatusFromError(error) !== 400) {
+	// Never persist dumps under the test runner: providers exercise the 400 path
+	// with mocked fetch responses, which would otherwise litter the real ~/.prometheus logs.
+	if (!dump || isBunTestRuntime() || extractHttpStatusFromError(error) !== 400) {
 		return message;
 	}
 
@@ -65,12 +63,6 @@ export async function finalizeErrorMessage(
 	return appendRawHttpRequestDumpFor400(message, error, rawRequestDump);
 }
 
-export function withHttpStatus(error: unknown, status: number): Error {
-	const wrapped = error instanceof Error ? error : new Error(String(error));
-	(wrapped as ErrorWithStatus).status = status;
-	return wrapped;
-}
-
 /**
  * Rewrite error message for GitHub Copilot request failures.
  * Must run AFTER finalizeErrorMessage since it replaces the message entirely.

package/src/utils/idle-iterator.ts

@@ -2,6 +2,8 @@ import { $env } from "@prometheus-ai/utils";
 
 const DEFAULT_STREAM_IDLE_TIMEOUT_MS = 120_000;
 const DEFAULT_STREAM_FIRST_EVENT_TIMEOUT_MS = 100_000;
+/** Re-mint persistent race promises every N iterations (see hoisted-racer comment). */
+const RACER_REMINT_INTERVAL = 1024;
 
 function normalizeIdleTimeoutMs(value: string | undefined, fallback: number): number | undefined {
 	if (value === undefined) return fallback;
@@ -130,8 +132,11 @@ export async function* iterateWithIdleTimeout<T>(
 		firstItemTimeoutMs !== undefined && firstItemTimeoutMs > 0 ? Date.now() + firstItemTimeoutMs : undefined;
 	const abortSignal = options.abortSignal;
 	const iterator = iterable[Symbol.asyncIterator]();
+	let iteratorClosed = false;
 
 	const closeIterator = (): void => {
+		if (iteratorClosed) return;
+		iteratorClosed = true;
 		const returnPromise = iterator.return?.();
 		if (returnPromise) {
 			void returnPromise.catch(() => {});
@@ -167,100 +172,266 @@ export async function* iterateWithIdleTimeout<T>(
 		(firstItemTimeoutMs === undefined || firstItemTimeoutMs <= 0) &&
 		(options.idleTimeoutMs === undefined || options.idleTimeoutMs <= 0);
 
-	while (true) {
-		let activeTimeoutMs: number | undefined;
-		if (awaitingFirstItem) {
-			if (firstItemDeadlineMs !== undefined) {
-				activeTimeoutMs = firstItemDeadlineMs - Date.now();
-				if (activeTimeoutMs <= 0) {
-					options.onFirstItemTimeout?.();
-					closeIterator();
-					throw new Error(options.firstItemErrorMessage ?? options.errorMessage);
-				}
-			}
-		} else if (options.idleTimeoutMs !== undefined && options.idleTimeoutMs > 0) {
-			activeTimeoutMs = options.idleTimeoutMs - (Date.now() - lastProgressAt);
-			if (activeTimeoutMs <= 0) {
-				options.onIdle?.();
-				closeIterator();
-				throw new Error(options.errorMessage);
-			}
-		}
-
-		const nextResultPromise = withRacy(iterator.next());
+	// Persistent racers, hoisted out of the per-item loop. The abort promise can
+	// only ever resolve once (abort latches), and a timeout resolution always
+	// precedes a throw — so neither needs per-item re-creation. This keeps the
+	// token hot path free of timer create/destroy and listener churn.
+	//
+	// Each Promise.race() call still attaches a reaction record to every pending
+	// racer, and those records live until the racer settles — so a never-firing
+	// abort/timeout promise would accumulate one record per streamed item for
+	// the stream's whole life. The loop re-mints both promises every
+	// RACER_REMINT_INTERVAL iterations to keep that retention bounded; the
+	// listener and timer callbacks resolve through late-bound variables so a
+	// re-mint never strands them.
+	let abortPromise: Promise<{ kind: "abort" }> | undefined;
+	let abortListener: (() => void) | undefined;
+	let resolveAbort: ((value: { kind: "abort" }) => void) | undefined;
+	if (abortSignal) {
+		const { promise, resolve } = Promise.withResolvers<{ kind: "abort" }>();
+		resolveAbort = resolve;
+		abortListener = () => resolveAbort?.({ kind: "abort" });
+		abortSignal.addEventListener("abort", abortListener, { once: true });
+		abortPromise = promise;
+	}
 
-		const racers: Array<
-			Promise<
-				| { kind: "next"; result: IteratorResult<T> }
-				| { kind: "error"; error: unknown }
-				| { kind: "timeout" }
-				| { kind: "abort" }
-			>
-		> = [nextResultPromise];
+	let timeoutPromise: Promise<{ kind: "timeout" }> | undefined;
+	let resolveTimeout: ((value: { kind: "timeout" }) => void) | undefined;
+	let timeoutFired = false;
+	let timer: NodeJS.Timeout | undefined;
+	let timerFireAtMs = Infinity;
 
-		let timer: NodeJS.Timeout | undefined;
-		let resolveTimeout: ((value: { kind: "timeout" }) => void) | undefined;
-		const enforceTimeout = !noTimeoutEnforced && activeTimeoutMs !== undefined && activeTimeoutMs > 0;
-		if (enforceTimeout) {
+	const currentDeadlineMs = (): number | undefined => {
+		if (awaitingFirstItem) return firstItemDeadlineMs;
+		if (options.idleTimeoutMs !== undefined && options.idleTimeoutMs > 0) {
+			return lastProgressAt + options.idleTimeoutMs;
+		}
+		return undefined;
+	};
+	const onTimerFire = (): void => {
+		timer = undefined;
+		timerFireAtMs = Infinity;
+		const deadlineMs = currentDeadlineMs();
+		if (deadlineMs === undefined) return;
+		const remainingMs = deadlineMs - Date.now();
+		if (remainingMs > 0) {
+			// Progress moved the deadline since this timer was armed — re-arm for
+			// the remainder. One stale wake per idle period, not one per item.
+			timerFireAtMs = deadlineMs;
+			timer = setTimeout(onTimerFire, remainingMs);
+			return;
+		}
+		timeoutFired = true;
+		resolveTimeout?.({ kind: "timeout" });
+	};
+	const armTimer = (deadlineMs: number): void => {
+		if (timeoutPromise === undefined || timeoutFired) {
+			// A fired-but-unconsumed resolution (the item won the same race) is
+			// stale — racing it again would fake a timeout, so mint a fresh one.
 			const { promise, resolve } = Promise.withResolvers<{ kind: "timeout" }>();
+			timeoutPromise = promise;
 			resolveTimeout = resolve;
-			timer = setTimeout(() => resolve({ kind: "timeout" }), activeTimeoutMs);
-			racers.push(promise);
+			timeoutFired = false;
 		}
-
-		let abortListener: (() => void) | undefined;
-		let resolveAbort: ((value: { kind: "abort" }) => void) | undefined;
-		if (abortSignal) {
-			const { promise, resolve } = Promise.withResolvers<{ kind: "abort" }>();
-			resolveAbort = resolve;
-			abortListener = () => resolve({ kind: "abort" });
-			abortSignal.addEventListener("abort", abortListener, { once: true });
-			racers.push(promise);
+		if (timer !== undefined) {
+			// An armed timer firing at or before the new deadline re-arms itself.
+			if (timerFireAtMs <= deadlineMs) return;
+			clearTimeout(timer);
 		}
+		timerFireAtMs = deadlineMs;
+		timer = setTimeout(onTimerFire, Math.max(0, deadlineMs - Date.now()));
+	};
 
-		try {
-			const outcome = await Promise.race(racers);
-			if (outcome.kind === "abort") {
-				closeIterator();
-				throw abortReason(abortSignal!);
+	try {
+		let raceCount = 0;
+		while (true) {
+			if (++raceCount % RACER_REMINT_INTERVAL === 0) {
+				if (abortPromise !== undefined && !abortSignal!.aborted) {
+					const { promise, resolve } = Promise.withResolvers<{ kind: "abort" }>();
+					resolveAbort = resolve;
+					abortPromise = promise;
+				}
+				if (timeoutPromise !== undefined && !timeoutFired) {
+					const { promise, resolve } = Promise.withResolvers<{ kind: "timeout" }>();
+					resolveTimeout = resolve;
+					timeoutPromise = promise;
+				}
 			}
-			if (outcome.kind === "timeout") {
-				if (!awaitingFirstItem) {
+			let activeTimeoutMs: number | undefined;
+			if (awaitingFirstItem) {
+				if (firstItemDeadlineMs !== undefined) {
+					activeTimeoutMs = firstItemDeadlineMs - Date.now();
+					if (activeTimeoutMs <= 0) {
+						options.onFirstItemTimeout?.();
+						closeIterator();
+						throw new Error(options.firstItemErrorMessage ?? options.errorMessage);
+					}
+				}
+			} else if (options.idleTimeoutMs !== undefined && options.idleTimeoutMs > 0) {
+				activeTimeoutMs = options.idleTimeoutMs - (Date.now() - lastProgressAt);
+				if (activeTimeoutMs <= 0) {
 					options.onIdle?.();
-				} else {
-					options.onFirstItemTimeout?.();
+					closeIterator();
+					throw new Error(options.errorMessage);
 				}
-				closeIterator();
-				throw new Error(
-					!awaitingFirstItem ? options.errorMessage : (options.firstItemErrorMessage ?? options.errorMessage),
-				);
 			}
-			if (outcome.kind === "error") {
-				throw outcome.error;
+
+			const nextResultPromise = withRacy(iterator.next());
+
+			const racers: Array<
+				Promise<
+					| { kind: "next"; result: IteratorResult<T> }
+					| { kind: "error"; error: unknown }
+					| { kind: "timeout" }
+					| { kind: "abort" }
+				>
+			> = [nextResultPromise];
+
+			const enforceTimeout = !noTimeoutEnforced && activeTimeoutMs !== undefined && activeTimeoutMs > 0;
+			if (enforceTimeout) {
+				armTimer(Date.now() + activeTimeoutMs!);
+				racers.push(timeoutPromise!);
+			}
+			if (abortPromise) {
+				racers.push(abortPromise);
 			}
-			if (outcome.result.done) {
-				markFirstItemReceived();
-				return;
+
+			// Tracks whether this iteration handed an item to the consumer and resumed
+			// normally. Any other exit — internal throw, `done` return, or the consumer
+			// abandoning us via `.return()`/`.throw()` at the `yield` below — must close
+			// the upstream iterator so the underlying SSE body / SDK stream (and its
+			// socket) is released instead of being left suspended.
+			let continuing = false;
+			try {
+				const outcome = await Promise.race(racers);
+				if (outcome.kind === "abort") {
+					closeIterator();
+					throw abortReason(abortSignal!);
+				}
+				if (outcome.kind === "timeout") {
+					if (!awaitingFirstItem) {
+						options.onIdle?.();
+					} else {
+						options.onFirstItemTimeout?.();
+					}
+					closeIterator();
+					throw new Error(
+						!awaitingFirstItem ? options.errorMessage : (options.firstItemErrorMessage ?? options.errorMessage),
+					);
+				}
+				if (outcome.kind === "error") {
+					throw outcome.error;
+				}
+				if (outcome.result.done) {
+					markFirstItemReceived();
+					return;
+				}
+				const item = outcome.result.value;
+				// Non-progress items (e.g. provider keepalives, synthetic `start` events that
+				// arrive before the model has produced any tokens) MUST NOT flip us out of
+				// `awaitingFirstItem`. Otherwise the next iteration switches from the (longer)
+				// first-item watchdog to the (shorter) idle watchdog while we're still waiting
+				// on the model's first real output.
+				if (isProgressItem(item)) {
+					markFirstItemReceived();
+					lastProgressAt = Date.now();
+				}
+				yield item;
+				continuing = true;
+			} finally {
+				if (!continuing) closeIterator();
 			}
-			const item = outcome.result.value;
-			// Non-progress items (e.g. provider keepalives, synthetic `start` events that
-			// arrive before the model has produced any tokens) MUST NOT flip us out of
-			// `awaitingFirstItem`. Otherwise the next iteration switches from the (longer)
-			// first-item watchdog to the (shorter) idle watchdog while we're still waiting
-			// on the model's first real output.
-			if (isProgressItem(item)) {
-				markFirstItemReceived();
-				lastProgressAt = Date.now();
+		}
+	} finally {
+		if (timer !== undefined) clearTimeout(timer);
+		// Settle the persistent racers so the final Promise.race releases them.
+		resolveTimeout?.({ kind: "timeout" });
+		if (abortListener && abortSignal) {
+			abortSignal.removeEventListener("abort", abortListener);
+		}
+		resolveAbort?.({ kind: "abort" });
+	}
+}
+
+export interface TerminalGraceIteratorOptions {
+	/**
+	 * Epoch-ms timestamp at which the consumer observed a logically terminal
+	 * item (e.g. a chat-completions chunk carrying `finish_reason`), or
+	 * `undefined` while the stream is still mid-response. Read before every
+	 * pull, so the consumer can flip it between yields.
+	 */
+	finishedAtMs: () => number | undefined;
+	/**
+	 * Post-terminal budget: how long after `finishedAtMs()` to keep draining
+	 * trailing items (e.g. a usage-only chunk or the `[DONE]` sentinel) before
+	 * ending the iteration cleanly. The deadline is fixed at
+	 * `finishedAtMs() + graceMs`; trailing items do not extend it, so
+	 * keepalive-only servers cannot hold the stream open.
+	 */
+	graceMs: number;
+	/**
+	 * Invoked when the grace window closes with the source still open. Use it
+	 * to abort the underlying request: the source generator is typically parked
+	 * mid-`next()` (not at a yield), so a queued `.return()` alone cannot reach
+	 * the transport until that pending read settles.
+	 */
+	onGraceEnd?: () => void;
+}
+
+/**
+ * Yields items from an async iterable until the consumer marks the stream
+ * logically finished AND the source stays silent past a short grace window.
+ *
+ * Misbehaving OpenAI-compatible servers deliver the terminal chunk but never
+ * send `[DONE]` nor close the connection; without this guard the consumer
+ * hangs on `iterator.next()` until the idle watchdog converts an
+ * already-successful turn into a timeout error. Grace expiry is a clean end
+ * of iteration, never an error.
+ */
+export async function* iterateWithTerminalGrace<T>(
+	iterable: AsyncIterable<T>,
+	options: TerminalGraceIteratorOptions,
+): AsyncGenerator<T> {
+	const iterator = iterable[Symbol.asyncIterator]();
+	try {
+		while (true) {
+			const finishedAtMs = options.finishedAtMs();
+			if (finishedAtMs === undefined) {
+				const result = await iterator.next();
+				if (result.done) return;
+				yield result.value;
+				continue;
+			}
+			const remainingMs = finishedAtMs + options.graceMs - Date.now();
+			if (remainingMs <= 0) {
+				options.onGraceEnd?.();
+				return;
 			}
-			yield item;
-		} finally {
-			if (timer !== undefined) clearTimeout(timer);
-			// Resolve dangling promises so the racers don't leak (Promise.race is one-shot).
-			resolveTimeout?.({ kind: "timeout" });
-			if (abortListener && abortSignal) {
-				abortSignal.removeEventListener("abort", abortListener);
+			const nextPromise = iterator.next();
+			let timer: NodeJS.Timeout | undefined;
+			const timeoutPromise = new Promise<"timeout">(resolve => {
+				timer = setTimeout(() => resolve("timeout"), remainingMs);
+			});
+			try {
+				const outcome = await Promise.race([nextPromise, timeoutPromise]);
+				if (outcome === "timeout") {
+					// The abandoned read settles (likely rejects) once onGraceEnd
+					// aborts the transport — mark it handled so it cannot surface
+					// as an unhandled rejection.
+					nextPromise.catch(() => {});
+					options.onGraceEnd?.();
+					return;
+				}
+				if (outcome.done) return;
+				yield outcome.value;
+			} finally {
+				if (timer !== undefined) clearTimeout(timer);
 			}
-			resolveAbort?.({ kind: "abort" });
+		}
+	} finally {
+		const returnPromise = iterator.return?.();
+		if (returnPromise) {
+			void Promise.resolve(returnPromise).catch(() => {});
 		}
 	}
 }