pi-ext

/**
 * Passive per-model and per-backend performance, measured only from real hops.
 * Quality order stays primary: measurements shorten deadlines, demote chronic
 * offenders, and promote proven backends one place; speed never promotes.
 */

/** Weight of the newest sample in each moving average. */
const ALPHA = 0.3;
/** Failures lose half their weight over this span, so demotion wears off. */
const FAILURE_HALF_LIFE_MS = 30 * 60_000;
/** User votes (manual cools) are judgements, so they fade over days. */
const VOTE_HALF_LIFE_MS = 7 * 24 * 60 * 60_000;
/** One vote does not demote; two recent votes do. */
const VOTE_WEIGHT = 0.35;
/** Below this many samples a model keeps its static rank and default deadline. */
const MIN_SAMPLES = 3;
/** Score above which a model or backend sorts behind its peers. */
const CHRONIC_FAILURE = 0.5;
/** A backend this reliable over this many requests may move up one place. */
const PROMOTE_SAMPLES = 20;
const PROMOTE_MAX_FAILURE = 0.1;
/** ...but only past a neighbor at least this unreliable. */
const UNRELIABLE_FAILURE = 0.3;
/** Persisted state keeps at most this many models, most recent first. */
const MAX_STORED = 500;
export const MIN_DEADLINE_MS = 3_000;
const DEADLINE_FACTOR = 3;

export type Sample =
	| { kind: "ok"; firstTokenMs: number; tokensPerSecond?: number }
	| { kind: "fail" };

interface Stat {
	failure: number;
	samples: number;
	timed: number;
	firstTokenMs?: number;
	tokensPerSecond?: number;
	at: number;
	votes?: number;
	votedAt?: number;
}

export interface ModelPerf {
	id: string;
	samples: number;
	failureRate: number;
	firstTokenMs?: number;
	tokensPerSecond?: number;
	demoted: boolean;
}

/** JSON shape of the routing state file. */
export interface PerfState {
	version: 1;
	models: Record<string, Stat>;
	backends: Record<string, Stat>;
}

const average = (previous: number | undefined, next: number) =>
	previous === undefined ? next : previous + ALPHA * (next - previous);

const finite = (value: unknown): value is number =>
	typeof value === "number" && Number.isFinite(value);

/** Accept only well-formed stats from disk; anything else is dropped. */
function validStat(value: unknown): Stat | undefined {
	const stat = value as Partial<Stat> | undefined;
	if (
		!stat ||
		!finite(stat.failure) ||
		!finite(stat.samples) ||
		!finite(stat.timed) ||
		!finite(stat.at)
	)
		return undefined;
	return {
		failure: Math.min(1, Math.max(0, stat.failure)),
		samples: Math.max(0, Math.floor(stat.samples)),
		timed: Math.max(0, Math.floor(stat.timed)),
		at: stat.at,
		...(finite(stat.firstTokenMs) ? { firstTokenMs: stat.firstTokenMs } : {}),
		...(finite(stat.tokensPerSecond)
			? { tokensPerSecond: stat.tokensPerSecond }
			: {}),
		...(finite(stat.votes) && finite(stat.votedAt)
			? { votes: Math.max(0, stat.votes), votedAt: stat.votedAt }
			: {}),
	};
}

