@thelioo/opencode-balancer 0.1.8 → 0.2.1

package/dist/server/native.ts

@@ -0,0 +1,49 @@
+import type { Database } from "bun:sqlite";
+import { suppressNativeAuthCapture as suppressDbNativeAuthCapture } from "../core/native-auth-suppression";
+import { isNativeConnectInProgress } from "../core/native-connect";
+import { now } from "../core/time";
+import type { AuthInfo } from "../core/types";
+
+const suppressAuthCaptureUntil = new Map<string, number>();
+
+export function suppressNativeAuthCapture(
+	providerID: string,
+	durationMs = 2_000,
+) {
+	suppressAuthCaptureUntil.set(providerID, now() + durationMs);
+}
+
+export function isNativeAuthCaptureSuppressed(providerID: string) {
+	const suppressedUntil = suppressAuthCaptureUntil.get(providerID) ?? 0;
+	if (suppressedUntil <= now()) {
+		suppressAuthCaptureUntil.delete(providerID);
+		return false;
+	}
+	return true;
+}
+
+export async function setNativeAuth(
+	client: any,
+	providerID: string,
+	auth: AuthInfo,
+	db?: Database,
+) {
+	if (db && isNativeConnectInProgress(db)) return;
+	suppressNativeAuthCapture(providerID);
+	if (db) suppressDbNativeAuthCapture(db, providerID);
+	try {
+		await client.auth.set({ body: auth, path: { id: providerID } });
+	} catch {}
+}
+
+export async function showToast(
+	client: any,
+	message: string,
+	variant: "info" | "success" | "warning" | "error" = "info",
+) {
+	try {
+		await client.tui.showToast({
+			body: { duration: 15_000, message, variant },
+		});
+	} catch {}
+}

package/dist/server/request-balancer.ts

@@ -0,0 +1,67 @@
+import type { Database } from "bun:sqlite";
+import { getAccount, listAccounts, normalizeAlias } from "../core/accounts";
+import { now } from "../core/time";
+import type { Account } from "../core/types";
+
+export const INTERNAL_REQUEST_HEADER = "x-opencode-balancer-request";
+export const BALANCER_METADATA_KEY = "opencodeBalancerCommand";
+
+export const RETRYABLE_STATUS = new Set([429, 500, 502, 503, 504, 529]);
+
+type PendingRequest = {
+	providerID: string;
+	account?: Account;
+};
+
+const pendingRequests = new Map<string, PendingRequest>();
+
+export function setPendingRequest(requestID: string, request: PendingRequest) {
+	pendingRequests.set(requestID, request);
+}
+
+export function takePendingRequest(requestID: string) {
+	const request = pendingRequests.get(requestID);
+	pendingRequests.delete(requestID);
+	return request;
+}
+
+export function __testGetPendingRequest(requestID: string) {
+	return pendingRequests.get(requestID);
+}
+
+export function __testClearPendingRequests() {
+	pendingRequests.clear();
+}
+
+export function markRateLimited(
+	db: Database,
+	providerID: string,
+	alias: string,
+	retryAfterMs = 60_000,
+) {
+	const account = getAccount(db, providerID, alias);
+	if (!account) return;
+
+	const timestamp = now();
+	db.query<unknown, [number, number, string, string]>(
+		`UPDATE accounts
+         SET failures = failures + 1,
+             rate_limited_until = ?,
+             updated_at = ?
+         WHERE provider_id = ? AND alias = ?`,
+	).run(timestamp + retryAfterMs, timestamp, providerID, account.alias);
+}
+
+export function chooseFailoverAccount(
+	db: Database,
+	providerID: string,
+	currentAlias: string,
+) {
+	const timestamp = now();
+	const normalizedCurrentAlias = normalizeAlias(currentAlias);
+	return listAccounts(db, providerID).find((account) => {
+		if (account.alias === normalizedCurrentAlias) return false;
+		if (account.disabled) return false;
+		return !account.rateLimitedUntil || account.rateLimitedUntil <= timestamp;
+	});
+}

package/dist/tui/actions.ts

@@ -1,144 +1,208 @@
-import { getAccount, getActiveAccount, getSelectedAccount, removeAccount, renameAccount, setActiveAccount } from "../core/accounts";
+import {
+	getAccount,
+	getActiveAccount,
+	getSelectedAccount,
+	removeAccount,
+	renameAccount,
+	setActiveAccount,
+} from "../core/accounts";
 import { appendEvent } from "../core/events";
