@maxiaochao/pi-toolkit 0.11.7 → 0.11.9

package/packages/openai-account-pool/tests/helpers/refresh-child.mjs

@@ -0,0 +1,29 @@
+import { writeFileSync } from "node:fs";
+import { join } from "node:path";
+
+import { ensureFreshAccountCredential, loadPoolSnapshot } from "../../src/auth.ts";
+import { withAuthTransaction } from "../../src/store.ts";
+import { fakeAccessToken } from "./fixtures.mjs";
+
+/** One pi process: refresh the pooled credential, and write an unrelated provider entry. */
+const [authPath, accountId, markerDir, label] = process.argv.slice(2);
+const snapshot = await loadPoolSnapshot(authPath);
+const account = snapshot.accounts.find((entry) => entry.id === accountId);
+
+const fresh = await ensureFreshAccountCredential({
+	authPath,
+	account,
+	refresh: async (credential) => {
+		// Record that this process really performed the token exchange.
+		writeFileSync(join(markerDir, `refresh-${label}`), credential.refresh, "utf8");
+		await new Promise((resolve) => setTimeout(resolve, 300));
+		return { type: "oauth", access: fakeAccessToken(accountId, `rotated-${label}`), refresh: `rotated-refresh-${label}`, expires: Date.now() + 3_600_000, accountId };
+	},
+});
+
+await withAuthTransaction(authPath, (entries) => ({
+	value: { ...entries, [`unrelated-${label}`]: { type: "api_key", key: `\$KEY_${label}` } },
+	result: undefined,
+}));
+
+process.stdout.write(JSON.stringify({ label, access: fresh.credential.access, refreshed: fresh.refreshed }));

package/packages/openai-account-pool/tests/integration-login-ownership.mjs

