@gtrabanco/pi-nan-provider 0.10.3 → 1.0.1

package/src/context-overflow-classifier.ts

@@ -23,22 +23,50 @@
  *
  * This module makes the package **re-check the request size on the way out**
  * (one of the two upstream fixes the tracking issue names). After the request is
- * sent and the gateway answers, if
- *   - the terminal assistant message is an error whose text carries NaN's
- *     generic 400 marker, AND
- *   - the request we sent was estimated to exceed the model's context window,
- * then the error message is rewritten into a form that matches pi-ai's
- * documented overflow patterns, so pi compacts and retries instead of wedging.
- * The original provider text is preserved in the rewritten message. The
- * reclassification is deliberately conservative: a generic 400 on a
- * within-window request is left untouched, so unrelated 400s are never
+ * sent and the gateway answers, two independent branches reclassify:
+ *
+ * 1. **HTTP 400 branch** — keys on the **model's context window**. When the
+ *    terminal error carries NaN's generic 400 marker AND the estimated input
+ *    tokens exceed the destination model's `contextWindow`, the message is
+ *    rewritten. A 400 may be caused by other schema violations, so the rewrite
+ *    conservatively requires the request to be over the window.
+ *
+ * 2. **HTTP 524 branch** — keys on a **fixed cold-cache ceiling**
+ *    (`NAN_COLD_CACHE_CEILING_TOKENS`, 200,000). Cloudflare's 524
+ *    origin_response_timeout fires when a large cold prompt hangs past the 120s
+ *    proxy read timeout. A 524 on an over-ceiling request is an overflow
+ *    masquerading as a timeout; a 524 under the ceiling is a genuine transient
+ *    timeout. The ceiling is a *practical endpoint limit* (Cloudflare 120s vs
+ *    cold-prompt prefill), NOT a model capability — so the model's
+ *    `contextWindow` is irrelevant to the 524 reclassify decision. The rewritten
+ *    message still mentions the model window for user clarity, but the
+ *    overflow-classification hinges only on the ceiling.
+ *
+ * The original provider text is preserved in rewritten messages. The
+ * reclassification is deliberately conservative so unrelated errors are never
  * mislabelled as overflow.
  *
