pi-freeflow 1.2.0 → 1.3.0

package/src/relay-state.ts

@@ -0,0 +1,218 @@
+/**
+ * Persistent relay state manager and failover ordering for pi-freeflow
+ *
+ * Handles atomic state writes to ~/.pi/agent/pi-freeflow-relay-state.json.
+ */
+
+import { randomUUID } from "node:crypto";
+import fs from "node:fs";
+import path from "node:path";
+import { DEFAULT_RELAY_URL, RELAY_STATE_FILE } from "./config.ts";
+import { logWarn } from "./logger.ts";
+import type { ExtensionUIContext, KnownRelay, RelayState } from "./types.ts";
+
+/**
+ * Load persisted relay state from disk.
+ */
+export function loadRelayState(): RelayState {
+	try {
+		if (!fs.existsSync(RELAY_STATE_FILE)) {
+			return { enabled: true, url: DEFAULT_RELAY_URL, relays: [] };
+		}
+		const s = JSON.parse(fs.readFileSync(RELAY_STATE_FILE, "utf8"));
+		const relays: KnownRelay[] = Array.isArray(s?.relays) ? s.relays : [];
+		// Auto-on by default if saved relays exist, unless explicitly set to false
+		const enabled = s?.enabled !== undefined ? Boolean(s.enabled) : relays.length > 0;
+		return {
+			enabled,
+			url: typeof s?.url === "string" ? s.url.trim() : (relays[0]?.url || DEFAULT_RELAY_URL),
+			relays,
+		};
+	} catch {
+		return { enabled: true, url: DEFAULT_RELAY_URL, relays: [] };
+	}
+}
+
+/**
+ * Atomically save relay state to disk using a temporary file and rename.
+ */
+export function saveRelayState(s: RelayState): void {
+	try {
+		const dir = path.dirname(RELAY_STATE_FILE);
+		if (!fs.existsSync(dir)) {
+			fs.mkdirSync(dir, { recursive: true });
+		}
+		const tmpPath = `${RELAY_STATE_FILE}.${randomUUID()}.tmp`;
+		fs.writeFileSync(tmpPath, JSON.stringify(s, null, 2), "utf8");
+		fs.renameSync(tmpPath, RELAY_STATE_FILE);
+	} catch (e) {
+		logWarn("Could not persist relay state", { error: String(e) });
+	}
+}
+
+/**
+ * Deduplicate and add a relay URL to the known relay list.
+ */
+export function ensureRelay(s: RelayState, url: string, label?: string): void {
+	if (!url || s.relays.some((r) => r.url === url)) {
+		return;
+	}
+	s.relays.push({ url, label, addedAt: new Date().toISOString() });
+}
+
+/**
+ * Remove a relay URL from the known relay list.
+ */
+export function removeRelay(s: RelayState, url: string): void {
+	s.relays = s.relays.filter((r) => r.url !== url);
+}
+
+/**
+ * Resolve relay state with defaults.
+ */
+export function resolveRelayState(): RelayState {
+	const s = loadRelayState();
+	// Seed the relay list from defaults when state lacks entries.
+	if (!s.relays.length) {
+		if (DEFAULT_RELAY_URL) {
+			ensureRelay(s, DEFAULT_RELAY_URL, "Default");
+		}
+		if (s.url && s.url !== DEFAULT_RELAY_URL) {
+			ensureRelay(s, s.url, "previous");
+		}
+	}
+	if (!s.url) {
+		s.url = DEFAULT_RELAY_URL;
+	}
+	return s;
+}
+
+// In-memory global relay state
+let activeRelayState: RelayState = resolveRelayState();
+// Monotonic counter to distribute primary relay across concurrent subagents
+let roundRobinCounter = 0;
+let activeStatusUi: ExtensionUIContext | null = null;
+/**
+ * Mtime of the on-disk state file at the moment we last read or wrote it.
+ * Lets worker processes pick up relay-pool changes persisted by another
+ * session's master daemon, while never clobbering this process's own
+ * unpersisted runtime overrides between external writes.
+ */
+let lastKnownStateMtimeMs = -1;
+
+function currentDiskStateMtimeMs(): number {
+	try {
+		return fs.statSync(RELAY_STATE_FILE).mtimeMs;
+	} catch {
+		return -1;
+	}
+}
+
+lastKnownStateMtimeMs = currentDiskStateMtimeMs();
+
+/**
+ * Get active in-memory relay state
+ */
+export function getActiveRelayState(): RelayState {
+	const currentMtime = currentDiskStateMtimeMs();
+	if (currentMtime > 0 && currentMtime > lastKnownStateMtimeMs) {
+		activeRelayState = resolveRelayState();
+		lastKnownStateMtimeMs = currentMtime;
+	}
+	return activeRelayState;
+}
+
+/**
+ * Set active in-memory relay state and persist to disk
+ */
+export function setActiveRelayState(s: RelayState, persist = true): void {
+	activeRelayState = s;
+	if (persist) {
+		saveRelayState(s);
+	}
+	lastKnownStateMtimeMs = currentDiskStateMtimeMs();
+}
+
+/**
+ * Register the Pi extension UI context for status line updates
+ */
+export function setStatusUi(ui: ExtensionUIContext | null): void {
+	activeStatusUi = ui;
+}
+
+/**
+ * Get the current Pi extension UI context
+ */
+export function getStatusUi(): ExtensionUIContext | null {
+	return activeStatusUi;
+}
+
+/**
+ * Generate a short, human-readable label for a relay URL.
+ */
+export function shortRelayLabel(url: string, relays?: KnownRelay[]): string {
+	const pool = relays || activeRelayState.relays;
+	try {
+		const hit = pool.find((r) => r.url === url);
+		if (hit?.label) return hit.label;
+		return new URL(url).host.split(".")[0];
+	} catch {
+		return url.slice(0, 18);
+	}
+}
+
+/**
+ * Get ordered candidate URLs for relay execution starting with the sticky active URL.
+ * Reloads from disk only when another process changed the state file (mtime moved),
+ * so cross-session relay-pool updates propagate to workers while this process's own
+ * unpersisted runtime overrides survive between external writes.
+ */
+export function getOrderedRelayUrls(): string[] {
+	const mtime = currentDiskStateMtimeMs();
+	if (mtime !== lastKnownStateMtimeMs) {
+		activeRelayState = loadRelayState();
+		lastKnownStateMtimeMs = mtime;
+	}
+
+	if (activeRelayState.relays && activeRelayState.relays.length > 0) {
+		const active = (activeRelayState.url || "").trim();
+		let activeIdx = activeRelayState.relays.findIndex((r) => r.url === active);
+		if (activeIdx < 0) {
+			activeIdx = 0;
+		}
+		// Rotate starting point per-request to avoid thundering herd when many
+		// subagents hit the shared 127.0.0.1 daemon at once — each request
+		// tries a different primary relay, but still rolls seamlessly on 429.
+		const startIdx = (activeIdx + (roundRobinCounter++ % activeRelayState.relays.length)) % activeRelayState.relays.length;
+		const ordered: string[] = [];
+		for (let i = 0; i < activeRelayState.relays.length; i++) {
+			const r = activeRelayState.relays[(startIdx + i) % activeRelayState.relays.length];
+			if (r?.url?.trim()) {
+				ordered.push(r.url.trim());
+			}
+		}
+		return ordered.length > 0 ? ordered : [DEFAULT_RELAY_URL];
+	}
+	if (activeRelayState.url?.trim()) {
+		return [activeRelayState.url.trim()];
+	}
+	return [DEFAULT_RELAY_URL];
+}
+
+/**
+ * Update the extension UI status widget with current active relay info
+ */
+export function updateRelayStatusUi(targetUrl?: string): void {
+	if (!activeStatusUi?.setStatus) {
+		return;
+	}
+	const currentUrl = targetUrl || activeRelayState.url;
+	if (!activeRelayState.enabled || !currentUrl) {
+		activeStatusUi.setStatus("freeflow", activeRelayState.enabled ? "relay: direct" : "relay: OFF");
+		return;
+	}
+	const label = shortRelayLabel(currentUrl);
+	const total = activeRelayState.relays.length || 1;
+	const pos = Math.max(1, activeRelayState.relays.findIndex((r) => r.url === currentUrl) + 1);
+	activeStatusUi.setStatus("freeflow", `relay: ON | ${label} ${pos}/${total}`);
+}

