@thelioo/opencode-balancer 0.1.8 → 0.2.1

package/dist/server/auth-watcher.ts

@@ -0,0 +1,318 @@
+import type { Database } from "bun:sqlite";
+import { existsSync, readFileSync } from "node:fs";
+import {
+	getActiveAccount,
+	listAccounts,
+	updateAccountAuth,
+} from "../core/accounts";
+import { isNativeAuthCaptureSuppressed as isDbNativeAuthCaptureSuppressed } from "../core/native-auth-suppression";
+import { nativeAuthPath } from "../core/path";
+import { createPendingConnection } from "../core/pending";
+import type { AuthInfo } from "../core/types";
+import {
+	showToast as defaultShowToast,
+	isNativeAuthCaptureSuppressed,
+} from "./native";
+
+let started = false;
+
+type NativeAuthReadResult =
+	| { ok: true; auth: Record<string, AuthInfo> }
+	| { ok: false };
+
+type ToastVariant = "info" | "success" | "warning" | "error";
+
+type AuthWatcherOptions = {
+	db: Database;
+	client: any;
+	initialSnapshot?: Record<string, AuthInfo>;
+	readAuth: () => NativeAuthReadResult;
+	showToast?: (
+		client: any,
+		message: string,
+		variant: ToastVariant,
+	) => Promise<void>;
+	createPending?: (
+		db: Database,
+		providerID: string,
+		auth: AuthInfo,
+		source: "auth-file",
+	) => unknown;
+	isSuppressed?: (providerID: string) => boolean;
+};
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return Boolean(value) && typeof value === "object" && !Array.isArray(value);
+}
+
+function stringMetadata(value: unknown) {
+	if (value === undefined) return { metadata: undefined, ok: true as const };
+	if (!isRecord(value)) return { ok: false as const };
+
+	const metadata: Record<string, string> = {};
+	for (const [key, entry] of Object.entries(value)) {
+		if (typeof entry !== "string") return { ok: false as const };
+		metadata[key] = entry;
+	}
+	return { metadata, ok: true as const };
+}
+
+function attachMetadata<T extends AuthInfo>(
+	auth: T,
+	value: unknown,
+): T | undefined {
+	const result = stringMetadata(value);
+	if (!result.ok) return undefined;
+	if (result.metadata) auth.metadata = result.metadata;
+	return auth;
+}
+
+function parseAuthInfo(value: unknown): AuthInfo | undefined {
+	if (!isRecord(value) || typeof value.type !== "string") return undefined;
+
+	if (value.type === "api" && typeof value.key === "string") {
+		return attachMetadata({ key: value.key, type: "api" }, value.metadata);
+	}
+
+	if (
+		value.type === "oauth" &&
+		typeof value.refresh === "string" &&
+		typeof value.access === "string" &&
+		typeof value.expires === "number" &&
+		Number.isFinite(value.expires)
+	) {
+		const auth: AuthInfo = {
+			access: value.access,
+			expires: value.expires,
+			refresh: value.refresh,
+			type: "oauth",
+		};
+		if (typeof value.accountId === "string") auth.accountId = value.accountId;
+		if (typeof value.enterpriseUrl === "string")
+			auth.enterpriseUrl = value.enterpriseUrl;
+		return attachMetadata(auth, value.metadata);
+	}
+
+	if (
+		value.type === "wellknown" &&
+		typeof value.key === "string" &&
+		typeof value.token === "string"
+	) {
+		return attachMetadata(
+			{ key: value.key, token: value.token, type: "wellknown" },
+			value.metadata,
+		);
+	}
+
+	return undefined;
+}
+
+function parseNativeAuthContent(content: string): NativeAuthReadResult {
+	let raw: unknown;
+	try {
+		raw = JSON.parse(content);
+	} catch {
+		return { ok: false };
+	}
+
+	if (!isRecord(raw)) return { ok: false };
+
+	const authByProvider: Record<string, AuthInfo> = {};
+	for (const [providerID, auth] of Object.entries(raw)) {
+		const parsed = parseAuthInfo(auth);
+		if (!parsed) return { ok: false };
+		authByProvider[providerID] = parsed;
+	}
+	return { auth: authByProvider, ok: true };
+}
+
+export function readNativeAuth(): NativeAuthReadResult {
+	const authContent = Bun.env.OPENCODE_AUTH_CONTENT;
+	if (authContent !== undefined) return parseNativeAuthContent(authContent);
+
+	const path = nativeAuthPath();
+	if (!existsSync(path)) return { auth: {}, ok: true };
+
+	try {
+		return parseNativeAuthContent(readFileSync(path, "utf8"));
+	} catch {
+		return { ok: false };
+	}
+}
+
+function authChanged(previous: AuthInfo | undefined, current: AuthInfo) {
+	return stableStringify(previous ?? null) !== stableStringify(current);
+}
+
+function stableStringify(value: unknown): string {
+	return JSON.stringify(sortKeys(value));
+}
+
+function sortKeys(value: unknown): unknown {
+	if (!isRecord(value)) return value;
+	return Object.fromEntries(
+		Object.entries(value)
+			.sort(([left], [right]) => left.localeCompare(right))
+			.map(([key, entry]) => [key, sortKeys(entry)]),
+	);
+}
+
+function decodeJwtPayload(token: string): Record<string, unknown> | undefined {
+	const payload = token.split(".")[1];
+	if (!payload) return undefined;
+	try {
+		const decoded = JSON.parse(
+			Buffer.from(payload, "base64url").toString("utf8"),
+		) as unknown;
+		return isRecord(decoded) ? decoded : undefined;
+	} catch {
+		return undefined;
+	}
+}
+
+function oauthAccountID(auth: Extract<AuthInfo, { type: "oauth" }>) {
+	if (auth.accountId) return auth.accountId;
+	const claim = decodeJwtPayload(auth.access)?.["https://api.openai.com/auth"];
+	if (!isRecord(claim)) return undefined;
+	const accountID = claim.chatgpt_account_id;
+	return typeof accountID === "string" && accountID.length > 0
+		? accountID
+		: undefined;
+}
+
+export function sameSavedAuth(saved: AuthInfo, current: AuthInfo) {
+	if (saved.type !== current.type) return false;
+	if (saved.type === "api" && current.type === "api")
+		return saved.key === current.key;
+	if (saved.type === "wellknown" && current.type === "wellknown")
+		return saved.key === current.key && saved.token === current.token;
+	if (saved.type === "oauth" && current.type === "oauth") {
+		const savedAccountID = oauthAccountID(saved);
+		const currentAccountID = oauthAccountID(current);
+		if (savedAccountID && currentAccountID)
+			return savedAccountID === currentAccountID;
+		return saved.refresh === current.refresh;
+	}
+	return false;
+}
+
+function saveKnownNativeAuth(db: Database, providerID: string, auth: AuthInfo) {
+	const account = listAccounts(db, providerID).find((candidate) =>
+		sameSavedAuth(candidate.auth, auth),
+	);
+	if (!account) return false;
+	updateAccountAuth(db, providerID, account.alias, auth);
+	return true;
+}
+
+export function captureUnknownNativeAuth(
+	db: Database,
+	providerID: string,
+	auth: AuthInfo,
+) {
+	if (saveKnownNativeAuth(db, providerID, auth)) return false;
+	createPendingConnection(db, providerID, auth, "auth-file");
+	return true;
+}
+
+function initialAuthSnapshot(
+	db: Database,
+	authByProvider: Record<string, AuthInfo>,
+) {
+	const snapshot: Record<string, AuthInfo> = {};
+	for (const [providerID, auth] of Object.entries(authByProvider)) {
+		if (
+			listAccounts(db, providerID).some((account) =>
+				sameSavedAuth(account.auth, auth),
+			)
+		) {
+			snapshot[providerID] = auth;
+		}
+	}
+	return snapshot;
+}
+
+function createAuthWatcher(options: AuthWatcherOptions) {
+	const lastAuthSnapshot = options.initialSnapshot ?? {};
+	let inFlight = false;
+	const notify = options.showToast ?? defaultShowToast;
+	const createPending = options.createPending ?? createPendingConnection;
+	const isSuppressed =
+		options.isSuppressed ??
+		((providerID: string) => {
+			return (
+				isNativeAuthCaptureSuppressed(providerID) ||
+				isDbNativeAuthCaptureSuppressed(options.db, providerID)
+			);
+		});
+
+	return {
+		async poll() {
+			if (inFlight) return;
+			inFlight = true;
+			try {
+				const result = options.readAuth();
+				if (!result.ok) return;
+
+				const next = result.auth;
+				const changed: Array<[string, AuthInfo]> = [];
+				for (const [providerID, auth] of Object.entries(next)) {
+					if (isSuppressed(providerID)) {
+						const active = getActiveAccount(options.db, providerID);
+						if (active && sameSavedAuth(active.auth, auth))
+							updateAccountAuth(options.db, providerID, active.alias, auth);
+						lastAuthSnapshot[providerID] = auth;
+						continue;
+					}
+					if (saveKnownNativeAuth(options.db, providerID, auth)) {
+						lastAuthSnapshot[providerID] = auth;
+						continue;
+					}
+					if (!authChanged(lastAuthSnapshot[providerID], auth)) continue;
+					changed.push([providerID, auth]);
+				}
+
+				for (const [providerID, auth] of changed) {
+					createPending(options.db, providerID, auth, "auth-file");
+					lastAuthSnapshot[providerID] = auth;
+					await notify(
+						options.client,
+						`Balancer detected a ${providerID} connection. Save an alias from the TUI sidebar.`,
+						"info",
+					);
+				}
+
+				for (const providerID of Object.keys(lastAuthSnapshot)) {
+					if (!(providerID in next)) delete lastAuthSnapshot[providerID];
+				}
+			} finally {
+				inFlight = false;
+			}
+		},
+	};
+}
+
+export function startAuthWatcher(client: any, db: Database) {
+	if (started) return;
+	started = true;
+	const initial = readNativeAuth();
+	const watcher = createAuthWatcher({
+		client,
+		db,
+		initialSnapshot: initial.ok ? initialAuthSnapshot(db, initial.auth) : {},
+		readAuth: readNativeAuth,
+	});
+
+	const timer = setInterval(() => void watcher.poll(), 1_000);
+	(timer as { unref?: () => void }).unref?.();
+}
+
+export const __testParseNativeAuthContent = parseNativeAuthContent;
+
+export const __testInitialAuthSnapshot = initialAuthSnapshot;
+
+export function __testCreateAuthWatcher(
+	options: Omit<AuthWatcherOptions, "client"> & { client?: any },
+) {
+	return createAuthWatcher({ ...options, client: options.client ?? {} });
+}