-import { completePendingConnection, removePendingConnection } from "../core/pending";
+import { suppressNativeAuthCapture } from "../core/native-auth-suppression";
+import {
+	completePendingConnection,
+	removePendingConnection,
+} from "../core/pending";
 import { refreshAccountUsage } from "../core/usage";
 import { saveUsageSnapshot } from "../core/usage/store";
-import { suppressNativeAuthCapture } from "../core/native-auth-suppression";
 import type { BalancerTuiState } from "./state";
 
 type AuthSetApi = {
-    client: {
-        auth: {
-            set: (input: any) => unknown;
-        };
-    };
-    keymap?: {
-        dispatchCommand?: (command: string) => unknown;
-    };
-    ui?: {
-        toast: (input: { variant: "success" | "warning" | "error" | "info"; message: string }) => unknown;
-    };
+	client: {
+		auth: {
+			set: (input: any) => unknown;
+		};
+	};
+	keymap?: {
+		dispatchCommand?: (command: string) => unknown;
+	};
+	ui?: {
+		toast: (input: {
+			variant: "success" | "warning" | "error" | "info";
+			message: string;
+		}) => unknown;
+	};
 };
 
 type ActivateAccountOptions = {
-    sessionProviderID?: string;
-    applyNativeProviderModel?: (providerID: string) => Promise<boolean>;
+	sessionProviderID?: string;
+	applyNativeProviderModel?: (providerID: string) => Promise<boolean>;
 };
 
 type ToastApi = {
-    ui: {
-        toast: (input: { variant: "success" | "warning" | "error" | "info"; message: string }) => unknown;
-    };
+	ui: {
+		toast: (input: {
+			variant: "success" | "warning" | "error" | "info";
+			message: string;
+		}) => unknown;
+	};
 };
 
 type RefreshUsageOptions = {
-    refreshUsage?: typeof refreshAccountUsage;
-    silent?: boolean;
+	refreshUsage?: typeof refreshAccountUsage;
+	silent?: boolean;
 };
 
