pi-freeflow 1.3.3 → 1.3.7

package/src/commands.ts

@@ -17,10 +17,14 @@ import {
 } from "./logger.ts";
 import {
 	ensureRelay,
+	findRelay,
 	getActiveRelayState,
+	getRelayHealth,
+	isRelayHealthy,
 	removeRelay,
 	saveRelayState,
 	setActiveRelayState,
+	setRelayLabel,
 	setStatusUi,
 	shortRelayLabel,
 } from "./relay-state.ts";
@@ -48,7 +52,10 @@ export function updateStatusBar(ui?: ExtensionUIContext): void {
 			relayState.relays.findIndex((r) => r.url === relayState.url) + 1,
 		);
 		const total = relayState.relays.length;
-		ui.setStatus("freeflow", `relay: ON | ${label} ${idx}/${total}`);
+		const modeLabel = relayState.mode === "on" ? "ON" : "AUTO (ON)";
+		ui.setStatus("freeflow", `relay: ${modeLabel} | ${label} ${idx}/${total}`);
+	} else if (relayState.mode === "off" || !relayState.enabled) {
+		ui.setStatus("freeflow", "relay: OFF (direct)");
 	} else {
 		ui.setStatus("freeflow", undefined);
 	}
@@ -60,17 +67,21 @@ export function createCommandSpec(
 ): Omit<RegisteredCommand, "name"> {
 	return {
 		description:
-			"Relay egress: on | off | status | logs [level] [n] | debug on|off|status | url [URL] | deploy | list | use <URL> | remove <URL> | refresh",
+			"Relay egress: auto | on | off | status | add <URL> [name] | list | use <URL|name|index> [name] | label <target> <name> | remove <target> | logs [level] [n] | debug on|off | refresh | deploy",
 		getArgumentCompletions: (prefix: string) =>
 			[
+				"auto",
 				"on",
 				"off",
 				"status",
-				"url",
-				"deploy",
+				"add",
 				"list",
 				"use",
+				"label",
+				"rename",
 				"remove",
+				"url",
+				"deploy",
 				"refresh",
 				"models",
 				"logs",
@@ -88,37 +99,100 @@ export function createCommandSpec(
 			const relayState = getActiveRelayState();
 
 			const flash = () => {
-				const activeLabel = shortRelayLabel(relayState.url);
+				const activeLabel = shortRelayLabel(relayState.url, relayState.relays);
 				const activeIdx = Math.max(
 					1,
 					relayState.relays.findIndex((r) => r.url === relayState.url) + 1,
 				);
 				const total = relayState.relays.length || 1;
+				const modeStr = (relayState.mode || "auto").toUpperCase();
 				ctx.ui.notify(
-					`Relay ${relayState.enabled ? "ON" : "OFF"}${relayState.enabled ? ` → ${activeLabel} (${activeIdx}/${total})` : " (direct)"} | saved=${relayState.relays.length} (auto-fallback rolling)`,
+					`Relay mode: ${modeStr} (${relayState.enabled ? "ON" : "OFF"})${relayState.enabled ? ` → ${activeLabel} (${activeIdx}/${total})` : " (direct)"} | saved=${relayState.relays.length} (auto-fallback rolling)`,
 					"info",
 				);
 			};
 
 			const persist = () => {
-				saveRelayState(relayState);
+				setActiveRelayState(relayState, true);
 				setStatusUi(ctx.ui);
 				updateStatusBar(ctx.ui);
 			};
-
 			const setRelay = (
 				enabled: boolean,
 				url: string,
 				addLabel?: string,
 			) => {
 				relayState.enabled = enabled;
-				relayState.url = (url || "").trim() || DEFAULT_RELAY_URL;
-				if (relayState.url) {
-					ensureRelay(relayState, relayState.url, addLabel);
+				const cleanUrl = (url || "").trim() || DEFAULT_RELAY_URL;
+				relayState.url = cleanUrl;
+				if (cleanUrl) {
+					ensureRelay(relayState, cleanUrl, addLabel);
 				}
 				setActiveRelayState(relayState);
 			};
 
+			const addRelayInteractive = async () => {
+				const inputUrl = (
+					await ctx.ui.input(
+						"Custom relay URL (e.g. https://my-relay.workers.dev):",
+						"",
+					)
+				)?.trim();
+				if (!inputUrl) {
+					ctx.ui.notify("Cancelled — no URL provided", "warning");
+					return;
+				}
+				const defaultLabel = shortRelayLabel(inputUrl, relayState.relays);
+				const inputLabel = (
+					await ctx.ui.input(
+						"Short name / alias for TUI (e.g. cf-sg, vercel-1):",
+						defaultLabel,
+					)
+				)?.trim() || defaultLabel;
+
+				const added = ensureRelay(relayState, inputUrl, inputLabel);
+				setRelay(true, added.url, added.label);
+				persist();
+				flash();
+				ctx.ui.notify(
+					`✓ Added & activated relay [${added.label || shortRelayLabel(added.url)}]: ${added.url}`,
+					"info",
+				);
+			};
+
+			const editRelayLabelMenu = async () => {
+				if (!relayState.relays.length) {
+					ctx.ui.notify("No saved relays to rename", "warning");
+					return;
+				}
+				const fmt = (r: KnownRelay, idx: number) =>
+					`[${idx + 1}] ${r.label ? `[${r.label}] ` : ""}${r.url}`;
+				const opts = relayState.relays.map(fmt);
+				const choice = await ctx.ui.select("Select relay to rename / set short name", opts);
+				if (!choice) return;
+				const match = relayState.relays.find((r, idx) => fmt(r, idx) === choice);
+				if (!match) return;
+
+				const curLabel = match.label || shortRelayLabel(match.url, relayState.relays);
+				const newLabel = (
+					await ctx.ui.input(
+						`New short name for ${match.url}:`,
+						curLabel,
+					)
+				)?.trim();
+				if (!newLabel) {
+					ctx.ui.notify("Cancelled — short name not changed", "warning");
+					return;
+				}
+				setRelayLabel(relayState, match.url, newLabel);
+				persist();
+				flash();
+				ctx.ui.notify(
+					`✓ Relay short name updated to [${newLabel}] for ${match.url}`,
+					"info",
+				);
+			};
+
 			const doDeploy = async () => {
 				const defaultName = `relay-${Date.now().toString(36)}`;
 				const token = (
@@ -158,31 +232,46 @@ export function createCommandSpec(
 					ctx.ui.notify("No saved relays yet", "warning");
 					return;
 				}
-				const fmt = (r: KnownRelay) =>
-					`${r.url === relayState.url ? "★ " : "  "}${r.url}${r.label ? `  (${r.label})` : ""}`;
+				const fmt = (r: KnownRelay, idx: number) => {
+					const isAct = r.url === relayState.url ? "★ " : "  ";
+					const lbl = r.label ? `[${r.label}] ` : `[${shortRelayLabel(r.url, relayState.relays)}] `;
+					return `${isAct}[${idx + 1}] ${lbl}→ ${r.url}`;
+				};
 				const opts = relayState.relays.map(fmt);
-				const choice = await ctx.ui.select("Switch relay", opts);
+				const choice = await ctx.ui.select("Switch active relay", opts);
 				if (!choice) return;
-				const match = relayState.relays.find((r) => fmt(r) === choice);
+				const match = relayState.relays.find((r, idx) => fmt(r, idx) === choice);
 				if (!match) return;
-				setRelay(true, match.url);
+				setRelay(true, match.url, match.label);
 				persist();
 				flash();
 			};
 
 			const showList = () => {
 				if (!relayState.relays.length) {
-					ctx.ui.notify("No saved relays", "info");
+					ctx.ui.notify("No saved relays (direct upstream mode)", "info");
 					return;
 				}
-				const lines = relayState.relays.map(
-					(r) =>
-						`${r.url === relayState.url ? "★" : " "} ${r.url}${r.label ? `  [${r.label}]` : ""}`,
-				);
-				ctx.ui.notify(
-					`Saved relays (${relayState.relays.length}):\n${lines.join("\n")}`,
-					"info",
+				const lines = relayState.relays.map((r, idx) => {
+					const star = r.url === relayState.url ? "★" : " ";
+					const shortName = r.label ? `[${r.label}]` : `[${shortRelayLabel(r.url, relayState.relays)}]`;
+					const paddedName = shortName.padEnd(16, " ");
+					const health = getRelayHealth(r.url);
+					const isCooling = health && Date.now() < health.cooldownUntil;
+					const remainingSec = isCooling ? Math.ceil((health.cooldownUntil - Date.now()) / 1000) : 0;
+					const healthBadge = isCooling
+						? ` ⚠️ [cooling ${remainingSec}s: ${health.lastStatus ? `HTTP ${health.lastStatus}` : "error"}]`
+						: " ✓";
+					return `${star} [${idx + 1}] ${paddedName} → ${r.url}${healthBadge}`;
+				});
+				const activeLabel = shortRelayLabel(relayState.url, relayState.relays);
+				const activeIdx = Math.max(
+					1,
+					relayState.relays.findIndex((r) => r.url === relayState.url) + 1,
 				);
+				const modeStr = (relayState.mode || "auto").toUpperCase();
+				const header = `Saved Relays (${relayState.relays.length}):\n${lines.join("\n")}\n\nActive: [${activeIdx}] ${activeLabel} | Mode: ${modeStr} (${relayState.enabled ? "ON" : "OFF"})`;
+				ctx.ui.notify(header, "info");
 			};
 
 			const removeRelayMenu = async () => {
@@ -196,25 +285,35 @@ export function createCommandSpec(
 					);
 					return;
 				}
-				const fmt = (r: KnownRelay) =>
-					`${r.url}${r.label ? `  (${r.label})` : ""}`;
+				const fmt = (r: KnownRelay, idx: number) => {
+					const lbl = r.label ? `[${r.label}] ` : `[${shortRelayLabel(r.url, relayState.relays)}] `;
+					return `[${idx + 1}] ${lbl}→ ${r.url}`;
+				};
 				const choice = await ctx.ui.select(
 					"Remove relay",
 					removable.map(fmt),
 				);
 				if (!choice) return;
-				const match = removable.find((r) => fmt(r) === choice);
+				const match = removable.find((r, idx) => fmt(r, idx) === choice);
 				if (!match) return;
 				removeRelay(relayState, match.url);
 				persist();
-				ctx.ui.notify(`Removed: ${match.url}`, "info");
+				ctx.ui.notify(`Removed: [${shortRelayLabel(match.url, relayState.relays)}] ${match.url}`, "info");
 			};
 
-			if (sub === "on") {
+			if (sub === "auto") {
+				relayState.mode = "auto";
+				relayState.enabled = true;
+				setRelay(true, relayState.url || DEFAULT_RELAY_URL);
+				persist();
+				flash();
+			} else if (sub === "on") {
+				relayState.mode = "on";
 				setRelay(true, relayState.url || DEFAULT_RELAY_URL);
 				persist();
 				flash();
 			} else if (sub === "off") {
+				relayState.mode = "off";
 				relayState.enabled = false;
 				setActiveRelayState(relayState);
 				persist();
@@ -223,19 +322,74 @@ export function createCommandSpec(
 				flash();
 			} else if (sub === "list") {
 				showList();
+			} else if (sub === "add") {
+				if (!rest.trim()) {
+					await addRelayInteractive();
+				} else {
+					const tokens = rest.trim().split(/\s+/);
+					const targetUrl = tokens[0];
+					const customLabel = tokens.slice(1).join(" ").trim() || undefined;
+					const added = ensureRelay(relayState, targetUrl, customLabel);
+					setRelay(true, added.url, added.label);
+					persist();
+					flash();
+					ctx.ui.notify(
+						`✓ Added & activated relay [${added.label || shortRelayLabel(added.url)}]: ${added.url}`,
+						"info",
+					);
+				}
 			} else if (sub === "use") {
-				const url = (
-					rest ||
-					(await ctx.ui.input("Relay URL to activate:", "")) ||
-					""
-				).trim();
-				if (!url) {
-					ctx.ui.notify("No URL given", "warning");
-					return;
+				if (!rest.trim()) {
+					await switchRelay();
+				} else {
+					const tokens = rest.trim().split(/\s+/);
+					const targetToken = tokens[0];
+					const customLabel = tokens.slice(1).join(" ").trim() || undefined;
+					const matched = findRelay(relayState, targetToken);
+					if (matched) {
+						if (customLabel) {
+							matched.label = customLabel;
+						}
+						setRelay(true, matched.url, matched.label);
+						persist();
+						flash();
+					} else {
+						// New URL
+						const added = ensureRelay(relayState, targetToken, customLabel);
+						setRelay(true, added.url, added.label);
+						persist();
+						flash();
+					}
+				}
+			} else if (sub === "label" || sub === "rename") {
+				if (!rest.trim()) {
+					await editRelayLabelMenu();
+				} else {
+					const tokens = rest.trim().split(/\s+/);
+					const targetToken = tokens[0];
+					const newLabel = tokens.slice(1).join(" ").trim();
+					const matched = findRelay(relayState, targetToken);
+					if (!matched) {
+						ctx.ui.notify(`Relay '${targetToken}' not found in saved list`, "warning");
+						return;
+					}
+					let finalLabel = newLabel;
+					if (!finalLabel) {
+						finalLabel = (
+							await ctx.ui.input(
+								`New short name / alias for ${matched.url}:`,
+								matched.label || shortRelayLabel(matched.url, relayState.relays),
+							)
+						)?.trim() || "";
+					}
+					setRelayLabel(relayState, matched.url, finalLabel);
+					persist();
+					flash();
+					ctx.ui.notify(
+						`✓ Relay short name updated to [${shortRelayLabel(matched.url, relayState.relays)}] for ${matched.url}`,
+						"info",
+					);
 				}
-				setRelay(true, url, "manual");
-				persist();
-				flash();
 			} else if (sub === "debug") {
 				const arg = rest.trim().toLowerCase();
 				if (arg === "on" || arg === "enable" || arg === "true") {
@@ -362,29 +516,26 @@ export function createCommandSpec(
 					"info",
 				);
 			} else if (sub === "remove") {
-				const url = (
-					rest ||
-					(await ctx.ui.input("Relay URL to remove:", "")) ||
-					""
-				).trim();
-				if (!url) {
-					ctx.ui.notify("No URL given", "warning");
-					return;
-				}
-				if (url === relayState.url) {
-					ctx.ui.notify(
-						"Cannot remove the active relay — switch first",
-						"warning",
-					);
-					return;
-				}
-				if (!relayState.relays.some((r) => r.url === url)) {
-					ctx.ui.notify("Not in saved list", "warning");
-					return;
+				if (!rest.trim()) {
+					await removeRelayMenu();
+				} else {
+					const target = rest.trim();
+					const matched = findRelay(relayState, target);
+					if (!matched) {
+						ctx.ui.notify(`Relay '${target}' not in saved list`, "warning");
+						return;
+					}
+					if (matched.url === relayState.url) {
+						ctx.ui.notify(
+							"Cannot remove the active relay — switch to another relay first",
+							"warning",
+						);
+						return;
+					}
+					removeRelay(relayState, matched.url);
+					persist();
+					ctx.ui.notify(`Removed: [${shortRelayLabel(matched.url, relayState.relays)}] ${matched.url}`, "info");
 				}
-				removeRelay(relayState, url);
-				persist();
-				ctx.ui.notify(`Removed: ${url}`, "info");
 			} else if (sub === "url") {
 				const input =
 					rest ||
@@ -392,55 +543,62 @@ export function createCommandSpec(
 						"Relay URL (empty = default):",
 						relayState.url || DEFAULT_RELAY_URL,
 					));
+				const cleanInput = (input || "").trim() || DEFAULT_RELAY_URL;
 				setRelay(
 					relayState.enabled,
-					(input || "").trim() || DEFAULT_RELAY_URL,
-					"manual",
+					cleanInput,
+					shortRelayLabel(cleanInput, relayState.relays),
 				);
 				persist();
 				flash();
 			} else if (sub === "deploy") {
 				await doDeploy();
 			} else {
+				const currentMode = (relayState.mode || "auto").toUpperCase();
+				const activeLabel = shortRelayLabel(relayState.url, relayState.relays);
 				const choice = await ctx.ui.select("freeflow relay", [
-					`Relay: ${relayState.enabled ? "ON" : "OFF"} → ${relayState.url || "direct"}`,
-					"Turn ON",
-					"Turn OFF",
-					"Switch relay…",
+					`Mode: ${currentMode} (${relayState.enabled ? "ON" : "OFF"}) → ${activeLabel}`,
+					"Add custom relay (URL + Short name)…",
+					"Switch active relay…",
+					"Rename / Set relay short name…",
 					"Remove relay…",
-					"Set URL",
-					"Deploy Vercel relay…",
 					"List saved relays",
+					"Deploy Vercel relay…",
+					"Mode: AUTO (auto-detect on model select)",
+					"Mode: ON (always relay)",
+					"Mode: OFF (always direct)",
 				]);
-				if (choice === "Turn ON") {
+				if (choice === `Mode: ${currentMode} (${relayState.enabled ? "ON" : "OFF"}) → ${activeLabel}`) {
+					flash();
+				} else if (choice === "Add custom relay (URL + Short name)…") {
+					await addRelayInteractive();
+				} else if (choice === "Switch active relay…") {
+					await switchRelay();
+				} else if (choice === "Rename / Set relay short name…") {
+					await editRelayLabelMenu();
+				} else if (choice === "Remove relay…") {
+					await removeRelayMenu();
+				} else if (choice === "List saved relays") {
+					showList();
+				} else if (choice === "Deploy Vercel relay…") {
+					await doDeploy();
+				} else if (choice === "Mode: AUTO (auto-detect on model select)") {
+					relayState.mode = "auto";
+					relayState.enabled = true;
 					setRelay(true, relayState.url || DEFAULT_RELAY_URL);
 					persist();
 					flash();
-				} else if (choice === "Turn OFF") {
-					relayState.enabled = false;
-					setActiveRelayState(relayState);
+				} else if (choice === "Mode: ON (always relay)") {
+					relayState.mode = "on";
+					setRelay(true, relayState.url || DEFAULT_RELAY_URL);
 					persist();
 					flash();
-				} else if (choice === "Switch relay…") {
-					await switchRelay();
-				} else if (choice === "Remove relay…") {
-					await removeRelayMenu();
-				} else if (choice === "Set URL") {
-					const input = await ctx.ui.input(
-						"Relay URL (empty = default):",
-						relayState.url || DEFAULT_RELAY_URL,
-					);
-					setRelay(
-						relayState.enabled,
-						(input || "").trim() || DEFAULT_RELAY_URL,
-						"manual",
-					);
+				} else if (choice === "Mode: OFF (always direct)") {
+					relayState.mode = "off";
+					relayState.enabled = false;
+					setActiveRelayState(relayState);
 					persist();
 					flash();
-				} else if (choice === "Deploy Vercel relay…") {
-					await doDeploy();
-				} else if (choice === "List saved relays") {
-					showList();
 				}
 			}
 		},

package/src/index.ts

@@ -18,7 +18,7 @@ import {
 import { createCommandSpec, updateStatusBar } from "./commands.ts";
 import { DEFAULT_HOST, HOST, PORT } from "./config.ts";
 import { log, logInfo, logWarn } from "./logger.ts";
-import { ALL_MODELS, KILO_MODEL_IDS } from "./models.ts";
+import { ALL_MODELS, KILO_MODEL_IDS, MODEL_MAP, getAllRegisteredModels, resolveCanonicalModelId } from "./models.ts";
 import { isProxyAlive, startProxy } from "./proxy.ts";
 import { resetRateLimits } from "./rate-limiter.ts";
 import {
@@ -26,6 +26,7 @@ import {
 	resolveRelayState,
 	setActiveRelayState,
 	setStatusUi,
+	setFreeFlowModelActive,
 } from "./relay-state.ts";
 import type {
 	ExtensionAPI,
@@ -138,14 +139,14 @@ export default async function (pi: ExtensionAPI): Promise<void> {
 
 	// 3. Instant 0ms Static Catalog Registration
 	// Register static models immediately on boot so Pi/OMP picker is populated with zero latency!
-	const initialCatalog: RegisteredModel[] = ALL_MODELS.map((m) => ({
+	const registeredCatalog: RegisteredModel[] = getAllRegisteredModels().map((m) => ({
 		...m,
 		source: KILO_MODEL_IDS.has(m.id) ? "kilo" : "opencode",
 	}));
-	setAliveCatalog(initialCatalog);
+	setAliveCatalog(ALL_MODELS.map((m) => ({ ...m, source: KILO_MODEL_IDS.has(m.id) ? "kilo" : "opencode" })));
 	pi.registerProvider(
 		"freeflow",
-		buildProviderConfig(initialCatalog, actualPort),
+		buildProviderConfig(registeredCatalog, actualPort),
 	);
 
 	// 4. Background Catalog Refresh
@@ -156,7 +157,7 @@ export default async function (pi: ExtensionAPI): Promise<void> {
 				setAliveCatalog(aliveModels);
 				pi.registerProvider(
 					"freeflow",
-					buildProviderConfig(aliveModels, actualPort),
+					buildProviderConfig(registeredCatalog, actualPort),
 				);
 				logInfo(
 					`Catalog refreshed: ${aliveModels.length} models verified active`,
@@ -179,9 +180,27 @@ export default async function (pi: ExtensionAPI): Promise<void> {
 	pi.registerCommand("freeflow", commandSpec);
 
 	// 6. Lifecycle Listeners
+	function isFreeFlowModelMatch(provider?: string, modelId?: string): boolean {
+		if (provider === "freeflow") return true;
+		if (typeof modelId === "string" && modelId.trim()) {
+			const clean = modelId.trim();
+			const canonical = resolveCanonicalModelId(clean);
+			return (
+				MODEL_MAP.has(clean) ||
+				MODEL_MAP.has(canonical) ||
+				getAliveCatalog().some((m) => m.id === clean || m.id === canonical)
+			);
+		}
+		// When provider and modelId are not yet populated on session_start,
+		// default to true so the widget is present when starting with freeflow
+		if (!provider && !modelId) {
+			return true;
+		}
+		return false;
+	}
+
 	pi.on?.("session_start", async (_event, ctx: ExtensionContext) => {
 		const freshRelayState = resolveRelayState();
-		setActiveRelayState(freshRelayState, false);
 		setStatusUi(ctx.ui);
 
 		let provider: string | undefined;
@@ -195,9 +214,13 @@ export default async function (pi: ExtensionAPI): Promise<void> {
 				modelId = m.id;
 			}
 		}
-		const isFreeFlow =
-			provider === "freeflow" ||
-			Boolean(modelId && getAliveCatalog().some((m) => m.id === modelId));
+		const isFreeFlow = isFreeFlowModelMatch(provider, modelId);
+		setFreeFlowModelActive(isFreeFlow);
+
+		if (freshRelayState.mode !== "off") {
+			freshRelayState.enabled = freshRelayState.relays.length > 0;
+			setActiveRelayState(freshRelayState, false);
+		}
 
 		if (isFreeFlow) {
 			updateStatusBar(ctx.ui);
@@ -208,7 +231,6 @@ export default async function (pi: ExtensionAPI): Promise<void> {
 
 	pi.on?.("model_select", async (event, ctx: ExtensionContext) => {
 		const freshRelayState = resolveRelayState();
-		setActiveRelayState(freshRelayState, false);
 		setStatusUi(ctx.ui);
 		let provider: string | undefined;
 		let modelId: string | undefined;
@@ -220,10 +242,22 @@ export default async function (pi: ExtensionAPI): Promise<void> {
 			if ("id" in m && typeof m.id === "string") {
 				modelId = m.id;
 			}
+		} else if (ctx && typeof ctx === "object" && "model" in ctx && ctx.model && typeof ctx.model === "object") {
+			const m = ctx.model;
+			if ("provider" in m && typeof m.provider === "string") {
+				provider = m.provider;
+			}
+			if ("id" in m && typeof m.id === "string") {
+				modelId = m.id;
+			}
+		}
+		const isFreeFlow = isFreeFlowModelMatch(provider, modelId);
+		setFreeFlowModelActive(isFreeFlow);
+
+		if (freshRelayState.mode !== "off") {
+			freshRelayState.enabled = freshRelayState.relays.length > 0;
+			setActiveRelayState(freshRelayState, false);
 		}
-		const isFreeFlow =
-			provider === "freeflow" ||
-			Boolean(modelId && getAliveCatalog().some((m) => m.id === modelId));
 
 		if (isFreeFlow) {
 			updateStatusBar(ctx.ui);
@@ -231,7 +265,6 @@ export default async function (pi: ExtensionAPI): Promise<void> {
 			ctx.ui?.setStatus?.("freeflow", undefined);
 		}
 	});
-
 	pi.on?.("session_shutdown", () => {
 		if (server) {
 			logInfo("shutting down proxy daemon...");