export function createPerf(now: () => number = Date.now) {
	const models = new Map<string, Stat>();
	const backends = new Map<string, Stat>();
	let onChange: (() => void) | undefined;

	const failureOf = (stat: Stat) =>
		stat.failure * 0.5 ** ((now() - stat.at) / FAILURE_HALF_LIFE_MS);
	const votesOf = (stat: Stat) =>
		stat.votes && stat.votedAt !== undefined
			? stat.votes * 0.5 ** ((now() - stat.votedAt) / VOTE_HALF_LIFE_MS)
			: 0;

	/** Measured failures count once there is enough evidence; votes always do. */
	const score = (stat: Stat | undefined) =>
		stat === undefined
			? 0
			: (stat.samples >= MIN_SAMPLES ? failureOf(stat) : 0) + votesOf(stat);
	const chronic = (stat: Stat | undefined) => score(stat) > CHRONIC_FAILURE;

	const fresh = (): Stat => ({ failure: 0, samples: 0, timed: 0, at: now() });

	function update(map: Map<string, Stat>, key: string, sample: Sample): void {
		const stat = map.get(key) ?? fresh();
		const failure = failureOf(stat);
		stat.failure =
			failure + ALPHA * ((sample.kind === "fail" ? 1 : 0) - failure);
		stat.samples++;
		stat.at = now();
		if (sample.kind === "ok") {
			stat.timed++;
			stat.firstTokenMs = average(stat.firstTokenMs, sample.firstTokenMs);
			if (sample.tokensPerSecond !== undefined)
				stat.tokensPerSecond = average(
					stat.tokensPerSecond,
					sample.tokensPerSecond,
				);
		}
		map.set(key, stat);
	}

	/** Stable: healthy entries keep their static order ahead of chronic ones. */
	const demote = <T>(
		items: readonly T[],
		stat: (item: T) => Stat | undefined,
	) =>
		[...items].sort(
			(a, b) => Number(chronic(stat(a))) - Number(chronic(stat(b))),
		);

	const proven = (stat: Stat | undefined) =>
		stat !== undefined &&
		stat.samples >= PROMOTE_SAMPLES &&
		score(stat) < PROMOTE_MAX_FAILURE;
	const unreliable = (stat: Stat | undefined) =>
		stat !== undefined &&
		stat.samples >= MIN_SAMPLES &&
		score(stat) >= UNRELIABLE_FAILURE;

	return {
		record(backend: string, modelId: string, sample: Sample): void {
			update(models, modelId, sample);
			update(backends, backend, sample);
			onChange?.();
		},
		/** A user's manual cool: a quality judgement against the model. */
		vote(modelId: string): void {
			const stat = models.get(modelId) ?? fresh();
			stat.votes = votesOf(stat) + VOTE_WEIGHT;
			stat.votedAt = now();
			models.set(modelId, stat);
			onChange?.();
		},
		/** Three times the model's typical first token, within 3 s and `max`. */
		deadline(modelId: string, max: number): number {
			const stat = models.get(modelId);
			if (!stat?.firstTokenMs || stat.timed < MIN_SAMPLES) return max;
			return Math.min(
				max,
				Math.max(
					Math.min(MIN_DEADLINE_MS, max),
					DEADLINE_FACTOR * stat.firstTokenMs,
				),
			);
		},
		orderModels<T extends { id: string }>(candidates: readonly T[]): T[] {
			return demote(candidates, (model) => models.get(model.id));
		},
		/**
		 * Chronic backends go last; then one pass lets a proven backend overtake
		 * an unreliable neighbor, so each backend moves up at most one place.
		 */
		orderBackends<T extends string>(order: readonly T[]): T[] {
			const ordered = demote(order, (backend) => backends.get(backend));
			for (let index = 1; index < ordered.length; index++) {
				const ahead = ordered[index - 1] as T;
				const current = ordered[index] as T;
				if (
					proven(backends.get(current)) &&
					unreliable(backends.get(ahead)) &&
					!chronic(backends.get(ahead))
				) {
					ordered[index - 1] = current;
					ordered[index] = ahead;
					index++;
				}
			}
			return ordered;
		},
		/** Measured models, most used first. */
		snapshot(): ModelPerf[] {
			return [...models]
				.filter(([, stat]) => stat.samples > 0)
				.map(([id, stat]) => ({
					id,
					samples: stat.samples,
					failureRate: failureOf(stat),
					...(stat.firstTokenMs === undefined
						? {}
						: { firstTokenMs: stat.firstTokenMs }),
					...(stat.tokensPerSecond === undefined
						? {}
						: { tokensPerSecond: stat.tokensPerSecond }),
					demoted: chronic(stat),
				}))
				.sort((a, b) => b.samples - a.samples);
		},
		/** Serializable state, newest models first, bounded. */
		export(): PerfState {
			const newest = (a: [string, Stat], b: [string, Stat]) =>
				Math.max(b[1].at, b[1].votedAt ?? 0) -
				Math.max(a[1].at, a[1].votedAt ?? 0);
			return {
				version: 1,
				models: Object.fromEntries(
					[...models].sort(newest).slice(0, MAX_STORED),
				),
				backends: Object.fromEntries(backends),
			};
		},
		/** Merge persisted state; entries already measured this session win. */
		import(state: unknown): void {
			const parsed = state as Partial<PerfState> | undefined;
			if (parsed?.version !== 1) return;
			for (const [target, source] of [
				[models, parsed.models],
				[backends, parsed.backends],
			] as const) {
				if (!source || typeof source !== "object") continue;
				for (const [key, value] of Object.entries(source)) {
					const stat = validStat(value);
					if (stat && !target.has(key)) target.set(key, stat);
				}
			}
		},
		/** Called after every change, e.g. to schedule a save. */
		onChange(listener: () => void): void {
			onChange = listener;
		},
	};
}