@@ -0,0 +1,299 @@
+import assert from "node:assert/strict";
+import { readFileSync } from "node:fs";
+
+import {
+	accountCredential,
+	accountIdOf,
+	commitAccount,
+	disablePool,
+	ensureFreshAccountCredential,
+	isPoolReferenceEntry,
+	loadPoolSnapshot,
+	POOL_REFERENCE_TOKEN,
+	POOL_UNRESOLVED_API_KEY,
+	removeAccount,
+} from "../src/auth.ts";
+import { poolEntryKey } from "../src/store.ts";
+import { fakeCredential, makeAgentDir, seedAuthFile } from "./helpers/fixtures.mjs";
+import { NativeCredentialStore } from "./helpers/native-credential-store.mjs";
+import { piAiEntry } from "./helpers/pi-runtime.mjs";
+
+/**
+ * Blocking integration check for the account pool.
+ *
+ * It runs against the real pi-ai `Models` orchestration to answer two questions before any
+ * further implementation: does a second plain login to the same provider really overwrite the
+ * first one, and can the pool keep several accounts while pi resolves request auth from a
+ * non-secret reference and the provider substitutes the selected account?
+ */
+
+const { createModels, createProvider, createAssistantMessageEventStream, lazyStream } = await import(piAiEntry());
+
+const UNRELATED_ENTRY = { type: "api_key", key: "unrelated-provider-key" };
+
+const model = {
+	id: "gpt-5-codex",
+	name: "GPT-5 Codex",
+	api: "openai-codex-responses",
+	provider: "openai-codex",
+	baseUrl: "https://chatgpt.com/backend-api",
+	input: ["text"],
+	cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+	contextWindow: 400_000,
+	maxTokens: 128_000,
+	reasoning: true,
+};
+
+const request = { systemPrompt: "test", messages: [{ role: "user", content: "hi", timestamp: Date.now() }] };
+
+function recordingApi(seen) {
+	const run = (requestModel, _context, options) => {
+		seen.push({ apiKey: options?.apiKey });
+		const stream = createAssistantMessageEventStream();
+		const message = {
+			role: "assistant",
+			content: [{ type: "text", text: "ok" }],
+			api: requestModel.api,
+			provider: requestModel.provider,
+			model: requestModel.id,
+			usage: { input: 1, output: 1, cacheRead: 0, cacheWrite: 0, totalTokens: 2, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } },
+			stopReason: "stop",
+			timestamp: Date.now(),
+		};
+		stream.push({ type: "start", partial: message });
+		stream.push({ type: "done", reason: "stop", message });
+		stream.end(message);
+		return stream;
+	};
+	return { stream: run, streamSimple: run };
+}
+
+/** Codex-shaped provider: subscription OAuth, one chat model, a recording API implementation. */
+function codexShapedProvider({ login, seen, poolAware = false, refreshed = [] }) {
+	return createProvider({
+		id: "openai-codex",
+		name: "OpenAI Codex (integration check)",
+		baseUrl: model.baseUrl,
+		auth: {
+			oauth: {
+				name: "OpenAI (ChatGPT Plus/Pro)",
+				isSubscription: true,
+				login,
+				async refresh(credential) {
+					refreshed.push(credential.refresh);
+					return credential;
+				},
+				async toAuth(credential) {
+					// Pre-request resolution of a pool reference must not touch any account.
+					if (poolAware && isPoolReferenceEntry(credential)) return { apiKey: POOL_UNRESOLVED_API_KEY };
+					return { apiKey: credential.access };
+				},
+			},
+		},
+		models: [model],
+		api: recordingApi(seen),
+	});
+}
+
+/**
+ * Minimal stand-in for the streaming wrapper: it only proves that the provider can replace the
+ * unresolved placeholder with a selected account before the API implementation runs. Account
+ * selection and failover are implemented and tested separately.
+ */
+function withAccountSubstitution(native, authPath) {
+	const substitute = async (options) => {
+		if (options?.apiKey !== POOL_UNRESOLVED_API_KEY) return options;
+		const snapshot = await loadPoolSnapshot(authPath);
+		const [account] = snapshot.accounts;
+		const { credential } = await ensureFreshAccountCredential({
+			authPath,
+			account,
+			refresh: async () => { throw new Error("The integration check uses credentials that are still valid."); },
+		});
+		return { ...options, apiKey: credential.access };
+	};
+	return {
+		...native,
+		stream: (requestModel, context, options) => lazyStream(requestModel, async () => native.stream(requestModel, context, await substitute(options))),
+		streamSimple: (requestModel, context, options) => lazyStream(requestModel, async () => native.streamSimple(requestModel, context, await substitute(options))),
+	};
+}
+
+const interaction = { prompt: async () => { throw new Error("No prompt expected."); }, notify: () => {} };
+const readEntries = (authPath) => JSON.parse(readFileSync(authPath, "utf8"));
+/** Stored accounts also carry non-secret membership data, which is not part of the credential. */
+const storedCredential = (entries, accountId) => {
+	const entry = entries[poolEntryKey(accountId)];
+	return entry === undefined ? undefined : accountCredential(entry);
+};
+
+// ---------------------------------------------------------------------------
+// 1. A second plain login to the same provider overwrites the first account.
+// ---------------------------------------------------------------------------
+{
+	const agent = makeAgentDir();
+	try {
+		const credentialA = fakeCredential("account-aaaa");
+		const credentialB = fakeCredential("account-bbbb");
+		seedAuthFile(agent.authPath, { anthropic: UNRELATED_ENTRY, "openai-codex": credentialA });
+
+		const seen = [];
+		const models = createModels({ credentials: new NativeCredentialStore(agent.authPath) });
+		models.setProvider(codexShapedProvider({ login: async () => credentialB, seen }));
+
+		await models.login("openai-codex", "oauth", interaction);
+
+		const entries = readEntries(agent.authPath);
+		assert.deepEqual(entries["openai-codex"], credentialB, "the plain login must store the new credential");
+		assert.equal(
+			JSON.stringify(entries).includes(credentialA.refresh),
+			false,
+			"reproduced: pi keeps one credential per provider, so the first account is gone",
+		);
+		assert.deepEqual(entries.anthropic, UNRELATED_ENTRY, "other providers must survive a login");
+	} finally {
+		agent.dispose();
+	}
+}
+
+// ---------------------------------------------------------------------------
+// 2. The dedicated commit path keeps every account and publishes a reference.
+// ---------------------------------------------------------------------------
+{
+	const agent = makeAgentDir();
+	try {
+		const credentialA = fakeCredential("account-aaaa");
+		const credentialB = fakeCredential("account-bbbb");
+		const credentialC = fakeCredential("account-cccc");
+		seedAuthFile(agent.authPath, { anthropic: UNRELATED_ENTRY, "openai-codex": credentialA });
+
+		const first = await commitAccount({ authPath: agent.authPath, credential: credentialB });
+		assert.equal(first.outcome, "added");
+		assert.equal(first.migratedExistingLogin, true, "the existing plain login is kept as an account");
+		assert.equal(first.accountCount, 2);
+
+		const second = await commitAccount({ authPath: agent.authPath, credential: credentialC });
+		assert.equal(second.accountCount, 3);
+
+		let entries = readEntries(agent.authPath);
+		assert.deepEqual(storedCredential(entries, "account-aaaa"), credentialA, "account A must stay unchanged");
+		assert.deepEqual(storedCredential(entries, "account-bbbb"), credentialB);
+		assert.deepEqual(storedCredential(entries, "account-cccc"), credentialC);
+		assert.deepEqual(entries.anthropic, UNRELATED_ENTRY);
+
+		// The main entry is a non-secret reference, and no account token leaks into it.
+		const main = entries["openai-codex"];
+		assert.equal(isPoolReferenceEntry(main), true);
+		assert.equal(main.access, POOL_REFERENCE_TOKEN);
+		assert.equal(main.refresh, POOL_REFERENCE_TOKEN);
+		assert.equal(accountIdOf(main), undefined, "the reference carries no account claim, so the API cannot use it");
+		for (const credential of [credentialA, credentialB, credentialC]) {
+			assert.equal(JSON.stringify(main).includes(credential.access), false);
+			assert.equal(JSON.stringify(main).includes(credential.refresh), false);
+		}
+
+		// Re-authorizing the same account updates it instead of adding a quota source.
+		const rotatedB = fakeCredential("account-bbbb", { marker: "v2" });
+		const again = await commitAccount({ authPath: agent.authPath, credential: rotatedB });
+		assert.equal(again.outcome, "updated");
+		assert.equal(again.accountCount, 3);
+		assert.equal(again.account.ordinal, 2, "the account keeps its display number");
+
+		entries = readEntries(agent.authPath);
+		assert.deepEqual(storedCredential(entries, "account-bbbb"), rotatedB);
+		assert.deepEqual(storedCredential(entries, "account-aaaa"), credentialA);
+		assert.deepEqual(storedCredential(entries, "account-cccc"), credentialC);
+	} finally {
+		agent.dispose();
+	}
+}
+
+// ---------------------------------------------------------------------------
+// 3. pi resolves request auth from the reference; the provider substitutes an account.
+// ---------------------------------------------------------------------------
+{
+	const agent = makeAgentDir();
+	try {
+		const credentialA = fakeCredential("account-aaaa");
+		seedAuthFile(agent.authPath, { "openai-codex": credentialA });
+		await commitAccount({ authPath: agent.authPath, credential: fakeCredential("account-bbbb") });
+
+		const seen = [];
+		const refreshed = [];
+		const models = createModels({ credentials: new NativeCredentialStore(agent.authPath) });
+		const native = codexShapedProvider({ login: async () => { throw new Error("No login expected."); }, seen, poolAware: true, refreshed });
+		models.setProvider(withAccountSubstitution(native, agent.authPath));
+
+		const resolved = await models.getAuth(model);
+		assert.equal(resolved?.auth.apiKey, POOL_UNRESOLVED_API_KEY, "pre-request resolution returns the placeholder");
+		assert.deepEqual(refreshed, [], "resolving the reference must not refresh any account");
+
+		const message = await models.streamSimple(model, request).result();
+		assert.equal(message.stopReason, "stop");
+		assert.equal(seen.length, 1);
+		assert.equal(seen[0].apiKey, credentialA.access, "the selected account's token reaches the API implementation");
+		assert.notEqual(seen[0].apiKey, POOL_UNRESOLVED_API_KEY, "the placeholder must never reach the API implementation");
+	} finally {
+		agent.dispose();
+	}
+}
+
+// ---------------------------------------------------------------------------
+// 4. Off restores one native login, and plain login semantics come back with it.
+// ---------------------------------------------------------------------------
+{
+	const agent = makeAgentDir();
+	try {
+		const credentialA = fakeCredential("account-aaaa");
+		const credentialB = fakeCredential("account-bbbb");
+		seedAuthFile(agent.authPath, { anthropic: UNRELATED_ENTRY, "openai-codex": credentialA });
+		await commitAccount({ authPath: agent.authPath, credential: credentialB });
+
+		const off = await disablePool({ authPath: agent.authPath, accountId: "account-aaaa" });
+		assert.equal(off.dormantAccounts, 1);
+
+		let entries = readEntries(agent.authPath);
+		assert.deepEqual(entries["openai-codex"], credentialA, "off restores the chosen account as the plain login, without pool bookkeeping");
+		assert.equal(entries[poolEntryKey("account-aaaa")], undefined, "the restored account has exactly one refresh authority");
+		assert.deepEqual(storedCredential(entries, "account-bbbb"), credentialB, "the other account stays stored but dormant");
+		assert.deepEqual(entries.anthropic, UNRELATED_ENTRY);
+
+		// With the pool off, a plain login overwrites the main entry again, as pi normally does.
+		const replacement = fakeCredential("account-dddd");
+		const models = createModels({ credentials: new NativeCredentialStore(agent.authPath) });
+		models.setProvider(codexShapedProvider({ login: async () => replacement, seen: [] }));
+		await models.login("openai-codex", "oauth", interaction);
+
+		entries = readEntries(agent.authPath);
+		assert.deepEqual(entries["openai-codex"], replacement);
+		assert.deepEqual(storedCredential(entries, "account-bbbb"), credentialB, "a plain login never touches isolated account entries");
+	} finally {
+		agent.dispose();
+	}
+}
+
+// ---------------------------------------------------------------------------
+// 5. Removing the last account logs openai-codex out instead of leaving a dangling reference.
+// ---------------------------------------------------------------------------
+{
+	const agent = makeAgentDir();
+	try {
+		seedAuthFile(agent.authPath, { anthropic: UNRELATED_ENTRY });
+		await commitAccount({ authPath: agent.authPath, credential: fakeCredential("account-aaaa") });
+		const snapshot = await loadPoolSnapshot(agent.authPath);
+		assert.equal(snapshot.enabled, true);
+		assert.equal(snapshot.accounts.length, 1, "one account enables the pool but leaves no spare");
+
+		const removal = await removeAccount({ authPath: agent.authPath, accountId: "account-aaaa" });
+		assert.equal(removal.poolDisabled, true);
+
+		const entries = readEntries(agent.authPath);
+		assert.equal(entries["openai-codex"], undefined, "no reference may survive without accounts");
+		assert.equal(entries[poolEntryKey("account-aaaa")], undefined);
+		assert.deepEqual(entries.anthropic, UNRELATED_ENTRY);
+	} finally {
+		agent.dispose();
+	}
+}
+
+console.log("Login ownership integration checks passed");

