pi-ext

import {
	type Api,
	type AssistantMessage,
	type AssistantMessageEvent,
	type AssistantMessageEventStream,
	createAssistantMessageEventStream,
	isContextOverflow,
	type ProviderStreams,
	type SimpleStreamOptions,
	type TranscriptContext,
} from "@earendil-works/pi-ai";
// Pi's extension loader exposes API implementations only via the compat entry.
import {
	anthropicMessagesApi,
	azureOpenAIResponsesApi,
	bedrockConverseStreamApi,
	googleGenerativeAIApi,
	googleVertexApi,
	mistralConversationsApi,
	openAICodexResponsesApi,
	openAICompletionsApi,
	openAIResponsesApi,
	piMessagesApi,
} from "@earendil-works/pi-ai/compat";
import {
	AUTO_ID,
	BACKEND_INFO,
	BACKENDS,
	type Backend,
	backendOf,
	type GratisModel,
	PROVIDER,
	upstreamId,
} from "./catalog.ts";
import { span } from "./debug.ts";
import { createPerf } from "./perf.ts";

/** Every API implementation Pi ships; the catalog entry picks one per model. */
const APIS: Partial<Record<Api, ProviderStreams>> = {
	"anthropic-messages": anthropicMessagesApi(),
	"azure-openai-responses": azureOpenAIResponsesApi(),
	"bedrock-converse-stream": bedrockConverseStreamApi(),
	"google-generative-ai": googleGenerativeAIApi(),
	"google-vertex": googleVertexApi(),
	"mistral-conversations": mistralConversationsApi(),
	"openai-codex-responses": openAICodexResponsesApi(),
	"openai-completions": openAICompletionsApi(),
	"openai-responses": openAIResponsesApi(),
	"pi-messages": piMessagesApi(),
};

/** Sent instead of a key to keyless backends; the auth header is dropped. */
const ANONYMOUS_KEY = "grts-anonymous";

export const OVERFLOW_PREFIX = "context_length_exceeded";

/** `grts/auto` gives each hop this long to stream anything before moving on. */
const FIRST_RESPONSE_DEADLINE_MS = 10_000;
/** A candidate that just hit a rate limit or quota is skipped this long. */
const MODEL_COOLDOWN_MS = 60_000;
/** A silent hop means a congested backend (free queues); skip it this long. */
const SLOW_COOLDOWN_MS = 10 * 60_000;
/** A candidate the account cannot use (404, not found) is skipped this long. */
const UNAVAILABLE_COOLDOWN_MS = 60 * 60_000;

type Failure = "rate" | "quota" | "overflow" | "unavailable" | "other";

/** Listed models an account cannot use, e.g. NVIDIA's "Function ...: Not found for account". */
const UNAVAILABLE =
	/\b404\b|not found|model_not_found|does not exist|no such model|not available for/i;

const QUOTA =
	/quota|resource[_ ]exhausted|exceeded your current|per[- ]day|daily limit|credits|insufficient balance|limit exhausted/i;
const RATE =
	/\b429\b|\b503\b|rate[_ -]?limit|too many requests|unavailable|overloaded|high demand|capacity|saturat|timed? ?out|timeout|network|connection|socket hang up|fetch failed|ECONN|ETIMEDOUT/i;

/** Overflow phrasing Pi's own detector may miss on free gateways. */
const OVERFLOW =
	/context[_ ](length|window)|prompt is too long|input is too long|too many (input )?tokens|maximum (context|input)/i;

/** True when Pi would not yet treat this error as a context overflow. */
export function isUnrecognizedOverflow(message: AssistantMessage): boolean {
	const text = message.errorMessage ?? "";
	return (
		!text.includes(OVERFLOW_PREFIX) &&
		!isContextOverflow(message) &&
		OVERFLOW.test(text) &&
		!RATE.test(text)
	);
}

/** Classify by status and body text together; wrapped 429s carry the reason in text. */
function classifyFailure(
	message: AssistantMessage,
	contextWindow?: number,
): Failure {
	if (isContextOverflow(message, contextWindow)) return "overflow";
	if (isUnrecognizedOverflow(message)) return "overflow";
	const text = message.errorMessage ?? "";
	if (QUOTA.test(text)) return "quota";
	if (RATE.test(text)) return "rate";
	if (UNAVAILABLE.test(text)) return "unavailable";
	return "other";
}

