fix: prevent silent logout when multiple pi instances refresh OAuth tokens (#466)

This commit is contained in:
cursive 2026-01-05 23:55:05 +08:00 committed by GitHub
parent 80be0fc901
commit 6a8609ac56
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 162 additions and 16 deletions

38
package-lock.json generated
View file

@ -2898,6 +2898,16 @@
"undici-types": "~6.21.0" "undici-types": "~6.21.0"
} }
}, },
"node_modules/@types/proper-lockfile": {
"version": "4.1.4",
"resolved": "https://registry.npmjs.org/@types/proper-lockfile/-/proper-lockfile-4.1.4.tgz",
"integrity": "sha512-uo2ABllncSqg9F1D4nugVl9v93RmjxF6LJzQLMLDdPaXCUIDPeOJ21Gbqi43xNKzBi/WQ0Q0dICqufzQbMjipQ==",
"dev": true,
"license": "MIT",
"dependencies": {
"@types/retry": "*"
}
},
"node_modules/@types/retry": { "node_modules/@types/retry": {
"version": "0.12.0", "version": "0.12.0",
"resolved": "https://registry.npmjs.org/@types/retry/-/retry-0.12.0.tgz", "resolved": "https://registry.npmjs.org/@types/retry/-/retry-0.12.0.tgz",
@ -5756,6 +5766,32 @@
"integrity": "sha512-3ouUOpQhtgrbOa17J7+uxOTpITYWaGP7/AhoR3+A+/1e9skrzelGi/dXzEYyvbxubEF6Wn2ypscTKiKJFFn1ag==", "integrity": "sha512-3ouUOpQhtgrbOa17J7+uxOTpITYWaGP7/AhoR3+A+/1e9skrzelGi/dXzEYyvbxubEF6Wn2ypscTKiKJFFn1ag==",
"license": "MIT" "license": "MIT"
}, },
"node_modules/proper-lockfile": {
"version": "4.1.2",
"resolved": "https://registry.npmjs.org/proper-lockfile/-/proper-lockfile-4.1.2.tgz",
"integrity": "sha512-TjNPblN4BwAWMXU8s9AEz4JmQxnD1NNL7bNOY/AKUzyamc379FWASUhc/K1pL2noVb+XmZKLL68cjzLsiOAMaA==",
"license": "MIT",
"dependencies": {
"graceful-fs": "^4.2.4",
"retry": "^0.12.0",
"signal-exit": "^3.0.2"
}
},
"node_modules/proper-lockfile/node_modules/retry": {
"version": "0.12.0",
"resolved": "https://registry.npmjs.org/retry/-/retry-0.12.0.tgz",
"integrity": "sha512-9LkiTwjUh6rT555DtE9rTX+BKByPfrMzEAtnlEtdEwr3Nkffwiihqe2bWADg+OQRjt9gl6ICdmB/ZFDCGAtSow==",
"license": "MIT",
"engines": {
"node": ">= 4"
}
},
"node_modules/proper-lockfile/node_modules/signal-exit": {
"version": "3.0.7",
"resolved": "https://registry.npmjs.org/signal-exit/-/signal-exit-3.0.7.tgz",
"integrity": "sha512-wnD2ZE+l+SPC/uoS0vXeE9L1+0wuaMqKlfz9AMUo38JsyLSBWSFcHR1Rri62LZc12vLr1gb3jl7iwQhgwpAbGQ==",
"license": "ISC"
},
"node_modules/proxy-from-env": { "node_modules/proxy-from-env": {
"version": "1.1.0", "version": "1.1.0",
"resolved": "https://registry.npmjs.org/proxy-from-env/-/proxy-from-env-1.1.0.tgz", "resolved": "https://registry.npmjs.org/proxy-from-env/-/proxy-from-env-1.1.0.tgz",
@ -7176,6 +7212,7 @@
"glob": "^11.0.3", "glob": "^11.0.3",
"jiti": "^2.6.1", "jiti": "^2.6.1",
"marked": "^15.0.12", "marked": "^15.0.12",
"proper-lockfile": "^4.1.2",
"sharp": "^0.34.2" "sharp": "^0.34.2"
}, },
"bin": { "bin": {
@ -7184,6 +7221,7 @@
"devDependencies": { "devDependencies": {
"@types/diff": "^7.0.2", "@types/diff": "^7.0.2",
"@types/node": "^24.3.0", "@types/node": "^24.3.0",
"@types/proper-lockfile": "^4.1.4",
"typescript": "^5.7.3", "typescript": "^5.7.3",
"vitest": "^3.2.4" "vitest": "^3.2.4"
}, },

View file

@ -49,11 +49,13 @@
"glob": "^11.0.3", "glob": "^11.0.3",
"jiti": "^2.6.1", "jiti": "^2.6.1",
"marked": "^15.0.12", "marked": "^15.0.12",
"proper-lockfile": "^4.1.2",
"sharp": "^0.34.2" "sharp": "^0.34.2"
}, },
"devDependencies": { "devDependencies": {
"@types/diff": "^7.0.2", "@types/diff": "^7.0.2",
"@types/node": "^24.3.0", "@types/node": "^24.3.0",
"@types/proper-lockfile": "^4.1.4",
"typescript": "^5.7.3", "typescript": "^5.7.3",
"vitest": "^3.2.4" "vitest": "^3.2.4"
}, },

View file

