@henryqw/pi-multi-codex 3.1.0 → 3.1.1

package/extensions/multi-codex.ts

@@ -260,7 +260,6 @@ async function loadState(signal?: AbortSignal): Promise<UsageState> {
 
 async function saveState(state: UsageState, signal?: AbortSignal): Promise<void> {
 	signal?.throwIfAborted();
-	if (state.untrusted) return;
 	const file = cachePath();
 	const data = `${JSON.stringify({
 		slots: [...state.slots.values()].sort((a, b) => a.slot - b.slot),
@@ -272,11 +271,11 @@ async function saveState(state: UsageState, signal?: AbortSignal): Promise<void>
 async function withCacheMutex<T>(operation: (signal: AbortSignal) => Promise<T>, signal?: AbortSignal): Promise<T> {
 	signal?.throwIfAborted();
 	await mkdir(extensionConfigDir("pi-multi-codex"), { recursive: true, mode: 0o700 });
-	let release: (() => Promise<void>) | undefined;
-	while (!release) {
+	const compromised = new AbortController();
+	let release: () => Promise<void>;
+	for (;;) {
 		signal?.throwIfAborted();
 		try {
-			const compromised = new AbortController();
 			release = await lock(cachePath(), {
 				lockfilePath: cacheMutexPath(),
 				realpath: false,
@@ -285,19 +284,19 @@ async function withCacheMutex<T>(operation: (signal: AbortSignal) => Promise<T>,
 				retries: 0,
 				onCompromised: (error) => compromised.abort(error),
 			});
-			const operationSignal = signal ? AbortSignal.any([signal, compromised.signal]) : compromised.signal;
-			try {
-				operationSignal.throwIfAborted();
-				return await raceWithSignal(operation(operationSignal), operationSignal);
-			} finally {
-				await release().catch(() => undefined);
-			}
+			break;
 		} catch (error) {
 			if (!(error && typeof error === "object" && "code" in error && String(error.code) === "ELOCKED")) throw error;
 			await delay(CACHE_MUTEX_RETRY_MS, undefined, { signal, ref: false });
 		}
 	}
-	throw new Error("Cache mutex acquisition ended unexpectedly.");
+	const operationSignal = signal ? AbortSignal.any([signal, compromised.signal]) : compromised.signal;
+	try {
+		operationSignal.throwIfAborted();
+		return await raceWithSignal(operation(operationSignal), operationSignal);
+	} finally {
+		await release().catch(() => undefined);
+	}
 }
 
 function isMeasured(snapshot: UsageSnapshot | undefined): snapshot is MeasuredSnapshot {
@@ -434,15 +433,10 @@ class CodexQuotaStatus {
 
 	private change<T>(change: (state: UsageState) => { value: T; write: boolean }, signal?: AbortSignal): Promise<T | undefined> {
 		const operation = this.writes.then(() => withCacheMutex(async (operationSignal) => {
-			operationSignal.throwIfAborted();
 			const state = await loadState(operationSignal);
-			operationSignal.throwIfAborted();
 			if (state.untrusted) return undefined;
 			const result = change(state);
-			if (result.write) {
-				operationSignal.throwIfAborted();
-				await saveState(state, operationSignal);
-			}
+			if (result.write) await saveState(state, operationSignal);
 			return result.value;
 		}, signal));
 		this.writes = operation.then(() => undefined, () => undefined);
@@ -468,7 +462,6 @@ class CodexQuotaStatus {
 		const refresh = this.refreshContext;
 		if (!refresh) return;
 		void this.refreshDue(refresh)
-			.catch(() => undefined)
 			.then(() => {
 				if (this.refreshContext === refresh) this.scheduleRefresh();
 			});
@@ -507,20 +500,21 @@ class CodexQuotaStatus {
 			this.timer = undefined;
 			if (this.refreshContext === refresh) this.requestRefresh();
 		}, delay);
-		this.timer.unref?.();
+		this.timer.unref();
 	}
 
 	private async refreshDue(refresh: QuotaRefresh): Promise<void> {
 		const credentials = readCodexCredentials();
 		const state = await this.state();
-		await Promise.all([...credentials.entries()].map(async ([slot, credential]) => {
-			if (this.refreshContext !== refresh || this.active.has(slot)) return;
+		const now = Date.now();
+		for (const [slot, credential] of credentials) {
+			if (this.refreshContext !== refresh || this.active.has(slot)) continue;
 			const identity = identityFor(credential);
-			if (!identity) return;
+			if (!identity) continue;
 			const snapshot = state.slots.get(slot);
-			if (isFresh(snapshot, identity, Date.now()) || checkedRecently(snapshot, identity, Date.now())) return;
+			if (isFresh(snapshot, identity, now) || checkedRecently(snapshot, identity, now)) continue;
 			this.launch(slot, identity, refresh);
-		})).catch(() => undefined);
+		}
 	}
 
 	private launch(slot: number, identity: SlotIdentity, refresh: QuotaRefresh): void {
@@ -569,7 +563,7 @@ class CodexQuotaStatus {
 				if (!kept) this.stopHeartbeat(slot, owner);
 			}).catch(() => this.stopHeartbeat(slot, owner));
 		}, HEARTBEAT_MS);
-		timer.unref?.();
+		timer.unref();
 		this.heartbeats.set(slot, { owner, timer });
 	}
 
@@ -662,7 +656,7 @@ class CodexQuotaStatus {
 				signal: fetchSignal,
 			});
 			if (!response.ok) throw new Error(`Codex usage request failed (${response.status}).`);
-			const usage = parseCodexUsage(await raceWithSignal(response.json(), fetchSignal));
+			const usage = parseCodexUsage(await response.json());
 			if (!usage) throw new Error("Codex usage response has no usable seven-day window.");
 			if (!this.identityStillCurrent(slot, identity.accountHash) || !(await this.stillOwn(slot, owner, identity.accountHash))) {
 				discard = true;
@@ -703,10 +697,6 @@ class CodexQuotaStatus {
 	}
 }
 
-function nativeModel(model: CodexModel): CodexModel {
-	return model.provider === NATIVE_PROVIDER_ID ? model : { ...model, provider: NATIVE_PROVIDER_ID };
-}
-
 function aliasModel(model: CodexModel, provider: string): CodexModel {
 	return model.provider === provider ? model : { ...model, provider };
 }
@@ -717,10 +707,7 @@ function nativeContext(context: TranscriptContext, alias: string): TranscriptCon
 		messages: context.messages.map((message) => {
 			if (message.role !== "assistant") return message;
 			const provider = message.provider === alias ? NATIVE_PROVIDER_ID : message.provider === NATIVE_PROVIDER_ID ? alias : message.provider;
-			if (provider === message.provider) return message;
-			const rewritten = { ...message, provider };
-			if (message.deferred) rewritten.deferred = { ...message.deferred, provider };
-			return rewritten;
+			return provider === message.provider ? message : aliasMessage(message, provider);
 		}),
 	};
 }
@@ -788,28 +775,28 @@ export function createCodexAliasProvider(native: CodexProvider, slot: number): C
 		auth: native.auth,
 		getModels: () => models,
 		stream: (model, context, options?: ApiStreamOptions<"openai-codex-responses">) =>
-			aliasStream(native.stream(nativeModel(model), nativeContext(context, providerId), options), model, providerId),
+			aliasStream(native.stream(aliasModel(model, NATIVE_PROVIDER_ID), nativeContext(context, providerId), options), model, providerId),
 		streamSimple: (model, context, options?: SimpleStreamOptions) =>
-			aliasStream(native.streamSimple(nativeModel(model), nativeContext(context, providerId), options), model, providerId),
+			aliasStream(native.streamSimple(aliasModel(model, NATIVE_PROVIDER_ID), nativeContext(context, providerId), options), model, providerId),
 	};
 
 	if (native.filterModels) {
 		provider.filterModels = (available, credential) =>
 			native
-				.filterModels!(available.map(nativeModel), credential)
+				.filterModels!(available.map((model) => aliasModel(model, NATIVE_PROVIDER_ID)), credential)
 				.map((model) => aliasModel(model, providerId));
 	}
 	if (native.fetchDeferred) {
 		provider.fetchDeferred = (model, handle, options?: DeferredFetchOptions) =>
 			aliasStream(
-				native.fetchDeferred!(nativeModel(model), { ...handle, provider: NATIVE_PROVIDER_ID }, options),
+				native.fetchDeferred!(aliasModel(model, NATIVE_PROVIDER_ID), { ...handle, provider: NATIVE_PROVIDER_ID }, options),
 				model,
 				providerId,
 			);
 	}
 	if (native.cancelDeferred) {
 		provider.cancelDeferred = (model, handle, options?: DeferredCancelOptions) =>
-			native.cancelDeferred!(nativeModel(model), { ...handle, provider: NATIVE_PROVIDER_ID }, options);
+			native.cancelDeferred!(aliasModel(model, NATIVE_PROVIDER_ID), { ...handle, provider: NATIVE_PROVIDER_ID }, options);
 	}
 	return provider;
 }
@@ -824,9 +811,9 @@ function isManagedModel(model: Model<any> | undefined): model is CodexModel {
 }
 
 function sessionHasAgentWork(ctx: ExtensionContext): boolean {
-	return ctx.sessionManager?.getBranch?.().some((entry) =>
+	return ctx.sessionManager.getBranch().some((entry) =>
 		entry.type === "message" || (entry.type === "custom" && entry.customType === AGENT_STARTED_ENTRY),
-	) ?? false;
+	);
 }
 
 export default function multiCodex(pi: ExtensionAPI): void {
@@ -863,7 +850,7 @@ export default function multiCodex(pi: ExtensionAPI): void {
 	};
 
 	const allowsModel = (ctx: ExtensionContext, model: Model<any>, slot: number): boolean => {
-		const scopedModels = ctx.scopedModels ?? [];
+		const scopedModels = ctx.scopedModels;
 		return scopedModels.length === 0 || scopedModels.some(({ model: scoped }) => scoped.provider === providerForSlot(slot) && scoped.id === model.id);
 	};
 
@@ -919,17 +906,17 @@ export default function multiCodex(pi: ExtensionAPI): void {
 		const now = Date.now();
 		if (isFiveHourLimited(snapshot, now)) {
 			const text = `${prefix} · 5h limit · ${formatDuration(snapshot.limitedUntil! - now)}`;
-			return ctx.ui.theme?.fg ? ctx.ui.theme.fg("error", text) : text;
+			return ctx.ui.theme.fg("error", text);
 		}
 		if (!isFresh(snapshot, identity, now)) return `${prefix} · stale`;
 		const display = displayedReset(snapshot, now);
 		const text = `${prefix} · ${formatPercent(snapshot.remaining)}% · ${display.label} ${formatDuration(display.time - now)}`;
 		const color = snapshot.remaining >= 50 ? "success" : snapshot.remaining >= 25 ? "warning" : "error";
-		return ctx.ui.theme?.fg ? ctx.ui.theme.fg(color, text) : text;
+		return ctx.ui.theme.fg(color, text);
 	};
 
 	const updateFooter = (ctx: ExtensionContext | undefined = sessionContext): void => {
-		if (ctx) ctx.ui.setStatus?.("pi-multi-codex", footerText(ctx));
+		if (ctx) ctx.ui.setStatus("pi-multi-codex", footerText(ctx));
 	};
 
 	const setModelAutomatically = async (model: Model<any>): Promise<boolean> => {
@@ -1099,10 +1086,8 @@ export default function multiCodex(pi: ExtensionAPI): void {
 				return;
 			}
 
-			const authenticated = [...syncSlots()]
-				.filter((slot) => slot === 1 || registered.has(slot))
-				.sort((left, right) => left - right);
-			const scopedModels = ctx.scopedModels ?? [];
+			const authenticated = [...syncSlots()].sort((left, right) => left - right);
+			const scopedModels = ctx.scopedModels;
 			const slots = authenticated.filter((slot) => allowsModel(ctx, model, slot));
 			const currentSlot = slotForProvider(model.provider);
 			if (slots.length === 0 || (scopedModels.length > 0 && authenticated.some((slot) => !allowsModel(ctx, model, slot)) && slots.every((slot) => slot === currentSlot))) {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@henryqw/pi-multi-codex",
-  "version": "3.1.0",
+  "version": "3.1.1",
   "description": "Use multiple ChatGPT Codex OAuth accounts in Pi.",
   "keywords": [
     "pi-package",