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

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`;
+}