@ -1,6 +1,9 @@
/** /**
* Credential storage for API keys and OAuth tokens. * Credential storage for API keys and OAuth tokens.
* Handles loading, saving, and refreshing credentials from auth.json. * Handles loading, saving, and refreshing credentials from auth.json.
*
* Uses file locking to prevent race conditions when multiple pi instances
* try to refresh tokens simultaneously.
*/ */
import { import {
@ -16,6 +19,7 @@ import {
} from "@mariozechner/pi-ai"; } from "@mariozechner/pi-ai";
import { chmodSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "fs"; import { chmodSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "fs";
import { dirname } from "path"; import { dirname } from "path";
import lockfile from "proper-lockfile";
export type ApiKeyCredential = { export type ApiKeyCredential = {
type: "api_key"; type: "api_key";
@ -198,12 +202,93 @@ export class AuthStorage {
this.remove(provider); this.remove(provider);
} }
/**
* Refresh OAuth token with file locking to prevent race conditions.
* Multiple pi instances may try to refresh simultaneously when tokens expire.
* This ensures only one instance refreshes while others wait and use the result.
*/
private async refreshOAuthTokenWithLock(
provider: OAuthProvider,
): Promise<{ apiKey: string; newCredentials: OAuthCredentials } | null> {
// Ensure auth file exists for locking
if (!existsSync(this.authPath)) {
const dir = dirname(this.authPath);
if (!existsSync(dir)) {
mkdirSync(dir, { recursive: true, mode: 0o700 });
}
writeFileSync(this.authPath, "{}", "utf-8");
chmodSync(this.authPath, 0o600);
}
let release: (() => Promise<void>) | undefined;
try {
// Acquire exclusive lock with retry and timeout
// Use generous retry window to handle slow token endpoints
release = await lockfile.lock(this.authPath, {
retries: {
retries: 10,
factor: 2,
minTimeout: 100,
maxTimeout: 10000,
randomize: true,
},
stale: 30000, // Consider lock stale after 30 seconds
});
// Re-read file after acquiring lock - another instance may have refreshed
this.reload();
const cred = this.data[provider];
if (cred?.type !== "oauth") {
return null;
}
// Check if token is still expired after re-reading
// (another instance may have already refreshed it)
if (Date.now() < cred.expires) {
// Token is now valid - another instance refreshed it
const needsProjectId = provider === "google-gemini-cli" || provider === "google-antigravity";
const apiKey = needsProjectId
? JSON.stringify({ token: cred.access, projectId: cred.projectId })
: cred.access;
return { apiKey, newCredentials: cred };
}
// Token still expired, we need to refresh
const oauthCreds: Record<string, OAuthCredentials> = {};
for (const [key, value] of Object.entries(this.data)) {
if (value.type === "oauth") {
oauthCreds[key] = value;
}
}
const result = await getOAuthApiKey(provider, oauthCreds);
if (result) {
this.data[provider] = { type: "oauth", ...result.newCredentials };
this.save();
return result;
}
return null;
} finally {
// Always release the lock
if (release) {
try {
await release();
} catch {
// Ignore unlock errors (lock may have been compromised)
}
}
}
}
/** /**
* Get API key for a provider. * Get API key for a provider.
* Priority: * Priority:
* 1. Runtime override (CLI --api-key) * 1. Runtime override (CLI --api-key)
* 2. API key from auth.json * 2. API key from auth.json
* 3. OAuth token from auth.json (auto-refreshed) * 3. OAuth token from auth.json (auto-refreshed with locking)
* 4. Environment variable * 4. Environment variable
* 5. Fallback resolver (models.json custom providers) * 5. Fallback resolver (models.json custom providers)
*/ */
@ -221,23 +306,44 @@ export class AuthStorage {
} }
if (cred?.type === "oauth") { if (cred?.type === "oauth") {
// Filter to only oauth credentials for getOAuthApiKey // Check if token needs refresh
const oauthCreds: Record<string, OAuthCredentials> = {}; const needsRefresh = Date.now() >= cred.expires;
for (const [key, value] of Object.entries(this.data)) {
if (value.type === "oauth") {
oauthCreds[key] = value;
}
}
try { if (needsRefresh) {
const result = await getOAuthApiKey(provider as OAuthProvider, oauthCreds); // Use locked refresh to prevent race conditions
if (result) { try {
this.data[provider] = { type: "oauth", ...result.newCredentials }; const result = await this.refreshOAuthTokenWithLock(provider as OAuthProvider);
this.save(); if (result) {
return result.apiKey; return result.apiKey;
}
} catch (err) {
// Refresh failed - re-read file to check if another instance succeeded
this.reload();
const updatedCred = this.data[provider];
if (updatedCred?.type === "oauth" && Date.now() < updatedCred.expires) {
// Another instance refreshed successfully, use those credentials
const needsProjectId =
provider === "google-gemini-cli" || provider === "google-antigravity";
return needsProjectId
? JSON.stringify({ token: updatedCred.access, projectId: updatedCred.projectId })
: updatedCred.access;
}
// Refresh truly failed - DO NOT remove credentials
// User can retry or re-authenticate manually
const errorMessage = err instanceof Error ? err.message : String(err);
throw new Error(
`OAuth token refresh failed for ${provider}: ${errorMessage}. ` +
`Please try again or re-authenticate with /login.`,
);
} }
} catch { } else {
this.remove(provider); // Token not expired, use current access token
const needsProjectId = provider === "google-gemini-cli" || provider === "google-antigravity";
return needsProjectId
? JSON.stringify({ token: cred.access, projectId: cred.projectId })
: cred.access;
} }
} }