@gtrabanco/pi-nan-provider 0.10.3 → 1.0.0

package/src/provider-factory.ts

@@ -22,6 +22,12 @@
  * Streaming API resolution: bare root under pi's compat alias is the instance
  * anchor; otherwise a file URL from import.meta.resolve; never a bare subpath
  * (loud failure — issue #8). Full logic in src/pi-ai-loader.ts.
+ *
+ * The provider also exposes NaN image models (flux-2-klein, qwen-image-2.1)
+ * through an image API ("nan-images"): `baselineImageModels` are added to the
+ * model set and a `ProviderImages` implementation is wired via
+ * `createNanImagesApi`, so pi's model registry can route
+ * `generateImages()` calls through it.
  */
 
 import * as piAi from "@earendil-works/pi-ai";
@@ -31,11 +37,17 @@ import type {
 	RefreshModelsContext,
 } from "@earendil-works/pi-ai";
 import { withContextOverflowClassification } from "./context-overflow-classifier.ts";
+import { NAN_IMAGE_API, baselineImageModels, createNanImagesApi } from "./images.ts";
 import {
 	baselineModels,
 	DEFAULT_MODELS_TIMEOUT_MS,
 	resolveCatalog,
 	type CatalogSource,
+	overlayGeneratedCapabilities,
+	buildOverlayMap,
+	applyOverlay,
+	type OverlayEntry,
+	NAN_GENERATED_MODELS,
 } from "./fetch-models.ts";
 import { resolveOpenAICompletionsApi } from "./pi-ai-loader.ts";
 import { sanitizeOpenAICompatPayload } from "./openai-compat-sanitizer.ts";
@@ -126,6 +138,55 @@ export function wrapApiForStrictSanitization(api: ProviderStreams): ProviderStre
 	};
 }
 
