@victor-software-house/pi-multicodex 1.0.3 → 1.0.4

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@victor-software-house/pi-multicodex",
-	"version": "1.0.3",
+	"version": "1.0.4",
 	"description": "Codex account rotation extension for pi",
 	"license": "MIT",
 	"type": "module",
@@ -30,7 +30,20 @@
 	},
 	"author": "Victor",
 	"files": [
+		"abort-utils.ts",
+		"account-manager.ts",
+		"browser.ts",
+		"commands.ts",
+		"extension.ts",
+		"hooks.ts",
 		"index.ts",
+		"provider.ts",
+		"quota.ts",
+		"selection.ts",
+		"storage.ts",
+		"stream-wrapper.ts",
+		"usage-client.ts",
+		"usage.ts",
 		"README.md",
 		"LICENSE",
 		"assets/**"
@@ -41,8 +54,8 @@
 		"tsgo": "tsgo -p tsconfig.json",
 		"check": "pnpm lint && pnpm tsgo && pnpm test",
 		"pack:dry": "npm pack --dry-run",
-		"publish:dry": "bun ./scripts/publish.ts --dry-run",
-		"publish:release": "bun ./scripts/publish.ts"
+		"release:dry": "bun ./scripts/publish.ts --dry-run",
+		"release:prepare": "bun ./scripts/publish.ts"
 	},
 	"peerDependencies": {
 		"@mariozechner/pi-ai": "*",

package/provider.ts

@@ -0,0 +1,75 @@
+import {
+	type Api,
+	type AssistantMessageEventStream,
+	type Context,
+	getApiProvider,
+	getModels,
+	type Model,
+	type SimpleStreamOptions,
+} from "@mariozechner/pi-ai";
+import type { AccountManager } from "./account-manager";
+import { createStreamWrapper } from "./stream-wrapper";
+
+export const PROVIDER_ID = "multicodex";
+
+export interface ProviderModelDef {
+	id: string;
+	name: string;
+	reasoning: boolean;
+	input: ("text" | "image")[];
+	cost: {
+		input: number;
+		output: number;
+		cacheRead: number;
+		cacheWrite: number;
+	};
+	contextWindow: number;
+	maxTokens: number;
+}
+
+export function getOpenAICodexMirror(): {
+	baseUrl: string;
+	models: ProviderModelDef[];
+} {
+	const sourceModels = getModels("openai-codex");
+	return {
+		baseUrl: sourceModels[0]?.baseUrl || "https://chatgpt.com/backend-api",
+		models: sourceModels.map((model) => ({
+			id: model.id,
+			name: model.name,
+			reasoning: model.reasoning,
+			input: model.input,
+			cost: model.cost,
+			contextWindow: model.contextWindow,
+			maxTokens: model.maxTokens,
+		})),
+	};
+}
+
+export function buildMulticodexProviderConfig(accountManager: AccountManager): {
+	baseUrl: string;
+	apiKey: string;
+	api: "openai-codex-responses";
+	streamSimple: (
+		model: Model<Api>,
+		context: Context,
+		options?: SimpleStreamOptions,
+	) => AssistantMessageEventStream;
+	models: ProviderModelDef[];
+} {
+	const mirror = getOpenAICodexMirror();
+	const baseProvider = getApiProvider("openai-codex-responses");
+	if (!baseProvider) {
+		throw new Error(
+			"OpenAI Codex provider not available. Please update pi to include openai-codex support.",
+		);
+	}
+
+	return {
+		baseUrl: mirror.baseUrl,
+		apiKey: "managed-by-extension",
+		api: "openai-codex-responses",
+		streamSimple: createStreamWrapper(accountManager, baseProvider),
+		models: mirror.models,
+	};
+}

package/quota.ts

@@ -0,0 +1,5 @@
+export function isQuotaErrorMessage(message: string): boolean {
+	return /\b429\b|quota|usage limit|rate.?limit|too many requests|limit reached/i.test(
+		message,
+	);
+}

package/selection.ts

@@ -0,0 +1,69 @@
+import type { Account } from "./storage";
+import {
+	type CodexUsageSnapshot,
+	getWeeklyResetAt,
+	isUsageUntouched,
+} from "./usage";
+
+export function isAccountAvailable(account: Account, now: number): boolean {
+	return !account.quotaExhaustedUntil || account.quotaExhaustedUntil <= now;
+}
+
+function pickRandomAccount(accounts: Account[]): Account | undefined {
+	if (accounts.length === 0) return undefined;
+	return accounts[Math.floor(Math.random() * accounts.length)];
+}
+
+function pickEarliestWeeklyResetAccount(
+	accounts: Account[],
+	usageByEmail: Map<string, CodexUsageSnapshot>,
+): Account | undefined {
+	const candidates = accounts
+		.map((account) => ({
+			account,
+			resetAt: getWeeklyResetAt(usageByEmail.get(account.email)),
+		}))
+		.filter(
+			(entry): entry is { account: Account; resetAt: number } =>
+				typeof entry.resetAt === "number",
+		)
+		.sort((a, b) => a.resetAt - b.resetAt);
+
+	return candidates[0]?.account;
+}
+
+export function pickBestAccount(
+	accounts: Account[],
+	usageByEmail: Map<string, CodexUsageSnapshot>,
+	options?: { excludeEmails?: Set<string>; now?: number },
+): Account | undefined {
+	const now = options?.now ?? Date.now();
+	const available = accounts.filter(
+		(account) =>
+			isAccountAvailable(account, now) &&
+			!options?.excludeEmails?.has(account.email),
+	);
+	if (available.length === 0) return undefined;
+
+	const withUsage = available.filter((account) =>
+		usageByEmail.has(account.email),
+	);
+	const untouched = withUsage.filter((account) =>
+		isUsageUntouched(usageByEmail.get(account.email)),
+	);
+
+	if (untouched.length > 0) {
+		return (
+			pickEarliestWeeklyResetAccount(untouched, usageByEmail) ??
+			pickRandomAccount(untouched)
+		);
+	}
+
+	const earliestWeeklyReset = pickEarliestWeeklyResetAccount(
+		withUsage,
+		usageByEmail,
+	);
+	if (earliestWeeklyReset) return earliestWeeklyReset;
+
+	return pickRandomAccount(available);
+}

package/storage.ts

@@ -0,0 +1,49 @@
+import * as fs from "node:fs";
+import * as os from "node:os";
+import * as path from "node:path";
+
+export interface Account {
+	email: string;
+	accessToken: string;
+	refreshToken: string;
+	expiresAt: number;
+	accountId?: string;
+	lastUsed?: number;
+	quotaExhaustedUntil?: number;
+}
+
+export interface StorageData {
+	accounts: Account[];
+	activeEmail?: string;
+}
+
+export const STORAGE_FILE = path.join(
+	os.homedir(),
+	".pi",
+	"agent",
+	"codex-accounts.json",
+);
+
+export function loadStorage(): StorageData {
+	try {
+		if (fs.existsSync(STORAGE_FILE)) {
+			return JSON.parse(fs.readFileSync(STORAGE_FILE, "utf-8")) as StorageData;
+		}
+	} catch (error) {
+		console.error("Failed to load multicodex accounts:", error);
+	}
+
+	return { accounts: [] };
+}
+
+export function saveStorage(data: StorageData): void {
+	try {
+		const dir = path.dirname(STORAGE_FILE);
+		if (!fs.existsSync(dir)) {
+			fs.mkdirSync(dir, { recursive: true });
+		}
+		fs.writeFileSync(STORAGE_FILE, JSON.stringify(data, null, 2));
+	} catch (error) {
+		console.error("Failed to save multicodex accounts:", error);
+	}
+}

package/stream-wrapper.ts

@@ -0,0 +1,191 @@
+import {
+	type Api,
+	type AssistantMessage,
+	type AssistantMessageEvent,
+	type AssistantMessageEventStream,
+	type Context,
+	createAssistantMessageEventStream,
+	type Model,
+	type SimpleStreamOptions,
+} from "@mariozechner/pi-ai";
+import { createLinkedAbortController } from "./abort-utils";
+import type { AccountManager } from "./account-manager";
+import { isQuotaErrorMessage } from "./quota";
+
+const MAX_ROTATION_RETRIES = 5;
+
+type ApiProviderRef = {
+	streamSimple: (
+		model: Model<Api>,
+		context: Context,
+		options?: SimpleStreamOptions,
+	) => AssistantMessageEventStream;
+};
+
+function withProvider(
+	event: AssistantMessageEvent,
+	provider: string,
+): AssistantMessageEvent {
+	if ("partial" in event) {
+		return { ...event, partial: { ...event.partial, provider } };
+	}
+	if (event.type === "done") {
+		return { ...event, message: { ...event.message, provider } };
+	}
+	if (event.type === "error") {
+		return { ...event, error: { ...event.error, provider } };
+	}
+	return event;
+}
+
+function createErrorAssistantMessage(
+	model: Model<Api>,
+	message: string,
+): AssistantMessage {
+	return {
+		role: "assistant",
+		content: [],
+		api: model.api,
+		provider: model.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: "error",
+		errorMessage: message,
+		timestamp: Date.now(),
+	};
+}
+
+function getErrorMessage(error: unknown): string {
+	if (error instanceof Error) return error.message;
+	return typeof error === "string" ? error : JSON.stringify(error);
+}
+
+export function createStreamWrapper(
+	accountManager: AccountManager,
+	baseProvider: ApiProviderRef,
+) {
+	return (
+		model: Model<Api>,
+		context: Context,
+		options?: SimpleStreamOptions,
+	): AssistantMessageEventStream => {
+		const stream = createAssistantMessageEventStream();
+
+		(async () => {
+			try {
+				const excludedEmails = new Set<string>();
+				for (let attempt = 0; attempt <= MAX_ROTATION_RETRIES; attempt++) {
+					const now = Date.now();
+					const manual = accountManager.getAvailableManualAccount({
+						excludeEmails: excludedEmails,
+						now,
+					});
+					const usingManual = Boolean(manual);
+					let account = manual;
+					if (!account) {
+						if (accountManager.hasManualAccount()) {
+							accountManager.clearManualAccount();
+						}
+						account = await accountManager.activateBestAccount({
+							excludeEmails: excludedEmails,
+							signal: options?.signal,
+						});
+					}
+					if (!account) {
+						throw new Error(
+							"No available Multicodex accounts. Please use /multicodex-login.",
+						);
+					}
+
+					const token = await accountManager.ensureValidToken(account);
+					const abortController = createLinkedAbortController(options?.signal);
+
+					const internalModel: Model<"openai-codex-responses"> = {
+						...(model as Model<"openai-codex-responses">),
+						provider: "openai-codex",
+						api: "openai-codex-responses",
+					};
+
+					const inner = baseProvider.streamSimple(
+						{
+							...internalModel,
+							headers: {
+								...(internalModel.headers || {}),
+								"X-Multicodex-Account": account.email,
+							},
+						},
+						context,
+						{
+							...options,
+							apiKey: token,
+							signal: abortController.signal,
+						},
+					);
+
+					let forwardedAny = false;
+					let retry = false;
+
+					for await (const event of inner) {
+						if (event.type === "error") {
+							const msg = event.error.errorMessage || "";
+							const isQuota = isQuotaErrorMessage(msg);
+
+							if (isQuota && !forwardedAny && attempt < MAX_ROTATION_RETRIES) {
+								await accountManager.handleQuotaExceeded(account, {
+									signal: options?.signal,
+								});
+								if (usingManual) {
+									accountManager.clearManualAccount();
+								}
+								excludedEmails.add(account.email);
+								abortController.abort();
+								retry = true;
+								break;
+							}
+
+							stream.push(withProvider(event, model.provider));
+							stream.end();
+							return;
+						}
+
+						forwardedAny = true;
+						stream.push(withProvider(event, model.provider));
+
+						if (event.type === "done") {
+							stream.end();
+							return;
+						}
+					}
+
+					if (retry) {
+						continue;
+					}
+
+					stream.end();
+					return;
+				}
+			} catch (error) {
+				const message = getErrorMessage(error);
+				const errorEvent: AssistantMessageEvent = {
+					type: "error",
+					reason: "error",
+					error: createErrorAssistantMessage(
+						model,
+						`Multicodex failed: ${message}`,
+					),
+				};
+				stream.push(withProvider(errorEvent, model.provider));
+				stream.end();
+			}
+		})();
+
+		return stream;
+	};
+}

package/usage-client.ts

@@ -0,0 +1,50 @@
+import { createTimeoutController } from "./abort-utils";
+import { type CodexUsageSnapshot, parseCodexUsageResponse } from "./usage";
+
+interface WhamUsageResponse {
+	rate_limit?: {
+		primary_window?: {
+			reset_at?: number;
+			used_percent?: number;
+		};
+		secondary_window?: {
+			reset_at?: number;
+			used_percent?: number;
+		};
+	};
+}
+
+export async function fetchCodexUsage(
+	accessToken: string,
+	accountId: string | undefined,
+	options?: { signal?: AbortSignal; timeoutMs?: number },
+): Promise<CodexUsageSnapshot> {
+	const { controller, clear } = createTimeoutController(
+		options?.signal,
+		options?.timeoutMs ?? 10_000,
+	);
+
+	try {
+		const headers: Record<string, string> = {
+			Authorization: `Bearer ${accessToken}`,
+			Accept: "application/json",
+		};
+		if (accountId) {
+			headers["ChatGPT-Account-Id"] = accountId;
+		}
+
+		const response = await fetch("https://chatgpt.com/backend-api/wham/usage", {
+			headers,
+			signal: controller.signal,
+		});
+
+		if (!response.ok) {
+			throw new Error(`Usage request failed: ${response.status}`);
+		}
+
+		const data = (await response.json()) as WhamUsageResponse;
+		return { ...parseCodexUsageResponse(data), fetchedAt: Date.now() };
+	} finally {
+		clear();
+	}
+}

package/usage.ts

@@ -0,0 +1,86 @@
+interface CodexUsageWindow {
+	usedPercent?: number;
+	resetAt?: number;
+}
+
+export interface CodexUsageSnapshot {
+	primary?: CodexUsageWindow;
+	secondary?: CodexUsageWindow;
+	fetchedAt: number;
+}
+
+interface WhamUsageResponse {
+	rate_limit?: {
+		primary_window?: WhamUsageWindow;
+		secondary_window?: WhamUsageWindow;
+	};
+}
+
+type WhamUsageWindow = {
+	reset_at?: number;
+	used_percent?: number;
+};
+
+function normalizeUsedPercent(value?: number): number | undefined {
+	if (typeof value !== "number" || !Number.isFinite(value)) return undefined;
+	return Math.min(100, Math.max(0, value));
+}
+
+function normalizeResetAt(value?: number): number | undefined {
+	if (typeof value !== "number" || !Number.isFinite(value)) return undefined;
+	return value * 1000;
+}
+
+function parseUsageWindow(
+	window?: WhamUsageWindow,
+): CodexUsageWindow | undefined {
+	if (!window) return undefined;
+	const usedPercent = normalizeUsedPercent(window.used_percent);
+	const resetAt = normalizeResetAt(window.reset_at);
+	if (usedPercent === undefined && resetAt === undefined) return undefined;
+	return { usedPercent, resetAt };
+}
+
+export function parseCodexUsageResponse(
+	data: WhamUsageResponse,
+): Omit<CodexUsageSnapshot, "fetchedAt"> {
+	return {
+		primary: parseUsageWindow(data.rate_limit?.primary_window),
+		secondary: parseUsageWindow(data.rate_limit?.secondary_window),
+	};
+}
+
+export function isUsageUntouched(usage?: CodexUsageSnapshot): boolean {
+	const primary = usage?.primary?.usedPercent;
+	const secondary = usage?.secondary?.usedPercent;
+	if (primary === undefined || secondary === undefined) return false;
+	return primary === 0 && secondary === 0;
+}
+
+export function getNextResetAt(usage?: CodexUsageSnapshot): number | undefined {
+	const candidates = [
+		usage?.primary?.resetAt,
+		usage?.secondary?.resetAt,
+	].filter((value): value is number => typeof value === "number");
+	if (candidates.length === 0) return undefined;
+	return Math.min(...candidates);
+}
+
+export function getWeeklyResetAt(
+	usage?: CodexUsageSnapshot,
+): number | undefined {
+	const resetAt = usage?.secondary?.resetAt;
+	return typeof resetAt === "number" ? resetAt : undefined;
+}
+
+export function formatResetAt(resetAt?: number): string {
+	if (!resetAt) return "unknown";
+	const diffMs = resetAt - Date.now();
+	if (diffMs <= 0) return "now";
+	const diffMinutes = Math.max(1, Math.round(diffMs / 60000));
+	if (diffMinutes < 60) return `in ${diffMinutes}m`;
+	const diffHours = Math.round(diffMinutes / 60);
+	if (diffHours < 48) return `in ${diffHours}h`;
+	const diffDays = Math.round(diffHours / 24);
+	return `in ${diffDays}d`;
+}