pi-freeflow 1.5.0 → 1.7.0

package/src/relay.ts

@@ -6,20 +6,37 @@ import { randomUUID } from "node:crypto";
 import { Buffer } from "node:buffer";
 import { Agent } from "undici";
 import { isDebugEnabled, log } from "./logger.ts";
-import { getModelDef } from "./models.ts";
 
 export const agent = new Agent({ keepAliveTimeout: 30_000 });
+/**
+ * Node's global fetch runs on the node-process-bundled undici. Node 22/23
+ * bundle undici 6.x/7.x whose dispatcher interface rejects an npm-undici 8.x
+ * Agent passed as `dispatcher` (`invalid onRequestStart method`). Only attach
+ * the custom dispatcher when the bundled undici major matches (Node 25+);
+ * otherwise fall back to the built-in global dispatcher, which still pools
+ * keep-alive connections, so the relay proxy works on Node 22/23 too.
+ */
+const bundledUndiciMajor = Number((process.versions.undici ?? "0").split(".")[0]);
+export const canUseCustomDispatcher = bundledUndiciMajor >= 8;
 import {
 	getActiveRelayState,
 	getOrderedRelayUrls,
+	getStatusUi,
 	markRelayFailure,
 	markRelaySuccess,
-	saveRelayState,
-	setActiveRelayState,
 	shortRelayLabel,
 	updateRelayStatusUi,
+	withRelayState,
+	validateRelayUrl,
 } from "./relay-state.ts";
 