package/dist/server/commands.ts

@@ -0,0 +1,58 @@
+import type { Database } from "bun:sqlite";
+import {
+	getActiveAccount,
+	listAccounts,
+	setActiveAccount,
+} from "../core/accounts";
+import { listPendingConnections } from "../core/pending";
+
+function parseArgs(input: string) {
+	return (
+		input.match(/(?:[^\s"]+|"[^"]*")+/g)?.map((x) => x.replace(/^"|"$/g, "")) ??
+		[]
+	);
+}
+
+export function runFallbackBalancerCommand(db: Database, raw: string) {
+	const [action = "help", ...args] = parseArgs(raw);
+
+	if (["help", "--help", "-h"].includes(action)) {
+		return [
+			"Balancer is TUI-first. Use fallback commands only for compatibility or troubleshooting.",
+			"Fallback commands: list, status, use <provider> <alias>, active <provider>.",
+		].join("\n");
+	}
+
+	if (action === "list") {
+		const accounts = listAccounts(db);
+		if (accounts.length === 0) return "No accounts saved.";
+		return accounts
+			.map((account) => `${account.providerID}/${account.alias}`)
+			.join("\n");
+	}
+
+	if (action === "status") {
+		return `accounts=${listAccounts(db).length} pending=${listPendingConnections(db).length}`;
+	}
+
+	if (action === "use") {
+		const [providerID, alias] = args;
+		if (!providerID || !alias) return "Usage: /balancer use <provider> <alias>";
+		try {
+			const account = setActiveAccount(db, providerID, alias);
+			return `Active account changed to ${providerID}/${account.alias}`;
+		} catch (error) {
+			return error instanceof Error ? error.message : "Account not found.";
+		}
+	}
+
+	if (action === "active") {
+		const [providerID] = args;
+		if (!providerID) return "Usage: /balancer active <provider>";
+		const account = getActiveAccount(db, providerID);
+		if (!account) return "No active account for provider.";
+		return `${providerID}/${account.alias}`;
+	}
+
+	return `Unknown command: ${action}. Use /balancer help.`;
+}

package/dist/server/fetch-patch.ts

@@ -0,0 +1,162 @@
+import type { Database } from "bun:sqlite";
+import { setActiveAccount } from "../core/accounts";
+import { getBalancingEnabled } from "../core/priority";
+import type { Account } from "../core/types";
+import { setNativeAuth, showToast } from "./native";
+import {
+	chooseFailoverAccount,
+	INTERNAL_REQUEST_HEADER,
+	markRateLimited,
+	RETRYABLE_STATUS,
+	takePendingRequest,
+} from "./request-balancer";
+
+let fetchPatched = false;
+
+function headerEntries(headers: Headers) {
+	return Array.from(headers.entries()).map(
+		([key, value]) => [key.toLowerCase(), value] as const,
+	);
+}
+
+function applyAuthToHeaders(
+	headers: Headers,
+	account: Account,
+	options: { preserveExistingOAuth?: boolean } = {},
+) {
+	const auth = account.auth;
+	const entries = headerEntries(headers);
+
+	if (auth.type === "oauth") {
+		if (options.preserveExistingOAuth && headers.has("authorization")) return;
+		headers.set("authorization", `Bearer ${auth.access}`);
+		return;
+	}
+
+	if (auth.type === "wellknown") {
+		if (headers.has("authorization")) {
+			headers.set("authorization", `Bearer ${auth.token}`);
+		} else {
+			headers.set(auth.key, auth.token);
+		}
+		return;
+	}
+
+	const authHeaderNames = [
+		"authorization",
+		"x-api-key",
+		"api-key",
+		"x-goog-api-key",
+		"x-stainless-api-key",
+		"anthropic-api-key",
+		"cohere-api-key",
+	];
+
+	let changed = false;
+	for (const name of authHeaderNames) {
+		const current = entries.find(([key]) => key === name)?.[1];
+		if (!current) continue;
+		headers.set(
+			name,
+			name === "authorization" && current.toLowerCase().startsWith("bearer ")
+				? `Bearer ${auth.key}`
+				: auth.key,
+		);
+		changed = true;
+	}
+
+	if (!changed) headers.set("authorization", `Bearer ${auth.key}`);
+}
+
+export function __testResetFetchPatch() {
+	fetchPatched = false;
+}
+
+function headersFrom(input: RequestInfo | URL, init?: RequestInit) {
+	if (init?.headers) return new Headers(init.headers);
+	if (typeof Request !== "undefined" && input instanceof Request) {
+		return new Headers(input.headers);
+	}
+	return new Headers();
+}
+
+function cloneRequestInput(
+	input: RequestInfo | URL,
+	init: RequestInit | undefined,
+	headers: Headers,
+): [RequestInfo | URL, RequestInit | undefined] {
+	if (typeof Request !== "undefined" && input instanceof Request) {
+		return [new Request(input, { ...init, headers }), undefined];
+	}
+	return [input, { ...init, headers }];
+}
+
+function retryAfterMs(response: Response) {
+	const retryAfter = response.headers.get("retry-after");
+	if (!retryAfter) return 60_000;
+
+	const seconds = Number(retryAfter);
+	if (Number.isFinite(seconds)) return Math.max(1_000, seconds * 1000);
+
+	const date = Date.parse(retryAfter);
+	if (Number.isFinite(date)) return Math.max(1_000, date - Date.now());
+
+	return 60_000;
+}
+
+export function installFetchPatch(db: Database, client: any) {
+	if (fetchPatched) return;
+	fetchPatched = true;
+	const originalFetch = globalThis.fetch.bind(globalThis);
+
+	globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => {
+		const headers = headersFrom(input, init);
+		const requestID = headers.get(INTERNAL_REQUEST_HEADER);
+		if (!requestID) return originalFetch(input, init);
+
+		const pending = takePendingRequest(requestID);
+		headers.delete(INTERNAL_REQUEST_HEADER);
+		if (!pending?.account) {
+			const [nextInput, nextInit] = cloneRequestInput(input, init, headers);
+			return originalFetch(nextInput, nextInit);
+		}
+
+		let account = pending.account;
+		const maxAttempts = 3;
+		for (let attempt = 0; attempt < maxAttempts; attempt++) {
+			const attemptHeaders = new Headers(headers);
+			applyAuthToHeaders(attemptHeaders, account, {
+				preserveExistingOAuth: attempt === 0,
+			});
+			const [nextInput, nextInit] = cloneRequestInput(
+				input,
+				init,
+				attemptHeaders,
+			);
+			const response = await originalFetch(nextInput, nextInit);
+			if (!RETRYABLE_STATUS.has(response.status)) return response;
+
+			markRateLimited(
+				db,
+				account.providerID,
+				account.alias,
+				retryAfterMs(response),
+			);
+			if (!getBalancingEnabled(db)) return response;
+			if (attempt === maxAttempts - 1) return response;
+
+			const next = chooseFailoverAccount(db, account.providerID, account.alias);
+			if (!next) return response;
+
+			account = setActiveAccount(db, next.providerID, next.alias);
+			await setNativeAuth(client, account.providerID, account.auth, db);
+			await showToast(
+				client,
+				`Balancer: ${pending.providerID}/${pending.account.alias} is rate limited. Switching to ${account.alias}.`,
+				"warning",
+			);
+		}
+
+		throw new Error("Balancer retry loop exited unexpectedly");
+	}) as typeof globalThis.fetch;
+}

package/dist/server/index.ts

@@ -0,0 +1,134 @@
+import type { Database } from "bun:sqlite";
+import {
+	type Config,
+	type Hooks,
+	type Plugin,
+	tool,
+} from "@opencode-ai/plugin";
+import {
+	getActiveAccount,
+	getSelectedAccount,
+	getSelectedModel,
+	setActiveAccount,
+} from "../core/accounts";
+import { openBalancerDatabase } from "../core/database";
+import { storePath } from "../core/path";
+import { getBalancingEnabled, resolveActiveSelection } from "../core/priority";
+import { migrate } from "../core/schema";
+import { runFallbackBalancerCommand } from "./commands";
+import { installFetchPatch } from "./fetch-patch";
+import { setNativeAuth, showToast } from "./native";
+import {
+	BALANCER_METADATA_KEY,
+	INTERNAL_REQUEST_HEADER,
+	setPendingRequest,
+} from "./request-balancer";
+
+export function configureFallbackCommand(cfg: Config) {
+	if (!cfg.command?.balancer) return;
+}
+
+function runSafeFallbackBalancerCommand(db: Database, raw: string) {
+	try {
+		return runFallbackBalancerCommand(db, raw);
+	} catch (error) {
+		return error instanceof Error ? error.message : "Balancer command failed.";
+	}
+}
+
+export function createServerHooks({
+	db,
+	client,
+}: {
+	db: Database;
+	client: any;
+}): Hooks {
+	return {
+		"chat.headers": async (input, output) => {
+			const providerID = input.model.providerID;
+			const account = getActiveAccount(db, providerID);
+			if (!account) return;
+
+			if (account.auth.type === "oauth") {
+				await setNativeAuth(client, providerID, account.auth, db);
+			}
+
+			const requestID = crypto.randomUUID();
+			setPendingRequest(requestID, { account, providerID });
+			output.headers[INTERNAL_REQUEST_HEADER] = requestID;
+		},
+
+		"chat.message": async (_input, output) => {
+			// Balancing on: the priority list decides the provider/model for every
+			// message (recomputed per message -> failover and recovery for free).
+			if (getBalancingEnabled(db)) {
+				const selection = resolveActiveSelection(
+					db,
+					undefined,
+					output.message.model?.providerID,
+				);
+				if (!selection) return;
+				setActiveAccount(db, selection.providerID, selection.account.alias);
+				output.message.model = {
+					modelID: selection.modelID,
+					providerID: selection.providerID,
+				};
+				return;
+			}
+
+			// Balancing off: keep opencode's native choice; only fill when missing.
+			if (output.message.model?.providerID && output.message.model?.modelID)
+				return;
+
+			const selected = getSelectedAccount(db);
+			if (!selected) return;
+			const model = getSelectedModel(db, selected.providerID);
+			if (!model) return;
+
+			output.message.model = {
+				modelID: model.modelID,
+				providerID: model.providerID,
+			};
+		},
+
+		"command.execute.before": async (input, output) => {
+			if (input.command !== "balancer") return;
+			const result = runSafeFallbackBalancerCommand(db, input.arguments);
+			output.parts.length = 0;
+			await showToast(client, result.split("\n")[0] ?? result, "info");
+			throw new Error(`[balancer]\n${result}`);
+		},
+		config: async (cfg) => {
+			configureFallbackCommand(cfg);
+		},
+
+		"experimental.chat.messages.transform": async (_input, output) => {
+			output.messages = output.messages.filter((message) => {
+				return !message.parts.some((part: any) => {
+					return part?.metadata?.[BALANCER_METADATA_KEY] === true;
+				});
+			});
+		},
+
+		tool: {
+			balancer_command: tool({
+				args: {
+					command: tool.schema
+						.string()
+						.describe("Command arguments for /balancer."),
+				},
+				description: "Run fallback account balancer commands.",
+				execute: async (args) =>
+					runSafeFallbackBalancerCommand(db, args.command),
+			}),
+		},
+	};
+}
+
+export const serverPlugin = (async ({ client }) => {
+	const db = openBalancerDatabase(storePath());
+	migrate(db);
+	installFetchPatch(db, client);
+
+	return createServerHooks({ client, db });
+}) satisfies Plugin;