@hu3rror/pi-failover 0.2.0

package/src/index.ts

@@ -0,0 +1,346 @@
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { addBackupKey, type AddBackupKeyResult } from "./add-backup-key.ts";
+import { loadAuthCatalog, type AuthCatalog } from "./auth-catalog.ts";
+import { runFailoverLogin } from "./failover-login.ts";
+import {
+	backupIndexForSlot,
+	FailoverEngine,
+	type FailoverAttempt,
+	type FailoverDecision,
+	type FailureObservation,
+} from "./failover-engine.ts";
+import { applyNextModel } from "./model-planner.ts";
+import {
+	formatStatus,
+	notify,
+	notifyExhausted,
+	notifyKeySwitch,
+	notifyProviderSwitch,
+} from "./notification.ts";
+import { createPiRuntimeAdapter, type PiRuntimeAdapter } from "./pi-runtime.ts";
+
+export interface FailoverExtensionOptions {
+	loadCatalog?: () => AuthCatalog;
+	writeBackupKey?: (providerId: string, backupKey: string) => AddBackupKeyResult;
+	now?: () => number;
+}
+
+interface AttemptResponse extends FailoverAttempt {
+	status?: number;
+	retryAfterMs?: number;
+}
+
+type FailoverExecutionResult = "applied" | "exhausted" | "none";
+
+const OVERLOADED_TEXT = /\boverload(?:ed)?(?:[_ -]?error)?\b/i;
+const RETRY_MESSAGE = {
+	customType: "pi-failover-retry",
+	content: "Retry the current user request now using the failover credential or provider. Do not mention this internal retry.",
+	display: false,
+} as const;
+
+function classifyAttemptFailure(attempt: AttemptResponse, errorMessage: string): FailureObservation {
+	const retryAfter = attempt.retryAfterMs === undefined ? {} : { retryAfterMs: attempt.retryAfterMs };
+	if (attempt.status !== undefined) {
+		if (attempt.status === 429 && OVERLOADED_TEXT.test(errorMessage)) {
+			return { kind: "overloaded", ...retryAfter };
+		}
+		return { status: attempt.status, ...retryAfter };
+	}
+	if (OVERLOADED_TEXT.test(errorMessage) || /\b529\b/i.test(errorMessage)) return { kind: "overloaded", ...retryAfter };
+	if (/\btimeout|timed out|deadline exceeded/i.test(errorMessage)) return { kind: "network", ...retryAfter };
+	if (/\bnetwork|connection (?:reset|refused|closed)|fetch failed|socket hang up/i.test(errorMessage)) {
+		return { kind: "network", ...retryAfter };
+	}
+	if (/\b40[13]\b|unauthorized|forbidden/i.test(errorMessage)) return { kind: "unauthorized", ...retryAfter };
+	if (/\b429\b|rate.?limit|too many requests/i.test(errorMessage)) return { kind: "rate-limited", ...retryAfter };
+	if (/\b50[0234]\b|temporar(?:y|ily) unavailable/i.test(errorMessage)) return { kind: "provider-error", ...retryAfter };
+	return retryAfter;
+}
+
+function retryAfterMilliseconds(headers: Record<string, string>, now: number): number | undefined {
+	const value = Object.entries(headers).find(([name]) => name.toLowerCase() === "retry-after")?.[1]?.trim();
+	if (!value) return undefined;
+	const seconds = Number(value);
+	if (Number.isFinite(seconds) && seconds >= 0) return seconds * 1_000;
+	const date = Date.parse(value);
+	return Number.isFinite(date) ? Math.max(0, date - now) : undefined;
+}
+
+export function createFailoverExtension(options: FailoverExtensionOptions = {}) {
+	const readCatalog = options.loadCatalog ?? loadAuthCatalog;
+	const writeBackupKey = options.writeBackupKey ?? addBackupKey;
+	const now = options.now ?? Date.now;
+
+	return function registerFailover(pi: ExtensionAPI): void {
+		let catalog: AuthCatalog = { enabled: false, providers: [], diagnostics: [] };
+		let engine: FailoverEngine | undefined;
+		let runtime: PiRuntimeAdapter | undefined;
+		let attempt: AttemptResponse | undefined;
+		let activeContext: ExtensionContext | undefined;
+		const possiblyOwnedProviderIds = new Set<string>();
+		let compatibilityErrorReported = false;
+		let exhaustionPending = false;
+		let deferredAttempt: FailoverAttempt | undefined;
+
+		function clearPendingExhaustion(): void {
+			exhaustionPending = false;
+			deferredAttempt = undefined;
+		}
+
+		function nextProvider(current: { providerId: string; model: string }, unavailableProviderIds: readonly string[]) {
+			if (!activeContext) return undefined;
+			const order = catalog.providers.map((provider) => provider.provider);
+			const currentIndex = order.indexOf(current.providerId);
+			const candidates = currentIndex < 0
+				? order
+				: [...order.slice(currentIndex + 1), ...order.slice(0, currentIndex)];
+			const unavailable = new Set(unavailableProviderIds);
+			const models = activeContext.modelRegistry.getAvailable();
+			for (const providerId of candidates) {
+				if (providerId === current.providerId || unavailable.has(providerId)) continue;
+				if (!activeContext.modelRegistry.getProviderAuthStatus(providerId).configured) continue;
+				const providerModels = models.filter((model) => model.provider === providerId);
+				const model = providerModels.find((candidate) => candidate.id === current.model) ?? providerModels[0];
+				if (model) return { providerId, model: model.id };
+			}
+			return undefined;
+		}
+
+		function rebuild(ctx: ExtensionContext): void {
+			clearPendingExhaustion();
+			activeContext = ctx;
+			catalog = readCatalog();
+			runtime = createPiRuntimeAdapter(ctx.modelRegistry);
+			if (!runtime.supported) {
+				engine = undefined;
+				if (!compatibilityErrorReported) {
+					notify(ctx, "pi-failover: disabled because this Pi runtime lacks credential override support", "error");
+					compatibilityErrorReported = true;
+				}
+				attempt = undefined;
+				return;
+			}
+			engine = catalog.enabled
+				? new FailoverEngine({
+					providers: catalog.providers.map((provider) => ({
+						id: provider.provider,
+						backupKeyCount: provider.backupKeys?.length ?? 0,
+					})),
+					now,
+					nextProvider: ({ current, unavailableProviderIds }) => nextProvider(current, unavailableProviderIds),
+				})
+				: undefined;
+			attempt = undefined;
+		}
+
+		async function restoreOwnedOverrides(): Promise<void> {
+			if (!runtime) return;
+			for (const providerId of [...possiblyOwnedProviderIds]) {
+				const result = await runtime.restoreOriginalKey(providerId);
+				if (result.ok) possiblyOwnedProviderIds.delete(providerId);
+			}
+		}
+
+		async function execute(initialDecision: FailoverDecision, ctx: ExtensionContext): Promise<FailoverExecutionResult> {
+			let decision = initialDecision;
+			while (decision.kind === "switch-key" || decision.kind === "switch-model") {
+				const activeDecision = decision;
+				if (activeDecision.kind === "switch-model") {
+					const applied = await applyNextModel(
+						{ modelRegistry: ctx.modelRegistry, setModel: (model) => pi.setModel(model) },
+						{
+							authOrder: [activeDecision.providerId],
+							current: { providerId: "", model: activeDecision.model },
+							unavailableProviderIds: [],
+							cooldownProviderIds: [],
+						},
+					);
+					if (!applied) {
+						decision = engine?.observeFailure({ kind: "provider-error" }) ?? { kind: "exhausted" };
+						continue;
+					}
+					notifyProviderSwitch(ctx, applied.providerId, applied.model);
+				}
+				if (activeDecision.keySlot === "primary") {
+					if (!runtime?.hasOwnedOverride(activeDecision.providerId)) {
+						if (possiblyOwnedProviderIds.has(activeDecision.providerId)) {
+							const result = await runtime?.restoreOriginalKey(activeDecision.providerId);
+							if (result?.ok) possiblyOwnedProviderIds.delete(activeDecision.providerId);
+						}
+						clearPendingExhaustion();
+						return "applied";
+					}
+					const result = await runtime?.restoreOriginalKey(activeDecision.providerId);
+					if (result?.ok) {
+						possiblyOwnedProviderIds.delete(activeDecision.providerId);
+						clearPendingExhaustion();
+						return "applied";
+					}
+					decision = engine?.observeFailure({ kind: "provider-error" }) ?? { kind: "exhausted" };
+					continue;
+				}
+				const provider = catalog.providers.find((candidate) => candidate.provider === activeDecision.providerId);
+				const backupIndex = backupIndexForSlot(activeDecision.keySlot);
+				const backupKey = backupIndex === undefined ? undefined : provider?.backupKeys?.[backupIndex];
+				if (backupKey !== undefined && runtime?.ownsBackupKey(activeDecision.providerId, backupKey)) {
+					// The backup override is already in effect from a prior turn; re-applying
+					// would only duplicate the runtime write and re-announce the switch.
+					clearPendingExhaustion();
+					return "applied";
+				}
+				possiblyOwnedProviderIds.add(activeDecision.providerId);
+				const result = backupKey === undefined
+					? { ok: false as const }
+					: await runtime?.setBackupKey(activeDecision.providerId, backupKey);
+				if (result?.ok) {
+					notifyKeySwitch(ctx, activeDecision.providerId, activeDecision.keySlot);
+					clearPendingExhaustion();
+					return "applied";
+				}
+				if (!runtime?.hasOwnedOverride(activeDecision.providerId)) {
+					possiblyOwnedProviderIds.delete(activeDecision.providerId);
+				}
+				decision = engine?.observeFailure({ kind: "unauthorized" }) ?? { kind: "exhausted" };
+			}
+			if (decision.kind === "exhausted") {
+				// Pi decides whether to auto-retry only after message_end handlers finish.
+				// Keep the active credential in place until agent_settled so a built-in
+				// retry cannot fall back to a known-bad primary credential.
+				exhaustionPending = true;
+				return "exhausted";
+			}
+			return "none";
+		}
+
+		pi.on("session_start", (_event, ctx) => {
+			rebuild(ctx);
+		});
+
+		pi.on("turn_start", async (_event, ctx) => {
+			activeContext = ctx;
+			if (!engine || !ctx.model) return;
+			if (exhaustionPending) return;
+			await execute(engine.startTurn({ providerId: ctx.model.provider, model: ctx.model.id }), ctx);
+		});
+
+		pi.on("before_provider_request", (_event, ctx) => {
+			attempt = undefined;
+			const decision = engine?.currentDecision();
+			const activeAttempt = decision && (decision.kind === "switch-key" || decision.kind === "switch-model")
+				? decision
+				: exhaustionPending
+					? deferredAttempt
+					: undefined;
+			if (ctx.model && activeAttempt?.providerId === ctx.model.provider) {
+				attempt = {
+					providerId: activeAttempt.providerId,
+					model: activeAttempt.model,
+					keySlot: activeAttempt.keySlot,
+				};
+			}
+		});
+
+		pi.on("after_provider_response", (event) => {
+			if (!attempt) return;
+			attempt.status = event.status;
+			attempt.retryAfterMs = retryAfterMilliseconds(event.headers, now());
+			if (event.status >= 200 && event.status < 300) engine?.observeSuccess();
+		});
+
+		pi.on("message_end", async (event, ctx) => {
+			if (event.message.role !== "assistant") return;
+			if (event.message.stopReason !== "error") {
+				attempt = undefined;
+				clearPendingExhaustion();
+				return;
+			}
+			if (!attempt || !engine) return;
+			const failedAttempt = attempt;
+			attempt = undefined;
+			const errorMessage = event.message.errorMessage ?? "";
+			const result = await execute(
+				engine.observeFailure(classifyAttemptFailure(failedAttempt, errorMessage)),
+				ctx,
+			);
+			if (result === "exhausted") {
+				deferredAttempt = {
+					providerId: failedAttempt.providerId,
+					model: failedAttempt.model,
+					keySlot: failedAttempt.keySlot,
+				};
+				engine.resumeAttempt(deferredAttempt);
+			}
+			if (result !== "applied") return;
+
+			pi.sendMessage(RETRY_MESSAGE, { triggerTurn: true, deliverAs: "followUp" });
+			return {
+				message: {
+					...event.message,
+					content: [],
+					stopReason: "stop",
+					errorMessage: undefined,
+				},
+			};
+		});
+
+		pi.on("agent_settled", async () => {
+			if (!exhaustionPending) return;
+			clearPendingExhaustion();
+			await restoreOwnedOverrides();
+			if (activeContext) notifyExhausted(activeContext);
+		});
+
+		pi.on("session_shutdown", async () => {
+			await restoreOwnedOverrides();
+			engine = undefined;
+			attempt = undefined;
+			clearPendingExhaustion();
+			activeContext = undefined;
+		});
+
+		pi.registerCommand("failover-status", {
+			description: "show redacted failover status",
+			handler: async (_args, ctx) => {
+				notify(ctx, formatStatus(catalog, engine?.snapshot()), "info");
+			},
+		});
+
+		pi.registerCommand("failover-reload", {
+			description: "reload failover configuration from auth.json",
+			handler: async (_args, ctx) => {
+				await restoreOwnedOverrides();
+				rebuild(ctx);
+				notify(ctx, catalog.enabled ? "pi-failover: reloaded auth.json" : "pi-failover: auth.json reload disabled failover", catalog.enabled ? "info" : "warning");
+			},
+		});
+
+		pi.registerCommand("failover-login", {
+			description: "add a backup API key to a provider's api_key credential in auth.json",
+			getArgumentCompletions: (prefix) => {
+				const items = readCatalog()
+					.providers.filter((provider) => provider.type === "api_key")
+					.map((provider) => ({ value: provider.provider, label: provider.provider }));
+				const matches = items.filter((item) => item.value.startsWith(prefix));
+				return matches.length > 0 ? matches : null;
+			},
+			handler: async (args, ctx) => {
+				await runFailoverLogin(args, {
+					hasUI: ctx.hasUI,
+					select: (title, options) => ctx.ui.select(title, options),
+					input: (title, placeholder) => ctx.ui.input(title, placeholder),
+					confirm: (title, message) => ctx.ui.confirm(title, message),
+					readCatalog,
+					writeBackupKey,
+					rebuild: async () => {
+						await restoreOwnedOverrides();
+						rebuild(ctx);
+					},
+					notify: (message, level) => notify(ctx, message, level),
+				});
+			},
+		});
+	};
+}
+
+export default createFailoverExtension();

