@gtrabanco/pi-nan-provider 0.10.2 → 1.0.0

package/src/images.ts

@@ -0,0 +1,408 @@
+/**
+ * NaN image generation API — custom API id, model definitions, and a
+ * ProviderImages implementation talking to the OpenAI-compatible image
+ * endpoints at api.nan.builders.
+ *
+ * NaN exposes two OpenAI-compatible image endpoints:
+ *
+ *   POST /v1/images/generations  — text-to-image (JSON body)
+ *   POST /v1/images/edits        — image-to-image (multipart/form-data)
+ *
+ * Both require the inference-tier membership (403 otherwise) and share a
+ * 20 req/min / 100 req/month quota separate from the chat token budget
+ * (https://nan.builders/docs/models and https://nan.builders/openapi.json,
+ * checked 2026-10-01).
+ *
+ * We register a custom API id "nan-images" so pi's model registry can route
+ * generateImages calls through our implementation (createNanImagesApi), and
+ * register the resulting ProviderImages via createProvider({ images: ... }).
+ *
+ * Style note (per repo conventions): only type-only imports from the bare
+ * root "@earendil-works/pi-ai". No subpath imports — pi's extension loader
+ * aliases subpaths and the whole extension fails to load. No imageErrorResult
+ * import either: pi-ai exposes it only under the forbidden
+ * "@earendil-works/pi-ai/utils/..." subpath, so error results are built here.
+ */
+
+import type {
+	AssistantImages,
+	ImageModel,
+	ProviderImages,
+} from "@earendil-works/pi-ai";
+
+// ── Custom API id ──────────────────────────────────────────────────────────
+
+/** Custom image-api id for NaN image generation. */
+export const NAN_IMAGE_API = "nan-images";
+
+// ── Model definitions ─────────────────────────────────────────────────────
+
+/**
+ * Static image-model definition carrying input/output modalities and
+ * provenance notes. Kept local (not in the generated chat catalog) because
+ * these are image endpoints, not chat models: NON_CHAT_MODEL_IDS still keeps
+ * them out of the chat catalog, and this list is the image-side counterpart.
+ */
+export interface NanImageModelDefinition {
+	/** Stable model id, e.g. "flux-2-klein". */
+	id: string;
+	/** Human-readable model name, e.g. "FLUX 2 Klein". */
+	name: string;
+	/** Input modalities the model accepts. */
+	input: ("text" | "image")[];
+	/** Output modalities the model produces. */
+	output: ("text" | "image")[];
+	/** Provenance notes citing sources for every capability claim. */
+	notes: string[];
+}
+
+/**
+ * All NaN image models known to this provider. Two entries, ordered:
+ * flux-2-klein (text+image input), qwen-image-2.1 (text-only).
+ */
+export const NAN_IMAGE_MODELS: readonly NanImageModelDefinition[] = [
+	{
+		id: "flux-2-klein",
+		name: "FLUX 2 Klein",
+		input: ["text", "image"],
+		output: ["image"],
+		notes: [
+			"Image generation endpoint (text-to-image and image-to-image). " +
+				"Documented in https://nan.builders/docs/models and " +
+				"https://nan.builders/openapi.json (checked 2026-10-01): POST " +
+				"/v1/images/generations and POST /v1/images/edits (1-4 reference " +
+				"images; mask unsupported). Requires inference-tier membership " +
+				"(403 otherwise); image endpoints are rate-limited separately " +
+				"from chat (20 req/min, 100 req/month shared with qwen-image-2.1) " +
+				"and do not consume the chat token budget. models.dev provider " +
+				"nan does not list flux-2-klein (checked 2026-10-01), so this is " +
+				"a manual entry; per-token cost is 0 because NaN bills image " +
+				"generation by request quota on the same membership model as chat.",
+		],
+	},
+	{
+		id: "qwen-image-2.1",
+		name: "Qwen-Image-2.1",
+		input: ["text"],
+		output: ["image"],
+		notes: [
+			"Image generation endpoint (text-to-image only). Documented in " +
+				"https://nan.builders/docs/models and " +
+				"https://nan.builders/openapi.json (checked 2026-10-01): POST " +
+				"/v1/images/generations. models.dev provider nan lists it " +
+				"(checked 2026-10-01) with cost 0/0 and limit.output 0 (it " +
+				"returns image bytes, not tokens). Same inference-tier " +
+				"requirement and 100 req/month quota shared with " +
+				"flux-2-klein; image-to-image (/images/edits) is not documented " +
+				"for this model.",
+		],
+	},
+] as const;
+
+// ── Baseline models ────────────────────────────────────────────────────────
+
+/**
+ * Structural source shape (same as fetch-models.ts CatalogSource) declared
+ * locally so this module has zero circular dependency on fetch-models.
+ */
+interface ImageCatalogSource {
+	/** Provider id as registered in pi, e.g. "nan". */
+	providerId: string;
+	/** OpenAI-compatible base URL including version path. */
+	baseUrl: string;
+}
+
+/**
+ * Map every entry in NAN_IMAGE_MODELS to a pi-ai ImageModel for our custom
+ * API id. Per-token cost is 0: NaN bills image generation by request quota
+ * on the membership model, the same convention the generated chat catalog
+ * uses for NaN's quota-based pricing.
+ */
+export function baselineImageModels(
+	source: ImageCatalogSource,
+): ImageModel<typeof NAN_IMAGE_API>[] {
+	return NAN_IMAGE_MODELS.map(
+		(model): ImageModel<typeof NAN_IMAGE_API> => ({
+			id: model.id,
+			name: model.name,
+			type: "image",
+			api: NAN_IMAGE_API,
+			provider: source.providerId,
+			baseUrl: source.baseUrl,
+			input: [...model.input],
+			output: [...model.output],
+			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+		}),
+	);
+}
+
+// ── API implementation ─────────────────────────────────────────────────────
+
+/** Injectable options for createNanImagesApi (fetchImpl is for tests). */
+export interface CreateNanImagesApiOptions {
+	/** Injectable fetch for tests; defaults to globalThis.fetch. */
+	fetchImpl?: typeof fetch;
+	/** Default per-request timeout in ms when the caller does not supply one. */
+	timeoutMs?: number;
+}
+
+/**
+ * Build a ProviderImages implementation for NaN's OpenAI-compatible image
+ * endpoints (POST /v1/images/generations and POST /v1/images/edits).
+ *
+ * We always request b64_json (not "url") so the returned images carry their
+ * base64 data directly — no separate download step is needed. NaN returns
+ * raw base64 with no MIME type, so the output type is sniffed from the
+ * decoded magic bytes.
+ *
+ * Contract: never throws. Auth failures, invalid input, transport errors and
+ * non-OK responses all surface as an AssistantImages with stopReason
+ * "error" (or "aborted" when the caller's signal aborted).
+ */
+export function createNanImagesApi(
+	options?: CreateNanImagesApiOptions,
+): ProviderImages {
+	const fetchImpl = options?.fetchImpl ?? globalThis.fetch;
+
+	return {
+		async generateImages(model, context, requestOptions) {
+			const base: AssistantImages = {
+				api: model.api,
+				provider: model.provider,
+				model: model.id,
+				output: [],
+				stopReason: "stop",
+				timestamp: Date.now(),
+			};
+
+			try {
+				// ── Auth ────────────────────────────────────────────────
+				const apiKey = requestOptions?.apiKey;
+				if (!apiKey) {
+					throw new Error("No API key for provider: " + model.provider);
+				}
+
+				// ── Input extraction ────────────────────────────────────
+				const text = context.input
+					.filter((block) => block.type === "text")
+					.map((block) => block.text)
+					.join("\n")
+					.trim();
+				const referenceImages = context.input.filter(
+					(block) => block.type === "image",
+				);
+
+				if (text.length === 0 && referenceImages.length === 0) {
+					throw new Error(
+						"Image generation requires a text prompt or a reference image",
+					);
+				}
+
+				// Reference images need a model that accepts image input
+				// (flux-2-klein); qwen-image-2.1 is text-to-image only.
+				if (
+					referenceImages.length > 0 &&
+					!model.input.includes("image")
+				) {
+					throw new Error(
+						"Model " + model.id + " does not accept image input",
+					);
+				}
+
+				// ── Endpoint selection ──────────────────────────────────
+				const endpoint =
+					referenceImages.length > 0
+						? "/images/edits"
+						: "/images/generations";
+				const url = model.baseUrl.replace(/\/+$/, "") + endpoint;
+
+				// ── Headers ─────────────────────────────────────────────
+				// Merge model headers with caller headers, dropping null
+				// values (ProviderHeaders uses null to suppress a default).
+				// Bearer auth is applied last so caller values cannot
+				// override the provider credential.
+				const headers: Record<string, string> = {};
+				for (const [key, value] of Object.entries({
+					...model.headers,
+					...requestOptions?.headers,
+				})) {
+					if (typeof value === "string") headers[key] = value;
+				}
+				headers.Authorization = "Bearer " + apiKey;
+
+				// ── Optional metadata passthrough ───────────────────────
+				// n, size, seed and guidance are NaN extensions accepted on
+				// both endpoints; pass them through only with the right type.
+				const metadata = (requestOptions?.metadata ?? {}) as Record<
+					string,
+					unknown
+				>;
+				const optionalFields: Record<string, string | number> = {};
+				if (typeof metadata.n === "number") optionalFields.n = metadata.n;
+				if (typeof metadata.size === "string") optionalFields.size = metadata.size;
+				if (typeof metadata.seed === "number") optionalFields.seed = metadata.seed;
+				if (typeof metadata.guidance === "number") {
+					optionalFields.guidance = metadata.guidance;
+				}
+
+				// ── Request body ────────────────────────────────────────
+				let body: FormData | string;
+				if (referenceImages.length > 0) {
+					// Image-to-image: multipart/form-data. Content-Type is
+					// deliberately NOT set — fetch adds the boundary.
+					const form = new FormData();
+					form.set("model", model.id);
+					form.set("prompt", text);
+					form.set("response_format", "b64_json");
+					for (const [key, value] of Object.entries(optionalFields)) {
+						form.set(key, String(value));
+					}
+					for (const image of referenceImages) {
+						// Decode base64 with atob (no Buffer: the extension
+						// must not assume a Node-only global).
+						const bytes = Uint8Array.from(atob(image.data), (char) =>
+							char.charCodeAt(0),
+						);
+						form.append(
+							"image[]",
+							new Blob([bytes], {
+								type: image.mimeType || "image/png",
+							}),
+							"reference.png",
+						);
+					}
+					body = form;
+				} else {
+					headers["Content-Type"] = "application/json";
+					// fetch does not serialize plain objects: build the JSON text here.
+					body = JSON.stringify({
+						model: model.id,
+						prompt: text,
+						response_format: "b64_json",
+						...optionalFields,
+					});
+				}
+
+				// ── Signal (timeout + caller abort) ─────────────────────
+				const timeoutMs = requestOptions?.timeoutMs ?? options?.timeoutMs;
+				let signal: AbortSignal | undefined;
+				if (timeoutMs) {
+					const timeoutSignal = AbortSignal.timeout(timeoutMs);
+					signal = requestOptions?.signal
+						? AbortSignal.any([timeoutSignal, requestOptions.signal])
+						: timeoutSignal;
+				} else if (requestOptions?.signal) {
+					signal = requestOptions.signal;
+				}
+
+				// ── HTTP call ───────────────────────────────────────────
+				const response = await fetchImpl(url, {
+					method: "POST",
+					headers,
+					body: body as RequestInit["body"],
+					...(signal ? { signal } : {}),
+				});
+
+				if (!response.ok) {
+					// Best-effort error body: NaN usually returns
+					// {"error":{"message":"..."}} (403 for a missing
+					// inference tier, 429 for the quota).
+					let bodyText = "";
+					try {
+						bodyText = await response.text();
+					} catch {
+						// Ignore read errors — best effort only.
+					}
+					throw new Error(
+						"NaN image request failed: HTTP " +
+							response.status +
+							(bodyText ? " " + bodyText : ""),
+					);
+				}
+
+				// ── Parse response ──────────────────────────────────────
+				const payload = (await response.json()) as {
+					data?: Array<{ b64_json?: string }>;
+				};
+				for (const item of payload.data ?? []) {
+					if (
+						typeof item.b64_json === "string" &&
+						item.b64_json.length > 0
+					) {
+						base.output.push({
+							type: "image",
+							data: item.b64_json,
+							mimeType: sniffImageMimeType(item.b64_json),
+						});
+					}
+				}
+				if (base.output.length === 0) {
+					throw new Error("NaN image response contained no image data");
+				}
+
+				return base;
+			} catch (error) {
+				base.stopReason =
+					requestOptions?.signal?.aborted === true ? "aborted" : "error";
+				base.errorMessage =
+					error instanceof Error ? error.message : String(error);
+				return base;
+			}
+		},
+	};
+}
+
+// ── MIME-type sniffing ─────────────────────────────────────────────────────
+
+/**
+ * Sniff the MIME type from a base64-encoded image string by inspecting the
+ * leading magic bytes.
+ *
+ * PNG:  89 50 4E 47  → "image/png"
+ * JPEG: FF D8 FF     → "image/jpeg"
+ * WebP: RIFF....WEBP → "image/webp"
+ *
+ * NaN returns raw base64 without a MIME type, so we infer it from the
+ * signature instead of guessing. Falls back to "image/png" and never throws.
+ */
+function sniffImageMimeType(base64: string): string {
+	try {
+		// Only the first 16 bytes are needed to identify PNG/JPEG/WebP.
+		const decoded = atob(base64.slice(0, 16));
+		const bytes = new Uint8Array(decoded.length);
+		for (let i = 0; i < decoded.length; i++) {
+			bytes[i] = decoded.charCodeAt(i);
+		}
+
+		if (
+			bytes[0] === 0x89 &&
+			bytes[1] === 0x50 &&
+			bytes[2] === 0x4e &&
+			bytes[3] === 0x47
+		) {
+			return "image/png";
+		}
+		if (
+			bytes[0] === 0xff &&
+			bytes[1] === 0xd8 &&
+			bytes[2] === 0xff
+		) {
+			return "image/jpeg";
+		}
+		if (
+			bytes[0] === 0x52 &&
+			bytes[1] === 0x49 &&
+			bytes[2] === 0x46 &&
+			bytes[3] === 0x46 &&
+			bytes[8] === 0x57 &&
+			bytes[9] === 0x45 &&
+			bytes[10] === 0x42 &&
+			bytes[11] === 0x50
+		) {
+			return "image/webp";
+		}
+		return "image/png";
+	} catch {
+		return "image/png";
+	}
+}