+// Throttle user-facing roll notifications so a burst of failures surfaces
+// one warning instead of a wall of identical toasts.
+let lastRollNotify = 0;
+const ROLL_NOTIFY_MS = 5 * 60 * 1_000;
+/** Test-only: reset roll-notify throttle */
+export function _resetRollNotifyForTest(): void { lastRollNotify = 0; }
+
 /**
  * Determine if an HTTP status code indicates a temporary relay or upstream error
  * that warrants rolling to the next relay candidate.
@@ -34,20 +51,6 @@ export function isRetriableStatus(status: number): boolean {
 		(status >= 520 && status <= 530)
 	);
 }
-export function clampMaxTokens(modelId: string, requested: number): number {
-	const def = getModelDef(modelId);
-	const max = def ? def.maxTokens - 1024 : 32_000 - 1024;
-	// special case for big-pickle acceptance
-	const effectiveMax = modelId.includes("big-pickle") ? 32_000 : max;
-	const clamped = Math.min(requested, effectiveMax);
-	if (clamped !== requested) {
-		// logDebug
-		try { log("debug", `clamp maxTokens ${requested}->${clamped} for ${modelId}`); } catch {}
-	}
-	return clamped;
-}
-
-
 /**
  * Fetch a target URL through the active relay pool with rolling failover and direct fallback.
  *
@@ -65,8 +68,7 @@ export async function relayFetch(
 
 	if (!relayState.enabled) {
 		log("debug", `relayFetch: direct (relay disabled) -> ${url}`, undefined, rid);
-		return fetch(url, { ...opts, dispatcher: agent } as unknown as RequestInit);
-
+		return fetch(url, { ...opts, ...(canUseCustomDispatcher ? { dispatcher: agent } : {}) } as unknown as RequestInit);
 	}
 	const candidates = getOrderedRelayUrls();
 
@@ -74,7 +76,7 @@ export async function relayFetch(
 		// Empty pool: skip straight to upstream instead of logging a misleading
 		// "relays bypassed/exhausted" WARN on every request.
 		log("debug", `relayFetch: direct (empty relay pool) -> ${url}`, undefined, rid);
-		return fetch(url, { ...opts, dispatcher: agent } as unknown as RequestInit);
+		return fetch(url, { ...opts, ...(canUseCustomDispatcher ? { dispatcher: agent } : {}) } as unknown as RequestInit);
 	}
 
 	let lastResponse: Response | null = null;
@@ -111,6 +113,18 @@ export async function relayFetch(
 		const targetUrl = candidates[i];
 		const attemptStart = Date.now();
 		try {
+			// SSRF guard: reject private/loopback/non-https candidates the same
+			// way a deployed relay worker rejects an inbound x-relay-target.
+			const candidateCheck = validateRelayUrl(targetUrl);
+			if (!candidateCheck.ok) {
+				log(
+					"warn",
+					`relay ${targetUrl} skipped — ${candidateCheck.reason}`,
+					{ upstream: url },
+					rid,
+				);
+				continue;
+			}
 			let targetHost = "opencode.ai";
 			try {
 				if (targetUrl) targetHost = new URL(targetUrl).host;
@@ -121,9 +135,17 @@ export async function relayFetch(
 			headers.set("x-relay-path", relayPath);
 			headers.set("host", targetHost);
 			headers.set("x-request-id", rid);
+			// Per-relay shared secret set by /freeflow deploy. Legacy entries
+			// without auth keep working: no header at all.
+			const entry = getActiveRelayState().relays.find(
+				(r) => r.url === targetUrl.trim(),
+			);
+			if (entry?.auth) {
+				headers.set("x-relay-auth", entry.auth);
+			}
 
 			const signal = opts.signal || AbortSignal.timeout(300_000);
-			const res = await fetch(targetUrl, { ...opts, headers, signal, dispatcher: agent } as unknown as RequestInit);
+			const res = await fetch(targetUrl, { ...opts, headers, signal, ...(canUseCustomDispatcher ? { dispatcher: agent } : {}) } as unknown as RequestInit);
 			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.
@@ -136,6 +158,14 @@ export async function relayFetch(
 					{ upstream: url, sizeKB: bodySizeKB },
 					rid,
 				);
+				const now = Date.now();
+				if (now - lastRollNotify > ROLL_NOTIFY_MS) {
+					lastRollNotify = now;
+					const ui = getStatusUi();
+					if (ui?.notify) {
+						ui.notify(`relay ${shortRelayLabel(targetUrl)} hit HTTP 504 — falling back to direct`, "warning");
+					}
+				}
 				break;
 			}
 
@@ -149,10 +179,18 @@ export async function relayFetch(
 					{ upstream: url, status: res.status },
 					rid,
 				);
+				const now = Date.now();
+				if (now - lastRollNotify > ROLL_NOTIFY_MS) {
+					lastRollNotify = now;
+					const ui = getStatusUi();
+					if (ui?.notify) {
+						ui.notify(`relay ${shortRelayLabel(targetUrl)} failed (HTTP ${res.status}) — rolled to next relay`, "warning");
+					}
+				}
 				continue;
 			}
 
-			markRelaySuccess(targetUrl);
+			markRelaySuccess(targetUrl, Date.now() - attemptStart);
 
 			// SUCCESS or non-retriable client error (e.g. 200, 404):
 			// If we switched to a different relay because previous failed, update sticky active relay!
@@ -160,9 +198,12 @@ export async function relayFetch(
 				log("info", `active relay auto-switched to ${targetUrl}`, {
 					previous: relayState.url,
 				}, rid);
-				relayState.url = targetUrl;
-				saveRelayState(relayState);
-				setActiveRelayState(relayState, false);
+				// CAS: re-apply the sticky-active switch to the freshest disk state at
+				// write time so a concurrent session's pool edit is never clobbered.
+				withRelayState((s) => {
+					s.url = targetUrl;
+					return s;
+				});
 			}
 
 			log("info", `relay ${targetUrl} succeeded (HTTP ${res.status} in ${elapsed}s)`, undefined, rid);
@@ -177,6 +218,11 @@ export async function relayFetch(
 			updateRelayStatusUi(targetUrl);
 			return res;
 		} catch (err) {
+			// Client abort: do not mark the relay failed — the client cancelled the
+			// request, the relay itself is not at fault. Propagate immediately.
+			if ((err as Error)?.name === "AbortError") {
+				throw err;
+			}
 			const elapsed = ((Date.now() - attemptStart) / 1000).toFixed(1);
 			lastError = err;
 			const errMsg = (err as Error)?.message || String(err);
@@ -205,7 +251,7 @@ export async function relayFetch(
 		directHeaders.set("host", u.host);
 		directHeaders.set("x-request-id", rid);
 
-		const directRes = await fetch(url, { ...opts, headers: directHeaders, dispatcher: agent } as unknown as RequestInit);
+		const directRes = await fetch(url, { ...opts, headers: directHeaders, ...(canUseCustomDispatcher ? { dispatcher: agent } : {}) } as unknown as RequestInit);
 		// lastResponse holds an unread body that would otherwise leak its socket
 		// until GC; the salvage path below still needs it, so only cancel here.
 		lastResponse?.body?.cancel().catch(() => {});