+ * ---
+ *
+ * Cold-cache 524 (issue #18): large COLD prompts (≳1 MB body / ~220k tokens) to
+ * https://api.nan.builders/v1/chat/completions hang past Cloudflare's 120s Proxy
+ * Read Timeout → HTTP 524 origin_response_timeout with JSON body
+ *   {"status":524,"error_name":"origin_response_timeout","retryable":true,...}
+ * (sometimes a Cloudflare HTML page for even larger bodies). Measured: cold
+ * 208,036 tokens (960 KB) succeeds (TTFT 12.9 s); ~220k and ~286k tokens cold
+ * → 524 at ~126 s. pi-ai classifies 524 as RETRYABLE (retrying an identical cold
+ * payload always fails) and its `isContextOverflow()` matches nothing, so the
+ * session wedges. When the error is a 524 AND the estimated input tokens exceed
+ * `NAN_COLD_CACHE_CEILING_TOKENS`, the error is rewritten into a pi-recognizable
+ * overflow message so pi compacts and retries instead of wedging. The original
+ * provider text is preserved. A 524 on a small request (under the ceiling) is
+ * left untouched — it is a genuine transient timeout.
+ *
  * Estimation ratio: the same chars/3.47 the issue's offline session measurements
  * use (`chars/3.47 + ~31K tools + ~12K system`). The system prompt and tool
  * schemas are counted explicitly, so the constant is only applied to the
  * message/character budget. This is a heuristic for deciding whether a generic
- * 400 is plausibly an overflow — it never fabricates model metadata.
+ * 400 / 524 is plausibly an overflow — it never fabricates model metadata.
  */
 
 import type {
@@ -52,6 +80,19 @@ import type {
 /** The exact generic body NaN's gateway returns for an over-window request. */
 export const NAN_GENERIC_BAD_REQUEST = "Invalid request. Check your request parameters.";
 
+/**
+ * Cloudflare 524 origin_response_timeout cold-cache ceiling (tokens).
+ *
+ * Measured (issue #18 repro): cold 208,036 tokens (960 KB) succeeds
+ * (TTFT 12.9 s); ~220k and ~286k tokens cold → 524 at ~126 s. 200,000 is a
+ * conservative midpoint — well below the working 208k measurement — so the
+ * working cold cache path is not affected while larger cold requests hit this
+ * rewrite. This is a practical endpoint ceiling (Cloudflare 120s proxy read
+ * timeout vs cold prompt prefill), NOT a model capability — contextWindow in
+ * the catalog must NOT change.
+ */
+export const NAN_COLD_CACHE_CEILING_TOKENS = 200_000;
+
 /**
  * Characters per token used for the request-size estimate. Same ratio as the
  * issue's offline measurements (chars/3.47); a conservative, provider-agnostic
@@ -93,11 +134,34 @@ export function isGenericNanBadRequest(message: Pick<AssistantMessage, "errorMes
 	return typeof message?.errorMessage === "string" && message.errorMessage.includes(NAN_GENERIC_BAD_REQUEST);
 }
 
+/**
+ * True when a terminal error message carries a Cloudflare 524
+ * origin_response_timeout marker. Matches the 524 status followed by either
+ * `origin_response_timeout` in the JSON body, or Cloudflare HTML markers
+ * (`Error 524` / `A timeout occurred`). Conservative: only these known
+ * markers are matched, not any arbitrary 524 response.
+ */
+function isCloudflare524(message: Pick<AssistantMessage, "errorMessage"> | undefined): boolean {
+	const text = message?.errorMessage;
+	if (typeof text !== "string") return false;
+	// Match the status code + origin_response_timeout JSON body.
+	if (text.includes("524") && text.includes("origin_response_timeout")) return true;
+	// Cloudflare HTML error page variant (occurs for very large bodies).
+	if (text.includes("524") && (text.includes("Error 524") || text.includes("A timeout occurred"))) return true;
+	return false;
+}
+
 /**
  * Reclassify a terminal assistant error as a context overflow when — and only
  * when — NaN returned its generic 400 for a request estimated to exceed the
- * model's context window. Mutates and returns the message (the same object
- * flows through the error event and `result()`).
+ * model's context window, OR when a Cloudflare 524 fires for an over-ceiling
+ * cold request. Mutates and returns the message (the same object flows through
+ * the error event and `result()`).
+ *
+ * The two error paths use independent thresholds:
+ * - 400: estimated > contextWindow  (conservative; a 400 may be schema-related)
+ * - 524: estimated > CEILING        (a 524 is either overflow or transient;
+ *                                    the ceiling is the discriminator)
  */
 export function classifyContextOverflowError<T extends AssistantMessage>(
 	message: T,
@@ -105,16 +169,43 @@ export function classifyContextOverflowError<T extends AssistantMessage>(
 	estimatedInputTokens: number,
 ): T {
 	if (message.stopReason !== "error") return message;
-	if (!isGenericNanBadRequest(message)) return message;
-	const contextWindow = model?.contextWindow ?? 0;
-	if (!Number.isFinite(contextWindow) || contextWindow <= 0) return message;
-	if (!Number.isFinite(estimatedInputTokens) || estimatedInputTokens <= contextWindow) return message;
-
-	message.errorMessage =
-		`Requested token count exceeds the model's maximum context length of ${contextWindow} tokens ` +
-		`(estimated ${Math.ceil(estimatedInputTokens)} input tokens). ` +
-		`NaN's gateway answered HTTP 400 "${NAN_GENERIC_BAD_REQUEST}" instead of naming the overflow, ` +
-		`so the context must be compacted before retrying.`;
+
+	const contextWindow = model?.contextWindow;
+	const hasValidWindow = Number.isFinite(contextWindow) && contextWindow! > 0;
+
+	// --- 400 branch: keys on the model's context window ---
+	if (isGenericNanBadRequest(message) && hasValidWindow
+		&& Number.isFinite(estimatedInputTokens)
+		&& estimatedInputTokens > contextWindow!) {
+		message.errorMessage =
+			`Requested token count exceeds the model's maximum context length of ${contextWindow!} tokens ` +
+			`(estimated ${Math.ceil(estimatedInputTokens)} input tokens). ` +
+			`NaN's gateway answered HTTP 400 "${NAN_GENERIC_BAD_REQUEST}" instead of naming the overflow, ` +
+			`so the context must be compacted before retrying.`;
+		return message;
+	}
+
+	// --- 524 branch: keys on the cold-cache ceiling only ---
+	// A 524 origin_response_timeout on an over-ceiling request means the edge
+	// timed out while prefilling a large cold prompt — it is an overflow, not a
+	// transient error. Small-requests that hit 524 are genuine timeouts.
+	if (isCloudflare524(message)
+		&& Number.isFinite(estimatedInputTokens)
+		&& estimatedInputTokens > NAN_COLD_CACHE_CEILING_TOKENS) {
+		const windowClause = hasValidWindow
+			? ` (declared model window: ${contextWindow!} tokens)`
+			: "";
+		message.errorMessage =
+			`Requested token count exceeds the maximum context length of ${NAN_COLD_CACHE_CEILING_TOKENS} tokens ` +
+			`servable with a cold prompt cache through NaN's endpoint` +
+			windowClause +
+			`, so the context must be compacted before retrying. ` +
+			`(estimated ${Math.ceil(estimatedInputTokens)} input tokens; ` +
+			`Cloudflare answers HTTP 524 origin_response_timeout past the 120s proxy read timeout). ` +
+			`Original error: ${message.errorMessage}.`;
+		return message;
+	}
+
 	return message;
 }
 
@@ -178,4 +269,4 @@ export function withContextOverflowClassification(api: ProviderStreams): Provide
 		streamSimple: (model, context, options) =>
 			classifyStreamContextOverflow(api.streamSimple(model, context, options), model, estimateRequestTokens(context)),
 	};
-}
+}

package/src/fetch-models.ts

@@ -23,10 +23,12 @@
  * back to the generated catalog so startup is never blocked.
  */
 
-import type { Model, OpenAICompletionsCompat } from "@earendil-works/pi-ai";
-import type { ThinkingLevelMap } from "@earendil-works/pi-ai";
+import type { ImageModel, Model } from "@earendil-works/pi-ai";
+import type { ThinkingLevelMap, OpenAICompletionsCompat } from "@earendil-works/pi-ai";
 import { GENERATED_CATALOG_META, NAN_GENERATED_MODELS } from "../scripts/models.generated.ts";
 
+export { NAN_GENERATED_MODELS };
+
 export type { Model };
 
 /**
@@ -341,3 +343,146 @@ export async function fetchNanCompatibleModels(
 export function generatedCatalogMeta() {
 	return GENERATED_CATALOG_META;
 }
+
+// ── Stale catalog overlay (issue #19) ───────────────────────────────────────
+
+/**
+ * Build a lookup map from model id to generated catalog entry. Used by
+ * `overlayGeneratedCapabilities` so the overlay is a single Map lookup.
+ */
+export function buildGeneratedCatalogMap(
+	generated: readonly GeneratedModelEntry[] = NAN_GENERATED_MODELS,
+): Map<string, GeneratedModelEntry> {
+	return new Map(generated.map((entry) => [entry.id, entry]));
+}
+
+/**
+ * Pre-built lookup map for generated chat-catalog entries.
+ * Built once at module evaluation time so every read is a true O(1) Map.get.
+ * (The doc on `overlayGeneratedCapabilities` previously claimed O(1)
+ *  `buildGeneratedCatalogMap` was called on every call — this fixes that.)
+ */
+const _generatedCatalogMap: Map<string, GeneratedModelEntry> = buildGeneratedCatalogMap();
+
+/**
+ * Overlay generated-catalog capability data onto a stored/restored model when
+ * the model id exists in the generated catalog (and the model type matches).
+ *
+ * Returns the generated catalog entry converted to a pi-ai Model via `toModel`
+ * when there is an exact id match with the same api (openai-completions).
+ * Image models, classifier models, or ids not in the catalog pass through
+ * unchanged — the persisted store decides which ids exist, but only the
+ * generated catalog decides capability data (issue #19).
+ *
+ * When `prebuilt` is provided it must be a Map<string, GeneratedModelEntry>
+ * keyed by model id for chat entries only. This parameter exists so the
+ * provider wrapper can pass a prebuilt map instead of rebuilding it every
+ * call (module-level `_generatedCatalogMap` is used when omitted).
+ *
+ * This function is pure, stateless, and O(1) — a single Map.get lookup.
+ * It is called at read time from the provider wrapper so every consumer
+ * (pi, the model registry, downstream code) sees the latest capabilities.
+ */
+export function overlayGeneratedCapabilities(
+	model: Model<string> | ImageModel<string>,
+	generatedOrPrebuilt?: readonly GeneratedModelEntry[] | Map<string, GeneratedModelEntry>,
+): Model<string> | ImageModel<string> {
+	if (!("reasoning" in model && "contextWindow" in model)) {
+		// ImageModel and ClassifierModel do not have reasoning/contextWindow.
+		return model;
+	}
+	// Resolve the lookup map: prebuilt map > array > module-level singleton.
+	const byId: Map<string, GeneratedModelEntry> =
+		generatedOrPrebuilt instanceof Map
+			? generatedOrPrebuilt
+			: Array.isArray(generatedOrPrebuilt)
+				? buildGeneratedCatalogMap(generatedOrPrebuilt)
+				: _generatedCatalogMap;
+	const entry = byId.get(model.id);
+	// Only match when the generated catalog has this id AND the api is the
+	// chat api (openai-completions). Image models use a different api and
+	// a separate static list — they are not in this catalog.
+	if (entry && model.api === NAN_COMPAT_API) {
+		return toModel(entry, { providerId: model.provider, baseUrl: model.baseUrl });
+	}
+	return model;
+}
+
+// ── Combined overlay for chat + image models (GAP 1: getAllModels wrapping) ─
+
+/**
+ * Internal representation of an overlay entry: either a generated chat entry
+ * (GeneratedModelEntry → toModel result) or a baseline image model
+ * (ImageModel). Keys are model ids that are unique across chat/image
+ * namespaces.
+ */
+export interface OverlayEntry {
+	/** For chat models: the pi-ai Model produced by toModel(entry, source). */
+	chat?: Model<string>;
+	/** For image models: the ImageModel from baselineImageModels(source). */
+	image?: ImageModel<string>;
+}
+
+/**
+ * Build a combined overlay map keyed by model id. The overlay contains:
+ * - For every generated chat model id: the toModel() result (chat entry).
+ * - For every baseline image model: the ImageModel itself (image entry).
+ *
+ * This map is built once at factory time (in withStaleCatalogOverlay) and
+ * used by both `getModels` and `getAllModels` wrappers so every read path
+ * receives generated-catalog capability data regardless of whether the
+ * model is chat or image (issue #19, GAP 1).
+ */
+export function buildOverlayMap(
+	source: CatalogSource,
+	generated: readonly GeneratedModelEntry[] = NAN_GENERATED_MODELS,
+	imageModels?: ImageModel<string>[],
+): Map<string, OverlayEntry> {
+	const overlay = new Map<string, OverlayEntry>();
+	// Chat models from the generated catalog.
+	for (const entry of generated) {
+		overlay.set(entry.id, { chat: toModel(entry, source) });
+	}
+	// Image models from the baseline (these live under a different api id
+	// so they cannot coexist as chat entries — ids are unique across types).
+	if (imageModels) {
+		for (const img of imageModels) {
+			overlay.set(img.id, { image: img });
+		}
+	}
+	return overlay;
+}
+
+/**
+ * Internal model-like shape used by applyOverlay: every model type has id, api,
+ * and (for images) a `type` discriminator.
+ */
+interface _ModelLike {
+	id: string;
+	api: string;
+	type?: string;
+}
+
+/**
+ * Apply the combined overlay (chat + image) to a single model. Looks up the
+ * id in the provided map and returns the overlaid entry when present.
+ * Image models (type === "image") receive their ImageModel overlay; chat
+ * models receive their generated-catalog entry via `overlayGeneratedCapabilities`.
+ * ClassifierModel and other model types pass through unchanged.
+ */
+export function applyOverlay<T extends _ModelLike>(
+	model: T,
+	overlay: Map<string, OverlayEntry>,
+): T {
+	const entry = overlay.get(model.id);
+	if (!entry) return model;
+	// Image models: return the baseline ImageModel overlay.
+	if ((model as _ModelLike).type === "image") {
+		return (entry.image ?? model) as T;
+	}
+	// Chat models: use the standard overlay function (still O(1)).
+	if (entry.chat && model.api === NAN_COMPAT_API) {
+		return (entry.chat as unknown as T) ?? model;
+	}
+	return model;
+}