+/**
+ * Stale-catalog overlay (issue #19): pi-ai's `createProvider` builds an
+ * internal `refreshModels` where phase 1 restores `context.stored.models`
+ * verbatim and phase 2 calls the input `fetchModels` + persists. The
+ * `dynamicModels` list is a closure variable — external code cannot assign
+ * it. When pi restores a persisted store entry written by an older version
+ * of this package (e.g. without `thinkingLevelMap`), that stale copy is
+ * used as-is and never refreshed (because `checkedAt` is recent).
+ *
+ * `currentModels()` merges static `input.models` with `dynamicModels` and a
+ * dynamic entry REPLACES the baseline entry with the same id+type (the
+ * shadowing mechanism). The `ModelsImpl` in pi resolves models exclusively
+ * via `getModels()` / `getAllModels()` on the provider object — it never
+ * inspects the internal dynamicModels list.
+ *
+ * The fix: wrap `getModels` and `getAllModels` on the returned provider so
+ * every call re-checks the generated catalog and applies capability data
+ * when the id exists there. The store decides which ids exist; the
+ * generated catalog decides capability data. Read-time overlay is the only
+ * provider-side hook that wins over the stale restore.
+ *
+ * `overlayMap` is prebuilt once at factory time from BOTH the generated
+ * chat catalog AND `baselineImageModels(source)`. Both `getModels`
+ * and `getAllModels` use `applyOverlay` for O(1) per-model lookup (GAP 2).
+ * This also covers image models: pi-ai's image models are only visible
+ * through `getAllModels()`, so a stale stored IMAGE entry would otherwise
+ * replace the baseline image entry with no overlay (GAP 1).
+ *
+ * We do NOT override `refreshModels` — the built-in phases stay intact. The
+ * overlay only affects read paths.
+ */
+function withStaleCatalogOverlay(
+	provider: Provider<"openai-completions">,
+	overlayMap: Map<string, OverlayEntry>,
+): Provider<"openai-completions"> {
+	return {
+		...provider,
+		getModels: () => {
+			const models = provider.getModels();
+			return models.map((m) => applyOverlay(m, overlayMap)) as typeof models;
+		},
+		getAllModels: () => {
+			const all = provider.getAllModels?.();
+			if (!all) return [];
+			return all.map((m) => applyOverlay(m, overlayMap));
+		},
+	};
+}
+
 /**
  * Build a complete pi-ai Provider for an OpenAI-compatible endpoint:
  *
@@ -133,11 +194,15 @@ export function wrapApiForStrictSanitization(api: ProviderStreams): ProviderStre
  *   `/login <id>` prompts for the key (pi's `envApiKeyAuth` semantics — the
  *   same precedence the built-in providers use). No prompt is needed when the
  *   env var is set.
- * - models: the generated fallback catalog as static baseline, so models are
+ * - models: the generated chat-model fallback catalog plus NaN image models
+ *   (flux-2-klein, qwen-image-2.1) as static baseline, so all models are
  *   available with zero network.
  * - fetchModels: live `/models` IDs × generated capability data; falls back
  *   to the baseline when the endpoint is unreachable. pi's Models runtime
  *   drives refreshes (startup/periodic) and persists the overlay.
+ * - images: a `ProviderImages` implementation for NaN's OpenAI-compatible
+ *   image endpoints, so `generateImages()` calls are routed through pi's
+ *   model registry.
  * - api: the openai-completions streaming implementation resolved via
  *   `resolveOpenAICompletionsApi` (bare root under pi's compat alias as
  *   the instance anchor; otherwise a file URL from import.meta.resolve;
@@ -157,33 +222,40 @@ export async function createNanCompatibleProvider(
 	// `available` (e.g. premium-tier models you are not subscribed to).
 	let liveIds: Set<string> | undefined;
 
-	return piAi.createProvider({
-		id: config.id,
-		name: config.name,
-		baseUrl: config.baseUrl,
-		auth: { apiKey: piAi.envApiKeyAuth(`${config.name} API key`, config.envVars) },
-		models: baselineModels(source),
-		fetchModels: async (context: RefreshModelsContext) => {
-			const credential = context.credential;
-			const resolved = await resolveCatalog(source, {
-				apiKey: credential?.type === "api_key" ? credential.key : undefined,
-				timeoutMs: options.timeoutMs ?? DEFAULT_MODELS_TIMEOUT_MS,
-				fetchImpl: options.fetchImpl,
-			});
-			liveIds = resolved.liveIds;
-			return resolved.models;
-		},
-		filterModels: (models) => {
-			// Snapshot: TS can't prove `liveIds` unchanged across the closure boundary.
-			const current = liveIds;
-			return current ? models.filter((model) => current.has(model.id)) : models;
-		},
-		// Sanitize the payload for NaN's strict schema, then classify an opaque
-		// generic 400 as a context overflow when the request we just sent was over
-		// the model's window. The second layer keeps a replayed cross-model
-		// reasoning trace (or any other over-window request the context-hook guard
-		// cannot reach, e.g. NAN_THINKING_GUARD=0) recoverable instead of wedging
-		// the session — see src/context-overflow-classifier.ts.
-		api: withContextOverflowClassification(wrapApiForStrictSanitization(apiFactory())),
-	});
+	const imageModels = baselineImageModels(source);
+	const overlayMap = buildOverlayMap(source, NAN_GENERATED_MODELS, imageModels);
+
+	return withStaleCatalogOverlay(
+		piAi.createProvider({
+			id: config.id,
+			name: config.name,
+			baseUrl: config.baseUrl,
+			auth: { apiKey: piAi.envApiKeyAuth(`${config.name} API key`, config.envVars) },
+			models: [...baselineModels(source), ...baselineImageModels(source)],
+			images: { [NAN_IMAGE_API]: createNanImagesApi({ fetchImpl: options.fetchImpl, timeoutMs: options.timeoutMs }) },
+			fetchModels: async (context: RefreshModelsContext) => {
+				const credential = context.credential;
+				const resolved = await resolveCatalog(source, {
+					apiKey: credential?.type === "api_key" ? credential.key : undefined,
+					timeoutMs: options.timeoutMs ?? DEFAULT_MODELS_TIMEOUT_MS,
+					fetchImpl: options.fetchImpl,
+				});
+				liveIds = resolved.liveIds;
+				return resolved.models;
+			},
+			filterModels: (models) => {
+				// Snapshot: TS can't prove `liveIds` unchanged across the closure boundary.
+				const current = liveIds;
+				return current ? models.filter((model) => current.has(model.id)) : models;
+			},
+			// Sanitize the payload for NaN's strict schema, then classify an opaque
+			// generic 400 as a context overflow when the request we just sent was over
+			// the model's window. The second layer keeps a replayed cross-model
+			// reasoning trace (or any other over-window request the context-hook guard
+			// cannot reach, e.g. NAN_THINKING_GUARD=0) recoverable instead of wedging
+			// the session — see src/context-overflow-classifier.ts.
+			api: withContextOverflowClassification(wrapApiForStrictSanitization(apiFactory())),
+		}),
+		overlayMap,
+	);
 }