@gtrabanco/pi-nan-provider 0.10.2 → 1.0.0

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

@@ -12,19 +12,23 @@
  *     only model `id`s — no capability fields. It is used as the authoritative
  *     list of what the endpoint returns for the API key.
  *
- * Merge: live IDs × generated capability data — all live IDs surface:
- * a live ID with generated data keeps its generated capabilities, and an
- * uncatalogued live ID (e.g. a new or undocumented model) gets conservative
- * placeholder limits (the same defaults used in custom-provider.md's
- * dynamic-discovery example) with no reasoning support — capabilities stay
- * "unknown", nothing is fabricated. On fetch failure, timeout, or an unusable
- * response, callers fall back to the generated catalog so startup is never blocked.
+ * Merge: live IDs × generated capability data — non-chat IDs (see
+ * `NON_CHAT_MODEL_IDS`) are excluded from the registered model set and
+ * reported in `MergedCatalog.nonChat`; all other live IDs surface: a live ID
+ * with generated data keeps its generated capabilities, and an uncatalogued
+ * live ID (e.g. a new or undocumented model) gets conservative placeholder
+ * limits (the same defaults used in custom-provider.md's dynamic-discovery
+ * example) with no reasoning support — capabilities stay "unknown", nothing is
+ * fabricated. On fetch failure, timeout, or an unusable response, callers fall
+ * 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 };
 
 /**
@@ -81,6 +85,32 @@ export const UNKNOWN_MODEL_LIMITS = { contextWindow: 128_000, maxTokens: 4_096 }
 /** Timeout for the live /models fetch; matches the pi-synthetic-provider precedent (~3s). */
 export const DEFAULT_MODELS_TIMEOUT_MS = 3_000;
 