package/packages/openai-account-pool/tests/login-slot.mjs

@@ -0,0 +1,131 @@
+import assert from "node:assert/strict";
+
+import { LoginSlot } from "../src/login-slot.ts";
+
+/**
+ * The bug this guards: pi gives extensions one slot below the chat, and closing a selector puts
+ * the editor back there. A dialog mounted before the selector is therefore gone afterwards, and
+ * the Codex flow asks for the login method first, so the sign-in URL and the manual-code input
+ * would render into a component nobody can see.
+ */
+
+/** A slot that records every dialog call together with whether that dialog was still mounted. */
+function fakeUi(answers) {
+	const dialogs = [];
+	const selects = [];
+	const calls = [];
+	const record = (dialog, name, detail) => calls.push({ name, detail, visible: dialog.live });
+	return {
+		dialogs,
+		selects,
+		calls,
+		ui: {
+			mount: (register) =>
+				new Promise((resolve) => {
+					const state = { live: true };
+					dialogs.push(state);
+					register(
+						{
+							signal: new AbortController().signal,
+							showAuth: (url) => record(state, "auth", url),
+							showDeviceCode: (info) => record(state, "device_code", info.userCode),
+							showInfo: (message) => record(state, "info", message),
+							showProgress: (message) => record(state, "progress", message),
+							showWaiting: (message) => record(state, "waiting", message),
+							showManualInput: async (message) => {
+								record(state, "manual_code", message);
+								return answers.shift();
+							},
+							showPrompt: async (message) => {
+								record(state, "prompt", message);
+								return answers.shift();
+							},
+						},
+						() => {
+							state.live = false;
+							resolve(undefined);
+						},
+					);
+				}),
+			select: async (message, options) => {
+				selects.push({ message, visible: dialogs.filter((dialog) => dialog.live).length });
+				return options[answers.shift()];
+			},
+		},
+	};
+}
+
+const LOGIN_METHOD = {
+	type: "select",
+	message: "选择登录方式",
+	options: [
+		{ id: "browser", label: "浏览器登录" },
+		{ id: "device", label: "设备码登录" },
+	],
+};
+
+// ---------------------------------------------------------------------------
+// The Codex order: pick a method, then show the URL, then take the pasted code.
+// ---------------------------------------------------------------------------
+{
+	const { ui, dialogs, selects, calls } = fakeUi([0, "pasted-code"]);
+	const slot = new LoginSlot(ui);
+	const interaction = slot.interaction;
+
+	assert.equal(await interaction.prompt(LOGIN_METHOD), "browser");
+	assert.deepEqual(selects, [{ message: "选择登录方式", visible: 0 }], "no dialog is mounted while a selector owns the slot");
+
+	interaction.notify({ type: "auth_url", url: "https://login.invalid/authorize", instructions: "在浏览器中完成" });
+	assert.equal(await interaction.prompt({ type: "manual_code", message: "粘贴回调代码" }), "pasted-code");
+
+	assert.deepEqual(
+		calls.map((call) => call.name),
+		["auth", "manual_code"],
+		"the steps after the selector reach the dialog in order",
+	);
+	for (const call of calls) assert.ok(call.visible, `${call.name} rendered into a dialog the user can see`);
+	assert.equal(dialogs.length, 1, "one dialog is enough once the selector is out of the way");
+
+	await slot.close();
+	assert.deepEqual(dialogs.map((dialog) => dialog.live), [false], "closing the login gives the slot back");
+}
+
+// ---------------------------------------------------------------------------
+// A selector in the middle of a login re-mounts the dialog instead of losing it.
+// ---------------------------------------------------------------------------
+{
+	const { ui, dialogs, selects, calls } = fakeUi([1, "device-done"]);
+	const slot = new LoginSlot(ui);
+	const interaction = slot.interaction;
+
+	interaction.notify({ type: "progress", message: "正在准备授权…" });
+	assert.equal(await interaction.prompt(LOGIN_METHOD), "device");
+	assert.deepEqual(selects, [{ message: "选择登录方式", visible: 0 }], "the dialog steps aside for the selector");
+
+	interaction.notify({ type: "device_code", userCode: "ABCD-EFGH", verificationUri: "https://login.invalid/device" });
+	assert.equal(await interaction.prompt({ type: "manual_code", message: "粘贴回调代码" }), "device-done");
+
+	assert.deepEqual(
+		calls.map((call) => call.name),
+		["progress", "device_code", "waiting", "manual_code"],
+		"events keep the order the provider emitted them in",
+	);
+	for (const call of calls) assert.ok(call.visible, `${call.name} rendered into a dialog the user can see`);
+	assert.equal(dialogs.length, 2, "the dialog is mounted again after the selector took the slot");
+
+	await slot.close();
+	assert.deepEqual(dialogs.map((dialog) => dialog.live), [false, false]);
+}
+
+// ---------------------------------------------------------------------------
+// Backing out of the selector cancels the login rather than hanging on the slot.
+// ---------------------------------------------------------------------------
+{
+	const { ui, dialogs } = fakeUi([undefined]);
+	const slot = new LoginSlot(ui);
+	await assert.rejects(() => slot.interaction.prompt(LOGIN_METHOD), /Login cancelled/);
+	await slot.close();
+	assert.deepEqual(dialogs, [], "a cancelled method choice never opens a dialog");
+}
+
+console.log("Login slot tests passed");