interface RouterState {
	/** Resolved key, `undefined` for anonymous keyless use, `null` when unusable. */
	key(backend: Backend): string | undefined | null;
	/** `grts/auto` candidates for one backend, best first. */
	candidates(backend: Backend): readonly GratisModel[];
	resolveUrl(url: string): string;
	/** Test seam; defaults to FIRST_RESPONSE_DEADLINE_MS. */
	firstResponseDeadlineMs?: number;
	/** Test seam: transports by API; defaults to every implementation Pi ships. */
	apis?: Partial<Record<Api, ProviderStreams>>;
}

/** How a forwarded stream ended; errors after output began still reach the caller. */
type Ending = "ok" | "error" | "aborted";

type Outcome =
	| { kind: "streamed"; ending: Ending; served: string }
	| { kind: "failed"; message: AssistantMessage; failure: Failure };

/** What the most recent request did, for on-demand inspection. */
interface RouteReport {
	route: "pinned" | "auto";
	requested: string;
	/** `<backend>/<model>` that produced the response, when one did. */
	served?: string;
	/** The grts model that made the call; differs from `served` behind virtual routers. */
	candidate?: string;
	hops: number;
	outcome: Ending | "exhausted";
	notes: readonly string[];
	at: number;
}

export function createRouter(state: RouterState) {
	const coolingUntil = new Map<Backend, number>();
	const perf = createPerf();
	let lastRoute: RouteReport | undefined;

	function report(
		end: ReturnType<NonNullable<typeof span>> | undefined,
		route: RouteReport,
		backend?: Backend,
	): void {
		lastRoute = route;
		end?.("finish", {
			outcome: route.outcome,
			...(route.route === "auto" ? { hops: route.hops } : {}),
			...(backend ? { backend } : {}),
		});
	}

	const isActive = (until: number | undefined) => (until ?? 0) > Date.now();
	const cooling = (backend: Backend) => isActive(coolingUntil.get(backend));

	function upstream(
		model: GratisModel,
		requested: GratisModel,
		context: TranscriptContext,
		options: SimpleStreamOptions | undefined,
	): AssistantMessageEventStream {
		const backend = backendOf(model);
		const api = (state.apis ?? APIS)[model.api];
		const key = backend ? state.key(backend) : null;
		if (!backend || !api || key === null)
			return failedStream(
				requested,
				!api
					? `grts: no transport for API "${model.api}"`
					: `grts: backend for ${model.id} is not available`,
			);
		const id = upstreamId(model);
		// Same-model history keeps provider-native replay state such as thinking signatures.
		const messages = context.messages.map((message) =>
			message.role === "assistant" &&
			message.provider === PROVIDER &&
			message.model === model.id &&
			requested.id === model.id
				? { ...message, model: id }
				: message,
		);
		const maxTokens = Math.min(
			options?.maxTokens ?? model.maxTokens,
			model.maxTokens,
		);
		return api.streamSimple(
			{ ...model, id, baseUrl: state.resolveUrl(model.baseUrl) },
			{ ...context, messages },
			{
				...options,
				maxTokens,
				apiKey: key ?? ANONYMOUS_KEY,
				headers:
					key === undefined
						? { ...options?.headers, Authorization: null }
						: options?.headers,
			},
		);
	}

	/**
	 * Forward one hop. Before any content arrives an error is returned instead
	 * of emitted, so the cascade can retry the same prompt elsewhere.
	 */
	async function forward(
		source: AssistantMessageEventStream,
		target: AssistantMessageEventStream,
		requested: GratisModel,
		hop: GratisModel,
		holdErrors: boolean,
		onProgress?: () => void,
	): Promise<Outcome> {
		let held: AssistantMessageEvent | undefined;
		let committed = !holdErrors;
		let ending: Ending = "ok";
		const backend = backendOf(hop) ?? "unknown";
		let served = `${backend}/${upstreamId(hop)}`;
		// Passive measurement of this real hop; nothing is ever sent just to measure.
		const startedAt = performance.now();
		let firstAt: number | undefined;
		let outputTokens = 0;
		let failure: Failure | undefined;
		for await (const event of source) {
			if (event.type === "error") {
				ending = event.reason;
				failure = classifyFailure(event.error, hop.contextWindow);
			} else if (event.type === "done")
				outputTokens = event.message.usage.output;
			else if (event.type !== "start") firstAt ??= performance.now();
			const message = eventMessage(event);
			message.provider = PROVIDER;
			message.model = requested.id;
			// Only the final message: transports fill the upstream's own model id
			// (what routers like kilo-auto/free resolved to) while streaming.
			if (event.type === "done" || event.type === "error") {
				served = `${backend}/${message.responseModel ?? upstreamId(hop)}`;
				if (served === requested.id) delete message.responseModel;
				else message.responseModel = served;
			}
			if (!committed) {
				if (event.type === "start") {
					held = event;
					continue;
				}
				if (event.type === "error") {
					if (event.reason !== "aborted" && failure !== "overflow")
						perf.record(backend, hop.id, { kind: "fail" });
					return {
						kind: "failed",
						message: event.error,
						failure: failure ?? "other",
					};
				}
				committed = true;
				onProgress?.();
				if (held) target.push(held);
			}
			target.push(event);
		}
		target.end();
		if (ending === "ok" && firstAt !== undefined) {
			const streamingMs = performance.now() - firstAt;
			perf.record(backend, hop.id, {
				kind: "ok",
				firstTokenMs: firstAt - startedAt,
				...(outputTokens > 0 && streamingMs >= 50
					? { tokensPerSecond: outputTokens / (streamingMs / 1000) }
					: {}),
			});
		} else if (ending === "error" && failure !== "overflow")
			perf.record(backend, hop.id, { kind: "fail" });
		return { kind: "streamed", ending, served };
	}

	function pinned(
		model: GratisModel,
		context: TranscriptContext,
		options?: SimpleStreamOptions,
	): AssistantMessageEventStream {
		const out = createAssistantMessageEventStream();
		const end = span?.("request", { route: "pinned" });
		const source = safeUpstream(model, () =>
			upstream(model, model, context, options),
		);
		void forward(source, out, model, model, false).then(
			(outcome) =>
				report(end, {
					route: "pinned",
					requested: model.id,
					...(outcome.kind === "streamed"
						? { served: outcome.served, candidate: model.id }
						: {}),
					hops: 1,
					outcome: outcome.kind === "streamed" ? outcome.ending : "error",
					notes: [],
					at: Date.now(),
				}),
			(error: unknown) => {
				end?.("error", { outcome: "error" });
				fail(out, model, errorText(error));
			},
		);
		return out;
	}

	const modelCoolingUntil = new Map<string, number>();
	const modelCooling = (model: GratisModel) =>
		isActive(modelCoolingUntil.get(model.id));
	// User-requested cooldowns; unlike automatic ones they also skip the last-resort pass.
	const manualUntil = new Map<string, number>();
	const manuallyCooled = (model: GratisModel) =>
		isActive(manualUntil.get(model.id));
	const coolModel = (model: GratisModel, ms: number) =>
		modelCoolingUntil.set(model.id, Date.now() + ms);
	const deadlineMs =
		state.firstResponseDeadlineMs ?? FIRST_RESPONSE_DEADLINE_MS;

	function auto(
		requested: GratisModel,
		context: TranscriptContext,
		options?: SimpleStreamOptions,
	): AssistantMessageEventStream {
		const out = createAssistantMessageEventStream();
		const end = span?.("request", { route: "auto" });
		void (async () => {
			const notes: string[] = [];
			let hops = 0;
			let lastOverflow = false;
			const done = (
				outcome: RouteReport["outcome"],
				backend?: Backend,
				served?: string,
				candidate?: string,
			) =>
				report(
					end,
					{
						route: "auto",
						requested: requested.id,
						...(served ? { served } : {}),
						...(candidate ? { candidate } : {}),
						hops,
						outcome,
						notes: [...notes],
						at: Date.now(),
					},
					backend,
				);

			/** One hop; with a deadline, silence before the first output counts as slow. */
			async function attempt(
				hop: GratisModel,
				deadline: number | undefined,
			): Promise<Outcome | "slow"> {
				hops++;
				const hopAbort = new AbortController();
				let timer =
					deadline === undefined
						? undefined
						: setTimeout(() => hopAbort.abort(), deadline);
				const signal = options?.signal
					? AbortSignal.any([options.signal, hopAbort.signal])
					: hopAbort.signal;
				const source = safeUpstream(hop, () =>
					upstream(hop, requested, context, {
						...options,
						signal,
						// Cascade instead of in-hop backoff; Pi retries the whole turn.
						maxRetries: 0,
					}),
				);
				try {
					const outcome = await forward(
						source,
						out,
						requested,
						hop,
						true,
						() => {
							clearTimeout(timer);
							timer = undefined;
						},
					);
					if (
						outcome.kind === "failed" &&
						hopAbort.signal.aborted &&
						!options?.signal?.aborted
					)
						return "slow";
					return outcome;
				} finally {
					clearTimeout(timer);
				}
			}

			type Step =
				| { stop: true }
				| { stop: false; failure: Exclude<Failure, "other"> | "slow" };

			const abort = (message = "Request was aborted"): Step => {
				done("aborted");
				fail(out, requested, message, "aborted");
				return { stop: true };
			};

			/** Run one hop and settle the request when it ends it. */
			async function step(
				hop: GratisModel,
				backend: Backend,
				deadline: number | undefined,
			): Promise<Step> {
				if (options?.signal?.aborted) return abort();
				const outcome = await attempt(hop, deadline);
				if (outcome === "slow") {
					perf.record(backend, hop.id, { kind: "fail" });
					coolModel(hop, SLOW_COOLDOWN_MS);
					return { stop: false, failure: "slow" };
				}
				if (outcome.kind === "streamed") {
					done(outcome.ending, backend, outcome.served, hop.id);
					return { stop: true };
				}
				if (outcome.message.stopReason === "aborted")
					return abort(outcome.message.errorMessage);
				if (outcome.failure === "other") {
					notes.push(`${backend} error`);
					done("error", backend);
					fail(out, requested, outcome.message.errorMessage ?? "");
					return { stop: true };
				}
				if (outcome.failure === "overflow") {
					lastOverflow = true;
					return { stop: false, failure: "overflow" };
				}
				lastOverflow = false;
				if (outcome.failure === "unavailable") {
					coolModel(hop, UNAVAILABLE_COOLDOWN_MS);
					return { stop: false, failure: "unavailable" };
				}
				coolModel(hop, MODEL_COOLDOWN_MS);
				return { stop: false, failure: outcome.failure };
			}

			// Best skipped or timed-out candidate per backend, for the last-resort pass.
			const lastResort = new Map<Backend, GratisModel>();
			// Quality order, with chronically failing backends and models behind their peers.
			for (const backend of perf.orderBackends(BACKENDS)) {
				const candidates = perf
					.orderModels(state.candidates(backend))
					.filter((hop) => !manuallyCooled(hop));
				if (candidates.length === 0) continue;
				if (cooling(backend)) {
					notes.push(`${backend} cooling`);
					lastResort.set(backend, candidates[0] as GratisModel);
					continue;
				}
				let exhausted: "rate" | "quota" | undefined;
				let slow = false;
				let overflowed = false;
				for (const hop of candidates) {
					if (modelCooling(hop)) {
						if (!lastResort.has(backend)) lastResort.set(backend, hop);
						continue;
					}
					const result = await step(
						hop,
						backend,
						perf.deadline(hop.id, deadlineMs),
					);
					if (result.stop) return;
					if (result.failure === "overflow") {
						overflowed = true;
						continue;
					}
					if (result.failure === "unavailable") continue;
					if (result.failure === "slow") {
						// One silent hop means a congested backend: never wait on its siblings too.
						if (!lastResort.has(backend)) lastResort.set(backend, hop);
						slow = true;
						break;
					}
					exhausted =
						exhausted === "quota" || result.failure === "quota"
							? "quota"
							: "rate";
				}
				if (slow) {
					coolingUntil.set(backend, Date.now() + SLOW_COOLDOWN_MS);
					notes.push(`${backend} slow`);
				} else if (exhausted) {
					// Every tried candidate failed; cool for the observed limit.
					const info = BACKEND_INFO[backend];
					coolingUntil.set(
						backend,
						Date.now() +
							(exhausted === "quota"
								? info.quotaCooldownMs
								: info.rateCooldownMs),
					);
					notes.push(`${backend} exhausted (${exhausted})`);
				} else if (overflowed) {
					notes.push(`${backend} context overflow`);
				}
			}

			// Cooling windows are estimates: give skipped and slow candidates one untimed try.
			for (const [backend, hop] of lastResort) {
				const result = await step(hop, backend, undefined);
				if (result.stop) return;
			}
			if (lastResort.size > 0) notes.push("last resort failed");

			done("exhausted");
			const summary = `grts/auto: no free backend completed the request: ${notes.join(", ") || "no backend available"}`;
			fail(
				out,
				requested,
				lastOverflow ? `${OVERFLOW_PREFIX}: ${summary}` : summary,
			);
		})().catch((error: unknown) => {
			end?.("error", { outcome: "error" });
			fail(out, requested, errorText(error));
		});
		return out;
	}

	return {
		streamSimple(
			model: GratisModel,
			context: TranscriptContext,
			options?: SimpleStreamOptions,
		): AssistantMessageEventStream {
			return model.id === AUTO_ID
				? auto(model, context, options)
				: pinned(model, context, options);
		},
		/** Last route plus backends still cooling, soonest recovery first. */
		status() {
			const now = Date.now();
			const coolingBackends = [...coolingUntil]
				.filter(([, until]) => until > now)
				.sort((a, b) => a[1] - b[1])
				.map(([backend, until]) => ({ backend, until }));
			const manual = [...manualUntil]
				.filter(([, until]) => until > now)
				.sort((a, b) => a[1] - b[1])
				.map(([model, until]) => ({ model, until }));
			return {
				last: lastRoute,
				cooling: coolingBackends,
				manual,
				models: perf.snapshot(),
			};
		},
		/** Skip a grts model in `grts/auto`, including its last-resort pass, for `ms`. */
		cool(modelId: string, ms: number): void {
			manualUntil.set(modelId, Date.now() + ms);
			// Also a quality vote, so a repeatedly rejected model stays behind its peers.
			perf.vote(modelId);
		},
		/** Routing measurements for persistence across restarts. */
		learned: {
			export: () => perf.export(),
			import: (state: unknown) => perf.import(state),
			onChange: (listener: () => void) => perf.onChange(listener),
		},
		/** End manual cooldowns: one model, or all; returns how many ended. */
		uncool(modelId?: string): number {
			const now = Date.now();
			const active = [...manualUntil].filter(
				([id, until]) =>
					until > now && (modelId === undefined || id === modelId),
			);
			for (const [id] of active) manualUntil.delete(id);
			return active.length;
		},
	};
}