+/**
+ * Live endpoint IDs classified as non-chat and excluded from the registered
+ * model set. Keys are model ids; values are provenance reasons citing the
+ * source that designated the endpoint as non-chat.
+ *
+ * Source: https://nan.builders/docs/models (checked 2026-10-01), plus
+ * qwen-image-2.1: https://models.dev API `limit.output: 0` (checked 2026-09-29).
+ * These are MCP-bridge territory (embedding / rerank / TTS / STT / image
+ * generation), not chat catalog models — see AGENTS.md classification notes
+ * and issue #15.
+ */
+export const NON_CHAT_MODEL_IDS: Readonly<Record<string, string>> = {
+	"qwen3-embedding":
+		"embedding endpoint (https://nan.builders/docs/models, checked 2026-10-01) — not a chat model",
+	rerank:
+		"rerank endpoint (https://nan.builders/docs/models, checked 2026-10-01) — not a chat model",
+	kokoro:
+		"text-to-speech (TTS) endpoint (https://nan.builders/docs/models, checked 2026-10-01) — not a chat model",
+	whisper:
+		"speech-to-text (STT) endpoint (https://nan.builders/docs/models, checked 2026-10-01) — not a chat model",
+	"flux-2-klein":
+		"image generation endpoint (https://nan.builders/docs/models, checked 2026-10-01) — not a chat model",
+	"qwen-image-2.1":
+		"image-generation model (text→image; output modality image) documented in https://nan.builders/docs/models and https://nan.builders/openapi.json (checked 2026-09-29) — not a chat model, out of the static chat catalog (models.dev reports limit.output 0)",
+} as const;
+
 export interface CatalogSource {
 	/** Provider id as registered in pi, e.g. "nan". */
 	providerId: string;
@@ -108,7 +138,11 @@ export function toModel(entry: GeneratedModelEntry, source: CatalogSource): Mode
 
 /** The generated fallback catalog as pi-ai Models for the given provider. */
 export function baselineModels(source: CatalogSource): Model<"openai-completions">[] {
-	return NAN_GENERATED_MODELS.map((entry) => toModel(entry, source));
+	// Defensive: the generated baseline never contains non-chat ids today,
+	// but the invariant must hold on every path (issue #15).
+	return NAN_GENERATED_MODELS.filter(
+		(entry) => !NON_CHAT_MODEL_IDS[entry.id],
+	).map((entry) => toModel(entry, source));
 }
 
 export interface LiveModelListOptions {
@@ -163,12 +197,18 @@ export interface MergedCatalog {
 	matched: string[];
 	/** Allowlisted live IDs kept with unknown capabilities (conservative limits). */
 	unknown: string[];
+	/**
+	 * Live model IDs classified as non-chat and excluded from the registered
+	 * model set, in live order. Populated from `NON_CHAT_MODEL_IDS` (issue #15).
+	 */
+	nonChat: string[];
 }
 
 /**
- * Merge live model IDs with the generated capability catalog. All live IDs
- * surface: live IDs with generated data keep their generated capabilities;
- * uncatalogued live IDs get conservative placeholder limits,
+ * Merge live model IDs with the generated capability catalog. Non-chat IDs
+ * (from `NON_CHAT_MODEL_IDS`) are excluded and reported in `nonChat`; all
+ * other live IDs surface: live IDs with generated data keep their generated
+ * capabilities; uncatalogued live IDs get conservative placeholder limits,
  * `reasoning: false`, and zero cost (documented defaults, not invented
  * capabilities).
  */
@@ -181,8 +221,15 @@ export function mergeLiveWithGenerated(
 	const models: Model<"openai-completions">[] = [];
 	const matched: string[] = [];
 	const unknown: string[] = [];
+	const nonChat: string[] = [];
 
 	for (const id of liveIds) {
+		// Issue #15: non-chat endpoint ids must never surface as chat models.
+		const nonChatReason = NON_CHAT_MODEL_IDS[id];
+		if (nonChatReason) {
+			nonChat.push(id);
+			continue;
+		}
 		const entry = byId.get(id);
 		if (entry) {
 			models.push(toModel(entry, source));
@@ -217,7 +264,7 @@ export function mergeLiveWithGenerated(
 		}
 	}
 
-	return { models, matched, unknown };
+	return { models, matched, unknown, nonChat };
 }
 
 export interface FetchModelsOptions {
@@ -240,13 +287,19 @@ export interface ResolvedCatalog {
 	liveIds: Set<string> | undefined;
 	/** Allowlisted live IDs kept with unknown capabilities (present only when liveIds is set). */
 	unknownIds: string[];
+	/**
+	 * Live model IDs classified as non-chat and excluded from the registered
+	 * model set (issue #15). Empty when the live fetch failed.
+	 */
+	nonChatIds: string[];
 }
 
 /**
  * Resolve the effective catalog: live `/models` IDs × generated capability
  * data. A successful live fetch is authoritative (tier-aware: it reflects
  * exactly what your NaN key can use); failure degrades to the generated
- * fallback so startup is never blocked.
+ * fallback so startup is never blocked. Non-chat IDs are excluded from
+ * models and reported in `nonChatIds` (issue #15).
  */
 export async function resolveCatalog(
 	source: CatalogSource,
@@ -260,11 +313,16 @@ export async function resolveCatalog(
 			timeoutMs: options.timeoutMs,
 			fetchImpl: options.fetchImpl,
 		});
-		if (!liveIds) return { models: baseline, liveIds: undefined, unknownIds: [] };
+		if (!liveIds) return { models: baseline, liveIds: undefined, unknownIds: [], nonChatIds: [] };
 		const merged = mergeLiveWithGenerated(liveIds, source);
-		return { models: merged.models, liveIds: new Set(liveIds), unknownIds: merged.unknown };
+		return {
+			models: merged.models,
+			liveIds: new Set(liveIds),
+			unknownIds: merged.unknown,
+			nonChatIds: merged.nonChat,
+		};
 	} catch {
-		return { models: baseline, liveIds: undefined, unknownIds: [] };
+		return { models: baseline, liveIds: undefined, unknownIds: [], nonChatIds: [] };
 	}
 }
 
@@ -285,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;
+}