package/src/model-planner.ts

@@ -0,0 +1,51 @@
+import type { ProviderPlan } from "./failover-engine.ts";
+
+export interface PlannerModel {
+	provider: string;
+	id: string;
+}
+
+export interface ModelPlannerContext<TModel extends PlannerModel = PlannerModel> {
+	modelRegistry: {
+		getAvailable(): readonly TModel[];
+		getProviderAuthStatus(providerId: string): { configured: boolean };
+	};
+	setModel(model: TModel): boolean | Promise<boolean>;
+}
+
+export interface ApplyNextModelRequest {
+	authOrder: readonly (string | { provider: string })[];
+	current: ProviderPlan;
+	unavailableProviderIds: readonly string[];
+	cooldownProviderIds: readonly string[];
+}
+
+export async function applyNextModel<TModel extends PlannerModel>(
+	ctx: ModelPlannerContext<TModel>,
+	request: ApplyNextModelRequest,
+): Promise<ProviderPlan | undefined> {
+	const providerIds = request.authOrder.map((entry) => typeof entry === "string" ? entry : entry.provider);
+	const currentIndex = providerIds.indexOf(request.current.providerId);
+	const ordered = currentIndex < 0
+		? providerIds
+		: [...providerIds.slice(currentIndex + 1), ...providerIds.slice(0, currentIndex)];
+	const available = ctx.modelRegistry.getAvailable();
+	const unavailable = new Set(request.unavailableProviderIds);
+	const cooling = new Set(request.cooldownProviderIds);
+
+	for (const providerId of ordered) {
+		if (unavailable.has(providerId) || cooling.has(providerId)) continue;
+		if (!ctx.modelRegistry.getProviderAuthStatus(providerId).configured) continue;
+		const providerModels = available.filter((candidate) => candidate.provider === providerId);
+		const model = providerModels.find((candidate) => candidate.id === request.current.model) ?? providerModels[0];
+		if (!model) continue;
+		try {
+			if (await ctx.setModel(model)) return { providerId, model: model.id };
+		} catch {
+			// A model may become unavailable between the registry snapshot and
+			// application. Continue the deterministic provider walk.
+		}
+	}
+
+	return undefined;
+}