-export function savePendingAlias(state: BalancerTuiState, pendingID: string, alias: string) {
-    const account = completePendingConnection(state.db, pendingID, alias);
-    state.refresh();
-    return account;
+export function savePendingAlias(
+	state: BalancerTuiState,
+	pendingID: string,
+	alias: string,
+) {
+	const account = completePendingConnection(state.db, pendingID, alias);
+	state.refresh();
+	return account;
 }
 
 export async function activateAccount(
-    api: AuthSetApi,
-    state: BalancerTuiState,
-    providerID: string,
-    alias: string,
-    options: ActivateAccountOptions = {},
+	api: AuthSetApi,
+	state: BalancerTuiState,
+	providerID: string,
+	alias: string,
+	options: ActivateAccountOptions = {},
 ) {
-    const previousProviderID = getSelectedAccount(state.db)?.providerID;
-    const account = setActiveAccount(state.db, providerID, alias) ?? getActiveAccount(state.db, providerID);
-
-    if (account) {
-        try {
-            suppressNativeAuthCapture(state.db, providerID);
-            await api.client.auth.set({ path: { id: providerID }, body: account.auth });
-        } catch {}
-    }
-
-    state.refresh();
-    const providerChanged = (options.sessionProviderID ?? previousProviderID) !== providerID;
-    if (providerChanged) {
-        await options.applyNativeProviderModel?.(providerID);
-    }
-    api.ui?.toast({ variant: "success", message: `Activated ${providerID}/${alias}.` });
+	const previousProviderID = getSelectedAccount(state.db)?.providerID;
+	const account =
+		setActiveAccount(state.db, providerID, alias) ??
+		getActiveAccount(state.db, providerID);
+
+	if (account) {
+		try {
+			suppressNativeAuthCapture(state.db, providerID);
+			await api.client.auth.set({
+				body: account.auth,
+				path: { id: providerID },
+			});
+		} catch {}
+	}
+
+	state.refresh();
+	const providerChanged =
+		(options.sessionProviderID ?? previousProviderID) !== providerID;
+	if (providerChanged) {
+		await options.applyNativeProviderModel?.(providerID);
+	}
+	api.ui?.toast({
+		message: `Activated ${providerID}/${alias}.`,
+		variant: "success",
+	});
 }
 
 export async function refreshUsageForAccount(
-    api: ToastApi,
-    state: BalancerTuiState,
-    providerID: string,
-    alias: string,
-    options: RefreshUsageOptions = {},
+	api: ToastApi,
+	state: BalancerTuiState,
+	providerID: string,
+	alias: string,
+	options: RefreshUsageOptions = {},
 ) {
-    const account = getAccount(state.db, providerID, alias) ?? state.accounts().find((candidate) => {
-        return candidate.providerID === providerID && candidate.alias === alias;
-    });
-
-    if (!account) {
-        api.ui.toast({ variant: "error", message: `Account not found: ${providerID}/${alias}` });
-        return;
-    }
-
-    try {
-        const snapshot = await (options.refreshUsage ?? refreshAccountUsage)(account);
-        saveUsageSnapshot(state.db, snapshot);
-        appendEvent(state.db, {
-            type: snapshot.confidence === "unavailable" ? "usage_unavailable" : "usage_refreshed",
-            providerID,
-            alias,
-            message: snapshot.message,
-        });
-        state.refresh();
-        if (!options.silent) {
-            api.ui.toast({
-                variant: snapshot.confidence === "unavailable" ? "warning" : "success",
-                message: snapshot.message,
-            });
-        }
-    } catch (error) {
-        const message = error instanceof Error ? error.message : "Usage refresh failed.";
-        api.ui.toast({ variant: "error", message });
-    }
+	const account =
+		getAccount(state.db, providerID, alias) ??
+		state.accounts().find((candidate) => {
+			return candidate.providerID === providerID && candidate.alias === alias;
+		});
+
+	if (!account) {
+		api.ui.toast({
+			message: `Account not found: ${providerID}/${alias}`,
+			variant: "error",
+		});
+		return;
+	}
+
+	try {
+		const snapshot = await (options.refreshUsage ?? refreshAccountUsage)(
+			account,
+		);
+		saveUsageSnapshot(state.db, snapshot);
+		appendEvent(state.db, {
+			alias,
+			message: snapshot.message,
+			providerID,
+			type:
+				snapshot.confidence === "unavailable"
+					? "usage_unavailable"
+					: "usage_refreshed",
+		});
+		state.refresh();
+		if (!options.silent) {
+			api.ui.toast({
+				message: snapshot.message,
+				variant: snapshot.confidence === "unavailable" ? "warning" : "success",
+			});
+		}
+	} catch (error) {
+		const message =
+			error instanceof Error ? error.message : "Usage refresh failed.";
+		api.ui.toast({ message, variant: "error" });
+	}
 }
 
-export function removeAccountFromTui(api: ToastApi, state: BalancerTuiState, providerID: string, alias: string) {
-    const removed = removeAccount(state.db, providerID, alias);
-    if (!removed) {
-        api.ui.toast({ variant: "error", message: `Account not found: ${providerID}/${alias}` });
-        return;
-    }
-
-    const message = `Removed account ${providerID}/${alias}.`;
-    appendEvent(state.db, {
-        type: "account_removed",
-        providerID,
-        alias,
-        message,
-    });
-    state.refresh();
-    state.removeAccountView(providerID, alias);
-    api.ui.toast({ variant: "success", message });
+export function removeAccountFromTui(
+	api: ToastApi,
+	state: BalancerTuiState,
+	providerID: string,
+	alias: string,
+) {
+	const removed = removeAccount(state.db, providerID, alias);
+	if (!removed) {
+		api.ui.toast({
+			message: `Account not found: ${providerID}/${alias}`,
+			variant: "error",
+		});
+		return;
+	}
+
+	const message = `Removed account ${providerID}/${alias}.`;
+	appendEvent(state.db, {
+		alias,
+		message,
+		providerID,
+		type: "account_removed",
+	});
+	state.refresh();
+	state.removeAccountView(providerID, alias);
+	api.ui.toast({ message, variant: "success" });
 }
 
-export function renameAccountFromTui(api: ToastApi, state: BalancerTuiState, providerID: string, alias: string, nextAlias: string) {
-    const account = renameAccount(state.db, providerID, alias, nextAlias);
-    state.refresh();
-    api.ui.toast({ variant: "success", message: `Renamed ${providerID}/${alias} to ${account.alias}.` });
-    return account;
+export function renameAccountFromTui(
+	api: ToastApi,
+	state: BalancerTuiState,
+	providerID: string,
+	alias: string,
+	nextAlias: string,
+) {
+	const account = renameAccount(state.db, providerID, alias, nextAlias);
+	state.refresh();
+	api.ui.toast({
+		message: `Renamed ${providerID}/${alias} to ${account.alias}.`,
+		variant: "success",
+	});
+	return account;
 }
 
-export function removePendingFromTui(api: ToastApi, state: BalancerTuiState, pendingID: string) {
-    const removed = removePendingConnection(state.db, pendingID);
-    if (!removed) {
-        api.ui.toast({ variant: "error", message: "Pending connection not found." });
-        return;
-    }
-
-    state.refresh();
-    state.removePendingView(pendingID);
-    api.ui.toast({ variant: "success", message: "Removed pending connection." });
+export function removePendingFromTui(
+	api: ToastApi,
+	state: BalancerTuiState,
+	pendingID: string,
+) {
+	const removed = removePendingConnection(state.db, pendingID);
+	if (!removed) {
+		api.ui.toast({
+			message: "Pending connection not found.",
+			variant: "error",
+		});
+		return;
+	}
+
+	state.refresh();
+	state.removePendingView(pendingID);
+	api.ui.toast({ message: "Removed pending connection.", variant: "success" });
 }

package/dist/tui/balancer-bar-sync.ts

@@ -1,56 +1,66 @@
 import type { TuiPluginApi } from "@opencode-ai/plugin/tui";
 import { getBalancingEnabled, resolveActiveSelection } from "../core/priority";
+import {
+	createNativeModelApplier,
+	type NativeModelApplier,
+} from "./native-model-apply";
 import type { BalancerTuiState } from "./state";
-import { createNativeModelApplier, type NativeModelApplier } from "./native-model-apply";
 
 export type BalancerBarSyncDeps = {
-    balancingEnabled: () => boolean;
-    dialogOpen: () => boolean;
-    activeSelection: () => { providerID: string; modelID: string } | undefined;
-    modelTitle: (providerID: string, modelID: string) => string | undefined;
-    apply: NativeModelApplier;
+	balancingEnabled: () => boolean;
+	dialogOpen: () => boolean;
+	activeSelection: () => { providerID: string; modelID: string } | undefined;
+	modelTitle: (providerID: string, modelID: string) => string | undefined;
+	apply: NativeModelApplier;
 };
 
 export function createBalancerBarSync(deps: BalancerBarSyncDeps) {
-    let lastApplied: string | undefined;
-    let applying = false;
-
-    const maybeSync = async () => {
-        if (!deps.balancingEnabled()) {
-            lastApplied = undefined;
-            return false;
-        }
-        if (applying || deps.dialogOpen()) return false;
-
-        const selection = deps.activeSelection();
-        if (!selection) return false;
-
-        const key = `${selection.providerID}/${selection.modelID}`;
-        if (key === lastApplied) return false;
-
-        const title = deps.modelTitle(selection.providerID, selection.modelID) ?? selection.modelID;
-        applying = true;
-        try {
-            const applied = await deps.apply(selection, title);
-            if (applied) lastApplied = key;
-            return applied;
-        } finally {
-            applying = false;
-        }
-    };
-
-    return { maybeSync };
+	let lastApplied: string | undefined;
+	let applying = false;
+
+	const maybeSync = async () => {
+		if (!deps.balancingEnabled()) {
+			lastApplied = undefined;
+			return false;
+		}
+		if (applying || deps.dialogOpen()) return false;
+
+		const selection = deps.activeSelection();
+		if (!selection) return false;
+
+		const key = `${selection.providerID}/${selection.modelID}`;
+		if (key === lastApplied) return false;
+
+		const title =
+			deps.modelTitle(selection.providerID, selection.modelID) ??
+			selection.modelID;
+		applying = true;
+		try {
+			const applied = await deps.apply(selection, title);
+			if (applied) lastApplied = key;
+			return applied;
+		} finally {
+			applying = false;
+		}
+	};
+
+	return { maybeSync };
 }
 
-export function createTuiBalancerBarSync(api: TuiPluginApi, state: BalancerTuiState) {
-    return createBalancerBarSync({
-        balancingEnabled: () => getBalancingEnabled(state.db),
-        dialogOpen: () => api.ui.dialog.open,
-        activeSelection: () => resolveActiveSelection(state.db),
-        modelTitle: (providerID, modelID) => {
-            const provider = api.state.provider.find((item) => item.id === providerID);
-            return provider?.models?.[modelID]?.name;
-        },
-        apply: createNativeModelApplier(api),
-    });
+export function createTuiBalancerBarSync(
+	api: TuiPluginApi,
+	state: BalancerTuiState,
+) {
+	return createBalancerBarSync({
+		activeSelection: () => resolveActiveSelection(state.db),
+		apply: createNativeModelApplier(api),
+		balancingEnabled: () => getBalancingEnabled(state.db),
+		dialogOpen: () => api.ui.dialog.open,
+		modelTitle: (providerID, modelID) => {
+			const provider = api.state.provider.find(
+				(item) => item.id === providerID,
+			);
+			return provider?.models?.[modelID]?.name;
+		},
+	});
 }