@gtrabanco/pi-nan-provider 0.10.3 → 1.0.0

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

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;