package/src/notification.ts

@@ -0,0 +1,43 @@
+import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
+import type { AuthCatalog } from "./auth-catalog.ts";
+import type { FailoverSnapshot, KeySlot } from "./failover-engine.ts";
+
+export type NotificationLevel = "info" | "warning" | "error";
+export type NotificationContext = Pick<ExtensionContext, "hasUI" | "ui">;
+
+export function notify(ctx: NotificationContext, message: string, level: NotificationLevel): void {
+	if (!ctx.hasUI) return;
+	try {
+		ctx.ui.notify(message, level);
+	} catch {
+		// A disappearing UI must not interrupt provider failover.
+	}
+}
+
+export function notifyKeySwitch(ctx: NotificationContext, providerId: string, slot: KeySlot): void {
+	notify(ctx, `pi-failover: ${providerId} switched to ${slot} credential`, "warning");
+}
+
+export function notifyProviderSwitch(ctx: NotificationContext, providerId: string, model: string): void {
+	notify(ctx, `pi-failover: switched to ${providerId}/${model}`, "warning");
+}
+
+export function notifyExhausted(ctx: NotificationContext): void {
+	notify(ctx, "pi-failover: all configured providers exhausted; preserving the original provider error", "error");
+}
+
+export function formatStatus(catalog: AuthCatalog, snapshot: FailoverSnapshot | undefined): string {
+	if (!catalog.enabled || !snapshot) return "pi-failover: disabled";
+	const lines = ["pi-failover: active"];
+	if (snapshot.active) {
+		lines.push(`current: ${snapshot.active.providerId}/${snapshot.active.model}; key=${snapshot.active.keySlot}`);
+	}
+	for (const provider of snapshot.providers) {
+		const providerCooldown = provider.cooldownUntil === undefined ? "" : ` until=${provider.cooldownUntil}`;
+		const keys = provider.keys.map((key) => `${key.slot}=${key.status}${key.cooldownUntil === undefined ? "" : ` until=${key.cooldownUntil}`}`).join(", ");
+		lines.push(`${provider.providerId}: ${provider.status}${providerCooldown}; ${keys}`);
+	}
+	lines.push(`visited providers: ${snapshot.visitedProviders.join(", ") || "none"}`);
+	lines.push(`visited keys: ${snapshot.visitedKeys.map((key) => `${key.providerId}:${key.keySlot}`).join(", ") || "none"}`);
+	return lines.join("\n");
+}