package/src/relay.ts

@@ -0,0 +1,193 @@
+/**
+ * High-resiliency multi-cloud relay client and failover dispatcher for pi-freeflow
+ *
+ * Implements sticky preference, rolling failover across distributed egress relays,
+ * 25-second fast fallback for Vercel 504 timeouts, and automatic direct upstream execution.
+ */
+
+import { randomUUID } from "node:crypto";
+import { Buffer } from "node:buffer";
+import { isDebugEnabled, log } from "./logger.ts";
+import {
+	getActiveRelayState,
+	getOrderedRelayUrls,
+	saveRelayState,
+	setActiveRelayState,
+	shortRelayLabel,
+	updateRelayStatusUi,
+} from "./relay-state.ts";
+
+/**
+ * Determine if an HTTP status code indicates a temporary relay or upstream error
+ * that warrants rolling to the next relay candidate.
+ */
+export function isRetriableStatus(status: number): boolean {
+	return (
+		status === 429 ||
+		status === 408 ||
+		status === 500 ||
+		status === 502 ||
+		status === 503 ||
+		status === 504 ||
+		(status >= 520 && status <= 530)
+	);
+}
+
+/**
+ * Fetch a target URL through the active relay pool with rolling failover and direct fallback.
+ *
+ * @param url Full upstream destination URL (e.g. https://opencode.ai/zen/v1/chat/completions)
+ * @param opts Standard fetch RequestInit options
+ * @param reqId Optional correlation request ID for end-to-end tracing
+ */
+export async function relayFetch(
+	url: string,
+	opts: RequestInit = {},
+	reqId?: string,
+): Promise<Response> {
+	const rid = reqId || randomUUID().slice(0, 8);
+	const relayState = getActiveRelayState();
+
+	if (!relayState.enabled) {
+		log("debug", `relayFetch: direct (relay disabled) -> ${url}`, undefined, rid);
+		return fetch(url, opts);
+	}
+
+	const candidates = getOrderedRelayUrls();
+	let lastResponse: Response | null = null;
+	let lastError: unknown = null;
+	const u = new URL(url);
+	const relayTarget = `${u.protocol}//${u.host}`;
+	const relayPath = `${u.pathname}${u.search}`;
+
+	const bodySizeKB =
+		typeof opts.body === "string"
+			? (opts.body.length / 1024).toFixed(1)
+			: Buffer.isBuffer(opts.body)
+				? (opts.body.length / 1024).toFixed(1)
+				: "0";
+
+	log("info", `request starting (${bodySizeKB}KB payload) -> ${url}`, undefined, rid);
+	if (isDebugEnabled()) {
+		try {
+			const bodyPreview = typeof opts.body === "string" ? opts.body.slice(0, 1200) : "";
+			const modelMatch = bodyPreview.match(/"model"\s*:\s*"([^"]+)"/);
+			const streamMatch = bodyPreview.match(/"stream"\s*:\s*(true|false)/);
+			log("debug", "request detail", {
+				model: modelMatch?.[1],
+				stream: streamMatch?.[1],
+				sizeKB: bodySizeKB,
+				relayTarget,
+				relayPath,
+				candidates: candidates.length,
+			}, rid);
+		} catch {}
+	}
+
+	for (let i = 0; i < candidates.length; i++) {
+		const targetUrl = candidates[i];
+		const attemptStart = Date.now();
+		try {
+			let targetHost = "opencode.ai";
+			try {
+				if (targetUrl) targetHost = new URL(targetUrl).host;
+			} catch {}
+
+			const headers = new Headers(opts.headers);
+			headers.set("x-relay-target", relayTarget);
+			headers.set("x-relay-path", relayPath);
+			headers.set("host", targetHost);
+			headers.set("x-request-id", rid);
+
+			const signal = opts.signal || AbortSignal.timeout(300_000);
+			const res = await fetch(targetUrl, { ...opts, headers, signal });
+			const elapsed = ((Date.now() - attemptStart) / 1000).toFixed(1);
+
+			// Vercel 504 Gateway Timeout on heavy prompts (>50KB or >25s):
+			// Fast fallback directly to upstream instead of cycling through multiple 25s timeouts.
+			if (res.status === 504) {
+				log(
+					"warn",
+					`relay ${targetUrl} hit HTTP 504 Gateway Timeout in ${elapsed}s (prompt evaluation exceeded Vercel 25s limit) — fast fallback to direct upstream`,
+					{ upstream: url, sizeKB: bodySizeKB },
+					rid,
+				);
+				break;
+			}
+
+			if (isRetriableStatus(res.status)) {
+				lastResponse = res;
+				log(
+					"warn",
+					`relay ${targetUrl} returned HTTP ${res.status} in ${elapsed}s — rolling to next relay`,
+					{ upstream: url, status: res.status },
+					rid,
+				);
+				continue;
+			}
+
+			// SUCCESS or non-retriable client error (e.g. 200, 404):
+			// If we switched to a different relay because previous failed, update sticky active relay!
+			if (relayState.url !== targetUrl) {
+				log("info", `active relay auto-switched to ${targetUrl}`, {
+					previous: relayState.url,
+				}, rid);
+				relayState.url = targetUrl;
+				saveRelayState(relayState);
+				setActiveRelayState(relayState, false);
+			}
+
+			log("info", `relay ${targetUrl} succeeded (HTTP ${res.status} in ${elapsed}s)`, undefined, rid);
+			if (isDebugEnabled()) {
+				log("debug", "relay headers", {
+					status: res.status,
+					contentType: res.headers.get("content-type"),
+					via: res.headers.get("via") || res.headers.get("x-vercel-id") || "direct",
+				}, rid);
+			}
+
+			updateRelayStatusUi(targetUrl);
+			return res;
+		} catch (err) {
+			const elapsed = ((Date.now() - attemptStart) / 1000).toFixed(1);
+			lastError = err;
+			log(
+				"warn",
+				`relay ${targetUrl} fetch error in ${elapsed}s — rolling to next relay`,
+				{ upstream: url, error: String(err) },
+				rid,
+			);
+			continue;
+		}
+	}
+
+	// Full fallback: attempt direct fetch to upstream
+	const directStart = Date.now();
+	try {
+		log("warn", "relays bypassed/exhausted — attempting direct fetch to upstream", {
+			upstream: url,
+			sizeKB: bodySizeKB,
+		}, rid);
+
+		const directHeaders = new Headers(opts.headers);
+		directHeaders.delete("x-relay-target");
+		directHeaders.delete("x-relay-path");
+		directHeaders.set("host", u.host);
+		directHeaders.set("x-request-id", rid);
+
+		const directRes = await fetch(url, { ...opts, headers: directHeaders });
+		const directElapsed = ((Date.now() - directStart) / 1000).toFixed(1);
+		log("info", `direct fetch returned HTTP ${directRes.status} in ${directElapsed}s`, undefined, rid);
+		return directRes;
+	} catch (directErr) {
+		const directElapsed = ((Date.now() - directStart) / 1000).toFixed(1);
+		log("error", `direct fallback also failed in ${directElapsed}s`, {
+			upstream: url,
+			error: String(directErr),
+		}, rid);
+		if (lastResponse) {
+			return lastResponse;
+		}
+		throw directErr || lastError;
+	}
+}