package/packages/openai-account-pool/tests/multiprocess.mjs

@@ -0,0 +1,73 @@
+import assert from "node:assert/strict";
+import { execFile } from "node:child_process";
+import { mkdirSync, readdirSync, readFileSync } from "node:fs";
+import { join } from "node:path";
+import { promisify } from "node:util";
+
+import { commitAccount } from "../src/auth.ts";
+import { poolEntryKey, withAuthTransaction } from "../src/store.ts";
+import { fakeCredential, makeAgentDir } from "./helpers/fixtures.mjs";
+
+const run = promisify(execFile);
+const CHILD = new URL("./helpers/refresh-child.mjs", import.meta.url).pathname;
+const ACCOUNT = "account-aaaa";
+
+// ---------------------------------------------------------------------------
+// Two real pi processes racing for a rotating refresh token exchange it once.
+// ---------------------------------------------------------------------------
+{
+	const agent = makeAgentDir();
+	try {
+		const markerDir = join(agent.directory, "markers");
+		mkdirSync(markerDir);
+		await commitAccount({ authPath: agent.authPath, credential: fakeCredential(ACCOUNT, { expiresInMs: -1_000 }) });
+
+		const [first, second] = await Promise.all([
+			run(process.execPath, ["--experimental-strip-types", CHILD, agent.authPath, ACCOUNT, markerDir, "one"]),
+			run(process.execPath, ["--experimental-strip-types", CHILD, agent.authPath, ACCOUNT, markerDir, "two"]),
+		]);
+
+		const results = [JSON.parse(first.stdout), JSON.parse(second.stdout)];
+		assert.deepEqual(readdirSync(markerDir).length, 1, "a rotating refresh token must be exchanged exactly once");
+		assert.equal(results[0].access, results[1].access, "both processes end up with the same live credential");
+		assert.equal(results.filter((result) => result.refreshed).length, 1, "only the winner reports a refresh");
+
+		const entries = JSON.parse(readFileSync(agent.authPath, "utf8"));
+		assert.equal(entries[poolEntryKey(ACCOUNT)].access, results[0].access, "the stored credential is the one both processes use");
+		assert.deepEqual(entries["unrelated-one"], { type: "api_key", key: "$KEY_one" });
+		assert.deepEqual(entries["unrelated-two"], { type: "api_key", key: "$KEY_two" }, "concurrent writes from another process are not lost");
+		assert.equal(entries[poolEntryKey(ACCOUNT)].piOpenAIAccountPoolMembership.ordinal, 1, "membership survives rotation");
+	} finally {
+		agent.dispose();
+	}
+}
+
+// ---------------------------------------------------------------------------
+// A write from this process interleaved with four child writes keeps every entry.
+// ---------------------------------------------------------------------------
+{
+	const agent = makeAgentDir();
+	try {
+		const markerDir = join(agent.directory, "markers");
+		mkdirSync(markerDir);
+		await commitAccount({ authPath: agent.authPath, credential: fakeCredential(ACCOUNT) });
+
+		const children = ["a", "b", "c", "d"].map((label) =>
+			run(process.execPath, ["--experimental-strip-types", CHILD, agent.authPath, ACCOUNT, markerDir, label]),
+		);
+		const local = withAuthTransaction(agent.authPath, (entries) => ({
+			value: { ...entries, anthropic: { type: "api_key", key: "$ANTHROPIC_API_KEY" } },
+			result: undefined,
+		}));
+		await Promise.all([...children, local]);
+
+		const entries = JSON.parse(readFileSync(agent.authPath, "utf8"));
+		for (const label of ["a", "b", "c", "d"]) assert.deepEqual(entries[`unrelated-${label}`], { type: "api_key", key: `$KEY_${label}` });
+		assert.deepEqual(entries.anthropic, { type: "api_key", key: "$ANTHROPIC_API_KEY" });
+		assert.equal(readdirSync(markerDir).length, 0, "a valid credential is never refreshed");
+	} finally {
+		agent.dispose();
+	}
+}
+
+console.log("Multiprocess tests passed");