package/src/pi-runtime.ts

@@ -0,0 +1,192 @@
+export interface PiAuthStatus {
+	configured: boolean;
+	source?: string;
+}
+
+export interface PiRuntimeRegistry {
+	getApiKeyForProvider(providerId: string): Promise<string | undefined>;
+	getProviderAuthStatus(providerId: string): PiAuthStatus;
+}
+
+export type RuntimeOverrideResult =
+	| { ok: true; action: "set" | "removed" | "restored" | "unchanged" }
+	| { ok: false; reason: "unsupported" | "operation_failed" };
+
+export interface PiRuntimeAdapter {
+	readonly supported: boolean;
+	hasOwnedOverride(providerId: string): boolean;
+	ownsBackupKey(providerId: string, backupKey: string): boolean;
+	setBackupKey(providerId: string, backupKey: string): Promise<RuntimeOverrideResult>;
+	restoreOriginalKey(providerId: string): Promise<RuntimeOverrideResult>;
+}
+
+type RuntimeApi = {
+	setRuntimeApiKey(providerId: string, apiKey: string): void | Promise<void>;
+	removeRuntimeApiKey(providerId: string): void | Promise<void>;
+};
+
+type OwnedOverride =
+	| { ownedKey: string; originalSourceWasRuntime: true; originalRuntimeKey: string }
+	| { ownedKey: string; originalSourceWasRuntime: false };
+
+const OWNERSHIP_SYMBOL = Symbol.for("pi-failover.runtime-override-ownership");
+
+function ownershipStore(): WeakMap<object, Map<string, OwnedOverride>> {
+	const root = globalThis as typeof globalThis & Record<PropertyKey, unknown>;
+	const existing = root[OWNERSHIP_SYMBOL];
+	if (existing instanceof WeakMap) return existing as WeakMap<object, Map<string, OwnedOverride>>;
+
+	const created = new WeakMap<object, Map<string, OwnedOverride>>();
+	root[OWNERSHIP_SYMBOL] = created;
+	return created;
+}
+
+function runtimeApi(registry: PiRuntimeRegistry): RuntimeApi | undefined {
+	// ModelRegistry.runtime is private in Pi's declarations. Keep the compatibility
+	// escape hatch contained here; all reads use the public registry facade.
+	const runtime = (registry as PiRuntimeRegistry & { runtime?: Partial<RuntimeApi> }).runtime;
+	if (
+		!runtime ||
+		typeof runtime.setRuntimeApiKey !== "function" ||
+		typeof runtime.removeRuntimeApiKey !== "function"
+	) {
+		return undefined;
+	}
+	return runtime as RuntimeApi;
+}
+
+export function createPiRuntimeAdapter(registry: PiRuntimeRegistry): PiRuntimeAdapter {
+	const runtime = runtimeApi(registry);
+
+	async function readRuntimeState(providerId: string): Promise<{ key: string | undefined; status: PiAuthStatus }> {
+		const key = await registry.getApiKeyForProvider(providerId);
+		return { key, status: registry.getProviderAuthStatus(providerId) };
+	}
+
+	function originalStateConfirmed(state: { key: string | undefined; status: PiAuthStatus }, owned: OwnedOverride): boolean {
+		return owned.originalSourceWasRuntime
+			? state.status.source === "runtime" && state.key === owned.originalRuntimeKey
+			: state.status.source !== "runtime" && state.key !== owned.ownedKey;
+	}
+
+	async function reconcileRejectedSet(
+		api: RuntimeApi,
+		providerId: string,
+		owned: OwnedOverride,
+		previousOwned: OwnedOverride | undefined,
+		byProvider: Map<string, OwnedOverride>,
+	): Promise<void> {
+		let state: { key: string | undefined; status: PiAuthStatus };
+		try {
+			state = await readRuntimeState(providerId);
+		} catch {
+			return;
+		}
+
+		if (previousOwned && state.status.source === "runtime" && state.key === previousOwned.ownedKey) {
+			byProvider.set(providerId, previousOwned);
+			return;
+		}
+		if (originalStateConfirmed(state, owned)) {
+			byProvider.delete(providerId);
+			return;
+		}
+		if (state.status.source !== "runtime" || state.key !== owned.ownedKey) return;
+
+		try {
+			if (owned.originalSourceWasRuntime) {
+				await api.setRuntimeApiKey(providerId, owned.originalRuntimeKey);
+			} else {
+				await api.removeRuntimeApiKey(providerId);
+			}
+		} catch {
+			// Pi can commit compensation before synchronization rejects. Confirm
+			// effective state below before deciding whether ownership can be cleared.
+		}
+
+		try {
+			state = await readRuntimeState(providerId);
+			if (originalStateConfirmed(state, owned)) byProvider.delete(providerId);
+		} catch {
+			// Indeterminate compensation keeps ownership for a later cleanup.
+		}
+	}
+
+	return {
+		supported: runtime !== undefined,
+		hasOwnedOverride(providerId) {
+			return runtime !== undefined && (ownershipStore().get(runtime)?.has(providerId) ?? false);
+		},
+		ownsBackupKey(providerId, backupKey) {
+			if (!runtime) return false;
+			return ownershipStore().get(runtime)?.get(providerId)?.ownedKey === backupKey;
+		},
+
+		async setBackupKey(providerId, backupKey) {
+			if (!runtime) return { ok: false, reason: "unsupported" };
+
+			try {
+				const status = registry.getProviderAuthStatus(providerId);
+				const currentKey = await registry.getApiKeyForProvider(providerId);
+				const store = ownershipStore();
+				const byProvider = store.get(runtime) ?? new Map<string, OwnedOverride>();
+				store.set(runtime, byProvider);
+				const existing = byProvider.get(providerId);
+				let owned: OwnedOverride;
+				if (existing && currentKey === existing.ownedKey) {
+					owned = existing.originalSourceWasRuntime
+						? { ownedKey: backupKey, originalSourceWasRuntime: true, originalRuntimeKey: existing.originalRuntimeKey }
+						: { ownedKey: backupKey, originalSourceWasRuntime: false };
+				} else if (status.source === "runtime") {
+					if (currentKey === undefined) return { ok: false, reason: "operation_failed" };
+					owned = { ownedKey: backupKey, originalSourceWasRuntime: true, originalRuntimeKey: currentKey };
+				} else {
+					owned = { ownedKey: backupKey, originalSourceWasRuntime: false };
+				}
+				byProvider.set(providerId, owned);
+				try {
+					await runtime.setRuntimeApiKey(providerId, backupKey);
+				} catch {
+					await reconcileRejectedSet(runtime, providerId, owned, existing, byProvider);
+					return { ok: false, reason: "operation_failed" };
+				}
+				return { ok: true, action: "set" };
+			} catch {
+				return { ok: false, reason: "operation_failed" };
+			}
+		},
+
+		async restoreOriginalKey(providerId) {
+			if (!runtime) return { ok: false, reason: "unsupported" };
+			const byProvider = ownershipStore().get(runtime);
+			const owned = byProvider?.get(providerId);
+			if (!owned) return { ok: true, action: "unchanged" };
+
+			try {
+				const currentKey = await registry.getApiKeyForProvider(providerId);
+				if (
+					currentKey === undefined &&
+					registry.getProviderAuthStatus(providerId).source === "runtime"
+				) {
+					return { ok: false, reason: "operation_failed" };
+				}
+				if (currentKey !== owned.ownedKey) {
+					byProvider?.delete(providerId);
+					return { ok: true, action: "unchanged" };
+				}
+
+				if (owned.originalSourceWasRuntime) {
+					await runtime.setRuntimeApiKey(providerId, owned.originalRuntimeKey);
+					byProvider?.delete(providerId);
+					return { ok: true, action: "restored" };
+				}
+
+				await runtime.removeRuntimeApiKey(providerId);
+				byProvider?.delete(providerId);
+				return { ok: true, action: "removed" };
+			} catch {
+				return { ok: false, reason: "operation_failed" };
+			}
+		},
+	};
+}

package/tsconfig.json

@@ -0,0 +1,19 @@
+{
+	"compilerOptions": {
+		"target": "ES2022",
+		"module": "NodeNext",
+		"moduleResolution": "NodeNext",
+		"strict": true,
+		"noEmit": true,
+		"skipLibCheck": true,
+		"forceConsistentCasingInFileNames": true,
+		"allowImportingTsExtensions": true
+	},
+	"include": [
+		"src/**/*.ts",
+		"tests/**/*.ts"
+	],
+	"exclude": [
+		"node_modules"
+	]
+}