package/src/stream-pipe.ts

@@ -0,0 +1,154 @@
+/**
+ * Upstream SSE streaming pipeline with thinking sniffing and exception isolation.
+ *
+ * Prevents upstream stream aborts or timeouts from crashing the host process.
+ * Sniffs thinking and reasoning chunks for debug trace logging without payload mutation.
+ */
+
+import { randomUUID } from "node:crypto";
+import type * as http from "node:http";
+import type { Readable } from "node:stream";
+import { isDebugEnabled, log } from "./logger.ts";
+
+/**
+ * Pipes an upstream readable stream to a client HTTP response.
+ *
+ * - Flushes HTTP headers immediately for low Time-to-First-Byte (TTFB).
+ * - Detects thinking/reasoning tokens in chunks for diagnostic logs.
+ * - Handles upstream errors (returns 502 if headers unsent) and ends response.
+ * - Cleans up resources and destroys upstream stream on client disconnect/abort.
+ */
+export function pipeUpstreamStream(
+	nodeStream: Readable,
+	res: http.ServerResponse,
+	req: http.IncomingMessage,
+	reqId?: string,
+): void {
+	const rid = reqId || randomUUID().slice(0, 8);
+	let totalChunks = 0;
+	let totalBytes = 0;
+	let thinkingChunks = 0;
+	let thinkingBytes = 0;
+	let firstChunkAt: number | null = null;
+	const startAt = Date.now();
+
+	const sniffThinking = (chunk: Buffer | string): boolean => {
+		const s =
+			typeof chunk === "string"
+				? chunk
+				: chunk.toString("utf8", 0, Math.min(chunk.length, 4000));
+		return (
+			s.includes("reasoning") ||
+			s.includes("thinking") ||
+			s.includes("<think>") ||
+			s.includes("reasoning_content") ||
+			s.includes('"type":"thinking"') ||
+			s.includes("thinking_delta")
+		);
+	};
+
+	try {
+		if (typeof res.flushHeaders === "function") {
+			res.flushHeaders();
+		}
+	} catch {}
+
+	nodeStream.on("data", (chunk: Buffer | string) => {
+		try {
+			if (firstChunkAt === null) {
+				firstChunkAt = Date.now();
+				const ttfb = firstChunkAt - startAt;
+				log("debug", `stream first chunk in ${ttfb}ms`, undefined, rid);
+			}
+			totalChunks++;
+			const chunkSize =
+				typeof chunk === "string" ? Buffer.byteLength(chunk) : chunk.length;
+			totalBytes += chunkSize;
+
+			if (sniffThinking(chunk)) {
+				thinkingChunks++;
+				thinkingBytes += chunkSize;
+				if (isDebugEnabled() && thinkingChunks <= 3) {
+					const preview =
+						typeof chunk === "string"
+							? chunk.slice(0, 600)
+							: chunk.toString("utf8", 0, 600);
+					log(
+						"debug",
+						`thinking chunk #${thinkingChunks}`,
+						{ preview: preview.slice(0, 400) },
+						rid,
+					);
+				}
+			}
+
+			res.write(chunk);
+			const maybeFlush = res as unknown as { flush?: () => void };
+			if (typeof maybeFlush.flush === "function") {
+				maybeFlush.flush();
+			}
+		} catch {}
+	});
+
+	nodeStream.on("error", (e: unknown) => {
+		log(
+			"error",
+			"upstream stream error",
+			{ error: String(e), totalChunks, thinkingChunks },
+			rid,
+		);
+		try {
+			if (!res.headersSent) {
+				res.writeHead(502, { "content-type": "application/json" });
+			}
+			if (!res.writableEnded) {
+				res.end();
+			}
+		} catch {}
+	});
+
+	nodeStream.on("end", () => {
+		const elapsed = ((Date.now() - startAt) / 1000).toFixed(1);
+		if (thinkingChunks > 0) {
+			log(
+				"info",
+				`stream ended in ${elapsed}s — ${totalChunks} chunks (${(totalBytes / 1024).toFixed(1)}KB), thinking: ${thinkingChunks} chunks (${(thinkingBytes / 1024).toFixed(1)}KB)`,
+				undefined,
+				rid,
+			);
+		} else if (isDebugEnabled()) {
+			log(
+				"debug",
+				`stream ended in ${elapsed}s — ${totalChunks} chunks (${(totalBytes / 1024).toFixed(1)}KB), no thinking detected`,
+				undefined,
+				rid,
+			);
+		}
+		try {
+			if (!res.writableEnded) res.end();
+		} catch {}
+	});
+
+	nodeStream.on("close", () => {
+		try {
+			if (!res.writableEnded) res.end();
+		} catch {}
+	});
+
+	req.on("aborted", () => {
+		log("warn", "client aborted — destroying upstream", { totalChunks }, rid);
+		if (!nodeStream.destroyed) nodeStream.destroy();
+	});
+
+	req.on("close", () => {
+		if (!nodeStream.destroyed) nodeStream.destroy();
+	});
+
+	res.on("close", () => {
+		if (!nodeStream.destroyed) nodeStream.destroy();
+	});
+
+	res.on("error", () => {
+		if (!nodeStream.destroyed) nodeStream.destroy();
+	});
+}