package/src/index.ts

@@ -17,11 +17,12 @@
  *
  *  2. MCP servers via pi's native registerMcpServer (pi >=0.99):
  *     - `nan-search` — the official NaN remote MCP server
- *       (https://api.nan.builders/mcp), exposing `mcp__nan-search__web_search`.
+ *       (https://api.nan.builders/mcp), exposing `mcp__nan_search__web_search`
+ *       (pi >=0.99.2; 0.99.0-0.99.1 kept the hyphen: `mcp__nan-search__web_search`).
  *       Default: enabled. NAN_MCP_TOOLS=0 to disable.
  *     - `nan-media` — the community stdio media server
  *       (flux-2-klein / kokoro / whisper), exposing
- *       `mcp__nan-media__generate_image` etc. Default: enabled,
+ *       `mcp__nan_media__generate_image` etc. Default: enabled,
  *       NAN_MEDIA_MCP=0 to disable.
  *
  *     Both are session-scoped (visible in /mcp with source "extension").
@@ -109,6 +110,7 @@ async function registerWebSearchMcpServer(pi: ExtensionAPI, apiKey: string | und
 		url: "https://api.nan.builders/mcp",
 		headers: { Authorization: `Bearer ${apiKey}` },
 		exposure: "direct",
+		description: "Web search through NaN's gateway (official NaN MCP server)",
 	});
 	return true;
 }
@@ -130,6 +132,7 @@ function registerMediaMcpServer(pi: ExtensionAPI, apiKey: string | undefined): b
 		args: mediaArgs.slice(1),
 		env: { [NAN_API_KEY_ENV]: apiKey },
 		exposure: "direct",
+		description: "NaN media tools: image generation/editing, text-to-speech and speech-to-text via the community nan-mcp-server",
 		timeout: mediaMcpTimeoutSec(),
 	});
 	return true;

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,
+	);
 }