function safeUpstream(
	model: GratisModel,
	start: () => AssistantMessageEventStream,
): AssistantMessageEventStream {
	try {
		return start();
	} catch (error) {
		return failedStream(model, errorText(error));
	}
}

function eventMessage(event: AssistantMessageEvent): AssistantMessage {
	return "partial" in event
		? event.partial
		: event.type === "done"
			? event.message
			: event.error;
}

function fail(
	out: AssistantMessageEventStream,
	model: GratisModel,
	errorMessage: string,
	reason: "error" | "aborted" = "error",
): void {
	const message: AssistantMessage = {
		role: "assistant",
		content: [],
		api: model.api,
		provider: PROVIDER,
		model: model.id,
		usage: {
			input: 0,
			output: 0,
			cacheRead: 0,
			cacheWrite: 0,
			totalTokens: 0,
			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
		},
		stopReason: reason,
		errorMessage,
		timestamp: Date.now(),
	};
	out.push({ type: "error", reason, error: message });
	out.end();
}

function failedStream(
	model: GratisModel,
	errorMessage: string,
): AssistantMessageEventStream {
	const out = createAssistantMessageEventStream();
	fail(out, model, errorMessage);
	return out;
}

function errorText(error: unknown): string {
	return error instanceof Error ? error.message : String(error);
}