pi-free 2.0.14 → 2.1.0

package/providers/kilo/kilo.ts

@@ -25,6 +25,7 @@ import {
 } from "../../config.ts";
 import { URL_KILO_TOS } from "../../constants.ts";
 import { isFreeModel, registerWithGlobalToggle } from "../../lib/registry.ts";
+import { wrapSessionStartHandler } from "../../lib/session-start-metrics.ts";
 import { cleanModelName, logWarning } from "../../lib/util.ts";
 import {
 	createCtxReRegister,
@@ -35,6 +36,119 @@ import {
 import { loginKilo, refreshKiloToken } from "./kilo-auth.ts";
 import { fetchKiloModels, KILO_GATEWAY_BASE } from "./kilo-models.ts";
 
+/** Kilo Gateway compat overrides, borrowed from pi-kilo-provider. */
+const KILO_COMPAT = {
+	supportsStore: false,
+	supportsDeveloperRole: false,
+	supportsReasoningEffort: false,
+	supportsUsageInStreaming: false,
+	supportsStrictMode: false,
+	thinkingFormat: "openrouter" as const,
+	maxTokensField: "max_tokens" as const,
+};
+
+// =============================================================================
+// XML leak detection and auto-retry
+// =============================================================================
+
+/**
+ * Detect when a model outputs raw XML tool calls instead of using
+ * native function calling. This happens when gateways don't pass
+ * tool definitions to certain models (e.g., step-3.7-flash via Kilo).
+ *
+ * Pattern: <tool><name>tool_name</name><param>...</param></tool>
+ */
+function detectXmlToolLeak(text: string): boolean {
+	// Use simple string searches instead of regex to avoid ReDoS risks.
+	const lower = text.toLowerCase();
+	return (
+		lower.includes("<tool>") ||
+		lower.includes("<tool_call>") ||
+		lower.includes("<function_call>") ||
+		lower.includes("<invoke") ||
+		lower.includes("<antml:tool_use>")
+	);
+}
+
+function findTag(
+	text: string,
+	tag: string,
+	start = 0,
+): { start: number; end: number; content: string } | null {
+	const open = `<${tag}>`;
+	const close = `</${tag}>`;
+	const openIdx = text.indexOf(open, start);
+	if (openIdx === -1) return null;
+	const contentStart = openIdx + open.length;
+	const closeIdx = text.indexOf(close, contentStart);
+	if (closeIdx === -1) return null;
+	return {
+		start: openIdx,
+		end: closeIdx + close.length,
+		content: text.slice(contentStart, closeIdx),
+	};
+}
+
+/**
+ * Parse XML tool calls and convert to pi's tool call format.
+ * Returns null if parsing fails.
+ *
+ * Uses simple string scanning instead of regex with backreferences
+ * to avoid super-linear backtracking (ReDoS).
+ */
+function parseXmlToolCalls(
+	text: string,
+): Array<{ name: string; arguments: Record<string, unknown> }> | null {
+	try {
+		const calls: Array<{ name: string; arguments: Record<string, unknown> }> =
+			[];
+		let searchStart = 0;
+		while (true) {
+			const toolBlock = findTag(text, "tool", searchStart);
+			if (!toolBlock) break;
+			searchStart = toolBlock.end;
+
+			const nameTag = findTag(toolBlock.content, "name");
+			if (!nameTag) continue;
+			const name = nameTag.content.trim();
+			if (!name) continue;
+
+			const args: Record<string, unknown> = {};
+			// Skip the <name>...</name> block we already consumed.
+			let paramStart = nameTag.end;
+			const paramsText = toolBlock.content;
+			while (true) {
+				const openIdx = paramsText.indexOf("<", paramStart);
+				if (openIdx === -1) break;
+				const closeOpenIdx = paramsText.indexOf(">", openIdx);
+				if (closeOpenIdx === -1) break;
+				const tagName = paramsText.slice(openIdx + 1, closeOpenIdx).trim();
+				if (!tagName || tagName.startsWith("/")) {
+					paramStart = closeOpenIdx + 1;
+					continue;
+				}
+				const closeTag = `</${tagName}>`;
+				const closeIdx = paramsText.indexOf(closeTag, closeOpenIdx + 1);
+				if (closeIdx === -1) break;
+				const value = paramsText.slice(closeOpenIdx + 1, closeIdx).trim();
+				try {
+					args[tagName] = JSON.parse(value);
+				} catch {
+					args[tagName] = value;
+				}
+				paramStart = closeIdx + closeTag.length;
+			}
+			calls.push({ name, arguments: args });
+		}
+		return calls.length > 0 ? calls : null;
+	} catch {
+		return null;
+	}
+}
+
+// =============================================================================
+// Extension entry point
+// =============================================================================
 const KILO_PROVIDER_CONFIG = {
 	providerId: PROVIDER_KILO,
 	baseUrl: KILO_GATEWAY_BASE,
@@ -44,6 +158,19 @@ const KILO_PROVIDER_CONFIG = {
 	},
 };
 
+/** Apply Kilo-specific compat overrides while preserving provider/model values. */
+function applyKiloCompat<T extends { compat?: ProviderModelConfig["compat"] }>(
+	models: T[],
+): T[] {
+	return models.map((m) => ({
+		...m,
+		compat: {
+			...KILO_COMPAT,
+			...m.compat,
+		},
+	}));
+}
+
 export default async function kiloProvider(pi: ExtensionAPI) {
 	// Try to fetch ALL models at startup (like Cline/OpenRouter)
 	// If no API key, this will return free models only
@@ -77,9 +204,11 @@ export default async function kiloProvider(pi: ExtensionAPI) {
 	const stored: StoredModels = { free: freeModels, all: allModels };
 
 	// Create re-register function
-	const reRegister = createReRegister(pi, {
+	const baseReRegister = createReRegister(pi, {
 		...KILO_PROVIDER_CONFIG,
 	});
+	const reRegister = (models: ProviderModelConfig[]) =>
+		baseReRegister(applyKiloCompat(models));
 
 	// Register with global toggle system
 	registerWithGlobalToggle(
@@ -108,9 +237,11 @@ export default async function kiloProvider(pi: ExtensionAPI) {
 				stored.free = freeModels;
 
 				// Update global toggle registration with new lists
-				const globalReRegister = createReRegister(pi, {
+				const baseGlobalReRegister = createReRegister(pi, {
 					...KILO_PROVIDER_CONFIG,
 				});
+				const globalReRegister = (models: ProviderModelConfig[]) =>
+					baseGlobalReRegister(applyKiloCompat(models));
 				registerWithGlobalToggle(PROVIDER_KILO, stored, globalReRegister, true);
 
 				// If paid mode is enabled, show all models
@@ -132,21 +263,24 @@ export default async function kiloProvider(pi: ExtensionAPI) {
 			const template = models.find((m) => m.provider === PROVIDER_KILO);
 			if (!template) return models;
 			const nonKilo = models.filter((m) => m.provider !== PROVIDER_KILO);
-			const fullModels = allModels.map((m) => ({
-				...template,
-				id: m.id,
-				name: cleanModelName(m.name),
-				reasoning: m.reasoning,
-				input: m.input,
-				cost: m.cost,
-				contextWindow: m.contextWindow,
-				maxTokens: m.maxTokens,
-			}));
-			return [...nonKilo, ...fullModels];
+			const fullModels = applyKiloCompat(
+				allModels.map((m) => ({
+					...template,
+					id: m.id,
+					name: cleanModelName(m.name),
+					reasoning: m.reasoning,
+					input: m.input,
+					cost: m.cost,
+					contextWindow: m.contextWindow,
+					maxTokens: m.maxTokens,
+				})),
+			);
+			return [...nonKilo, ...fullModels] as Model<"openai-completions">[];
 		},
 	};
 
 	// Register initial provider (default to free models)
+	const modelsWithCompat = applyKiloCompat(currentModels);
 	pi.registerProvider(PROVIDER_KILO, {
 		baseUrl: KILO_GATEWAY_BASE,
 		apiKey: "$KILO_API_KEY",
@@ -155,7 +289,7 @@ export default async function kiloProvider(pi: ExtensionAPI) {
 			"X-KILOCODE-EDITORNAME": "Pi",
 			"User-Agent": "pi-free-providers",
 		},
-		models: enhanceWithCI(currentModels),
+		models: enhanceWithCI(modelsWithCompat),
 		oauth: oauthConfig,
 	});
 
@@ -228,36 +362,129 @@ export default async function kiloProvider(pi: ExtensionAPI) {
 		}
 	});
 
-	// Refresh models on session start if authenticated
-	pi.on("session_start", async (_event, ctx) => {
-		const cred = ctx.modelRegistry.authStorage.get(PROVIDER_KILO);
+	// ── XML leak detection and auto-retry ─────────────────────────
+	//
+	// When a model outputs raw XML tool calls (<tool><name>...</name></tool>)
+	// instead of native function calling, detect it and rewrite the message
+	// to force the model to use proper tool calling on the next turn.
 
-		if (cred?.type === "oauth") {
-			try {
-				const newModels = await fetchKiloModels({
-					token: cred.access,
-					freeOnly: false,
-				});
-				allModels = newModels;
-				stored.all = allModels;
-				freeModels = allModels.filter((m) =>
-					isFreeModel({ ...m, provider: PROVIDER_KILO }, allModels),
-				);
-				stored.free = freeModels;
+	let xmlLeakRetryCount = 0;
+	const MAX_XML_LEAK_RETRIES = 2;
 
-				// Update global toggle registration
-				const ctxReRegister = createCtxReRegister(ctx as any, {
-					...KILO_PROVIDER_CONFIG,
-				});
-				registerWithGlobalToggle(PROVIDER_KILO, stored, ctxReRegister, true);
+	(pi as any).on("message_end", (event: any, ctx: any) => {
+		if (ctx.model?.provider !== PROVIDER_KILO) return;
 
-				// Apply current view mode
-				if (showPaidModels && !getKiloFreeOnly()) {
-					ctxReRegister(allModels);
-				}
-			} catch (error) {
-				logWarning("kilo", "Failed to refresh models at session start", error);
-			}
+		const msg = event.message;
+		if (msg.role !== "assistant") return;
+
+		// Extract text content from the message
+		let text = "";
+		if (typeof msg.content === "string") {
+			text = msg.content;
+		} else if (Array.isArray(msg.content)) {
+			text = msg.content
+				.filter((p: any) => p?.type === "text" && typeof p?.text === "string")
+				.map((p: any) => p.text)
+				.join("\n");
+		}
+
+		if (!text || !detectXmlToolLeak(text)) {
+			xmlLeakRetryCount = 0; // Reset on clean response
+			return;
 		}
+
+		// XML leak detected
+		if (xmlLeakRetryCount >= MAX_XML_LEAK_RETRIES) {
+			xmlLeakRetryCount = 0;
+			logWarning("kilo", "XML tool leak persisted after retries, giving up");
+			return;
+		}
+
+		xmlLeakRetryCount++;
+		logWarning(
+			"kilo",
+			`XML tool leak detected (attempt ${xmlLeakRetryCount}/${MAX_XML_LEAK_RETRIES}), rewriting message`,
+		);
+
+		// Try to parse the XML tool calls
+		const parsedCalls = parseXmlToolCalls(text);
+		if (parsedCalls && parsedCalls.length > 0) {
+			// We parsed the tool calls - convert to proper toolCall format
+			const toolCalls = parsedCalls.map((call, i) => ({
+				type: "toolCall" as const,
+				id: `xml_leak_${Date.now()}_${i}`,
+				name: call.name,
+				arguments: call.arguments,
+			}));
+
+			return {
+				...msg,
+				content: [
+					{
+						type: "text",
+						text:
+							text.replace(/<tool>[\s\S]*?<\/tool>/g, "").trim() ||
+							"(parsed tool calls)",
+					},
+					...toolCalls,
+				],
+			};
+		}
+
+		// Can't parse - add a correction message to force retry
+		// We rewrite the message to include a note about using proper tool calling
+		return {
+			...msg,
+			content: [
+				{
+					type: "text",
+					text: `${text}\n\n---\n[SYSTEM: You outputted XML tool calls instead of using the function calling API. Please use the native tool/function calling format with JSON arguments, not XML tags like <tool>.]`,
+				},
+			],
+		};
 	});
+
+	// Refresh models on session start if authenticated
+	let refreshInFlight: Promise<void> | undefined;
+	pi.on(
+		"session_start",
+		wrapSessionStartHandler("kilo", (_event, ctx) => {
+			const cred = ctx.modelRegistry.authStorage.get(PROVIDER_KILO);
+			if (cred?.type !== "oauth" || refreshInFlight) return Promise.resolve();
+
+			refreshInFlight = fetchKiloModels({ token: cred.access, freeOnly: false })
+				.then((newModels) => {
+					allModels = newModels;
+					stored.all = allModels;
+					freeModels = allModels.filter((m) =>
+						isFreeModel({ ...m, provider: PROVIDER_KILO }, allModels),
+					);
+					stored.free = freeModels;
+
+					// Update global toggle registration
+					const baseCtxReRegister = createCtxReRegister(ctx as any, {
+						...KILO_PROVIDER_CONFIG,
+					});
+					const ctxReRegister = (models: ProviderModelConfig[]) =>
+						baseCtxReRegister(applyKiloCompat(models));
+					registerWithGlobalToggle(PROVIDER_KILO, stored, ctxReRegister, true);
+
+					// Apply current view mode
+					if (showPaidModels && !getKiloFreeOnly()) {
+						ctxReRegister(allModels);
+					}
+				})
+				.catch((error) => {
+					logWarning(
+						"kilo",
+						"Failed to refresh models at session start",
+						error instanceof Error ? error.message : String(error),
+					);
+				})
+				.finally(() => {
+					refreshInFlight = undefined;
+				});
+			return Promise.resolve();
+		}),
+	);
 }

package/providers/model-fetcher.ts

@@ -3,8 +3,9 @@
  * Consolidates duplicate logic from openrouter.ts and kilo-models.ts
  */
 
-import { DEFAULT_FETCH_TIMEOUT_MS, URL_MODELS_DEV } from "../constants.ts";
-import type { ModelsDevModel, ProviderModelConfig } from "../lib/types.ts";
+import { DEFAULT_FETCH_TIMEOUT_MS } from "../constants.ts";
+import { safeEnrichModelsWithModelsDev } from "../lib/model-metadata.ts";
+import type { ProviderModelConfig } from "../lib/types.ts";
 import { fetchWithRetry, mapOpenRouterModel } from "../lib/util.ts";
 
 interface OpenRouterCompatibleModel {
@@ -22,12 +23,17 @@ interface OpenRouterCompatibleModel {
 		input_modalities?: string[] | null;
 		output_modalities?: string[] | null;
 	};
-	top_provider?: { max_completion_tokens?: number | null };
+	top_provider?: {
+		context_length?: number | null;
+		max_completion_tokens?: number | null;
+	};
 	supported_parameters?: string[];
 	isFree?: boolean;
 }
 
 interface FetchModelsOptions {
+	/** Provider id for scoped models.dev enrichment (e.g., openrouter, kilo). */
+	providerId?: string;
 	/** Base URL for the API (e.g., https://api.openrouter.ai/api/v1) */
 	baseUrl: string;
 	/** API key for authentication (optional) */
@@ -93,11 +99,13 @@ export async function fetchOpenRouterCompatibleModels(
 		throw new Error("Invalid models response: missing data array");
 	}
 
-	return json.data
+	const models = json.data
 		.filter((m) => {
-			// Filter out image generation models
+			// Filter out models that cannot produce text output (image/video/audio
+			// generation, embedding-only, etc.). Keep models with no output
+			// modality info to avoid over-filtering older endpoints.
 			const outputMods = m.architecture?.output_modalities ?? [];
-			if (outputMods.includes("image")) return false;
+			if (outputMods.length > 0 && !outputMods.includes("text")) return false;
 
 			// Filter by provider flag when available, otherwise pricing.
 			if (freeOnly) {
@@ -110,6 +118,10 @@ export async function fetchOpenRouterCompatibleModels(
 			return true;
 		})
 		.map(mapOpenRouterModel);
+
+	return await safeEnrichModelsWithModelsDev(models, {
+		providerId: options.providerId,
+	});
 }
 
 /**
@@ -131,52 +143,3 @@ export async function fetchOpenRouterModelsWithFree(
 
 	return { free, all };
 }
-
-// =============================================================================
-// Models.dev metadata fetching
-// =============================================================================
-
-interface ModelsDevResponse {
-	[id: string]: {
-		id?: string;
-		models?: Record<string, ModelsDevModel>;
-	};
-}
-
-/**
- * Fetch model metadata from models.dev.
- * @param providerId - If specified, only return models for that provider
- * @returns Map of model ID to model metadata
- */
-export async function fetchModelsDevMeta(
-	providerId?: string,
-): Promise<Record<string, ModelsDevModel>> {
-	const response = await fetchWithRetry(
-		URL_MODELS_DEV,
-		{
-			headers: { "User-Agent": "pi-free-providers" },
-		},
-		3,
-		1000,
-		DEFAULT_FETCH_TIMEOUT_MS,
-	);
-
-	if (!response.ok) return {};
-
-	const json = (await response.json()) as ModelsDevResponse;
-
-	// If providerId specified, return only that provider's models
-	if (providerId) {
-		const provider = Object.values(json).find((p) => p?.id === providerId);
-		return provider?.models ?? {};
-	}
-
-	// Otherwise, return all models from all providers
-	const allModels: Record<string, ModelsDevModel> = {};
-	for (const provider of Object.values(json)) {
-		if (provider?.models) {
-			Object.assign(allModels, provider.models);
-		}
-	}
-	return allModels;
-}

package/providers/novita/novita.ts

@@ -32,12 +32,15 @@ import {
 	PROVIDER_NOVITA,
 } from "../../constants.ts";
 import { createLogger } from "../../lib/logger.ts";
+import { safeEnrichModelsWithModelsDev } from "../../lib/model-metadata.ts";
 import {
 	getProxyModelCompat,
 	isLikelyReasoningModel,
 } from "../../lib/provider-compat.ts";
+import { createProviderProbe } from "../../lib/provider-probe.ts";
 import { isFreeModel, registerWithGlobalToggle } from "../../lib/registry.ts";
-import { fetchWithRetry } from "../../lib/util.ts";
+import { wrapSessionStartHandler } from "../../lib/session-start-metrics.ts";
+import { fetchWithRetry, fetchWithTimeout } from "../../lib/util.ts";
 import { createReRegister, setupProvider } from "../../provider-helper.ts";
 
 const _logger = createLogger("novita");
@@ -96,7 +99,7 @@ async function fetchNovitaModels(
 
 		_logger.info(`[novita] Fetched ${models.length} models`);
 
-		return models.map((m): ProviderModelConfig => {
+		const mapped = models.map((m): ProviderModelConfig => {
 			const name = m.display_name || m.id.split("/").pop() || m.id;
 			const reasoning =
 				(m.features ?? []).includes("reasoning") ||
@@ -127,6 +130,10 @@ async function fetchNovitaModels(
 				_pricingKnown: hasPricing,
 			} as ProviderModelConfig & { _pricingKnown?: boolean };
 		});
+
+		return await safeEnrichModelsWithModelsDev(mapped, {
+			providerId: PROVIDER_NOVITA,
+		});
 	} catch (error) {
 		_logger.error("[novita] Failed to fetch models:", {
 			error: error instanceof Error ? error.message : String(error),
@@ -202,4 +209,64 @@ export default async function novitaProvider(pi: ExtensionAPI) {
 	const initialModels =
 		showPaid && stored.all.length > 0 ? stored.all : freeModels;
 	reRegister(initialModels);
+
+	// ── Probe support ──────────────────────────────────────────────
+	const probe = createProviderProbe({
+		providerId: PROVIDER_NOVITA,
+		probeModel: async (_apiKey: string, modelId: string) => {
+			try {
+				const response = await fetchWithTimeout(
+					`${BASE_URL_NOVITA}/chat/completions`,
+					{
+						method: "POST",
+						headers: {
+							Authorization: `Bearer ${apiKey}`,
+							"Content-Type": "application/json",
+							"User-Agent": "pi-free-providers",
+						},
+						body: JSON.stringify({
+							model: modelId,
+							messages: [{ role: "user", content: "hi" }],
+							max_tokens: 1,
+						}),
+					},
+					10_000,
+				);
+				if (response.status === 404 || response.status >= 500) return "broken";
+				if (response.status === 429) return "ok";
+				if (response.ok) return "ok";
+				return "ok";
+			} catch {
+				return "unknown";
+			}
+		},
+	});
+
+	// Probe command
+	pi.registerCommand(`probe-${PROVIDER_NOVITA}`, {
+		description: "Test all Novita AI models for availability",
+		handler: async (_args, ctx) => {
+			ctx.ui.notify(`Probing ${allModels.length} Novita AI models…`, "info");
+			const broken = await probe.run(apiKey, allModels, {
+				onBroken: (ids) => {
+					ctx.ui.notify(
+						`Found ${ids.length} broken models (auto-hidden):\n${ids.join("\n")}`,
+						"warning",
+					);
+				},
+			});
+			if (broken.length === 0) {
+				ctx.ui.notify("All Novita AI models are accessible ✅", "info");
+			}
+		},
+	});
+
+	// Lazy auto-probe on first session_start
+	pi.on(
+		"session_start",
+		wrapSessionStartHandler(
+			`${PROVIDER_NOVITA}-auto-probe`,
+			probe.autoProbeHandler(apiKey, freeModels),
+		),
+	);
 }