fix(grok): coalesce in-process refreshes and unexport GrokRefreshAuthError

Address review on #1634:

- keep GrokRefreshAuthError module-private, consistent with the updated
  Kimi implementation, avoiding new public API surface;
- dedupe concurrent resolveGrokAuth calls with a per-credential in-flight
  refresh map keyed by sourceFile + authRecordKey, mirroring
  kimiRefreshInFlight, so parallel requests share one refresh (and its
  peer-rotation adoption outcome) instead of racing the same stale
  refresh token;
- add a test proving concurrent refreshes coalesce into a single token
  request.
This commit is contained in:
Flavio martil
2026-08-06 01:14:09 -03:00
parent bb018e95cf
commit 01d275892e
2 changed files with 58 additions and 5 deletions
@@ -47,6 +47,8 @@ const grokDefaultOidcIssuer = "https://auth.x.ai";
const grokOauthDefaultTimeoutMs = 8_000;
const grokFallbackClientVersion = "0.2.93";
const grokRefreshInFlight = new Map<string, Promise<GrokTokenSet>>();
const grokBillingResetPaths = [
"$.billingPeriodEnd",
"$.currentPeriod.end",
@@ -182,7 +184,7 @@ const grokBillingMapping: ProviderAccountMappingConfig = {
]
};
export class GrokRefreshAuthError extends Error {
class GrokRefreshAuthError extends Error {
readonly status: number;
constructor(status: number, message: string) {
@@ -265,11 +267,25 @@ export async function resolveGrokAuth(): Promise<GrokTokenSet | undefined> {
if (!auth?.refreshToken || (auth.accessToken && !grokAccessTokenExpired(auth))) {
return auth;
}
try {
return await refreshGrokAuth(auth);
} catch (error) {
return adoptPeerRotatedGrokAuth(auth, error);
// Coalesce concurrent refreshes of the same credential: the refresh token
// rotates on every refresh, so every in-process caller must share a single
// refresh attempt (and its peer-rotation adoption outcome) instead of
// submitting the same stale token in parallel.
const key = `${auth.sourceFile}::${auth.authRecordKey ?? ""}`;
let refresh = grokRefreshInFlight.get(key);
if (!refresh) {
refresh = (async () => {
try {
return await refreshGrokAuth(auth);
} catch (error) {
return adoptPeerRotatedGrokAuth(auth, error);
}
})().finally(() => {
grokRefreshInFlight.delete(key);
});
grokRefreshInFlight.set(key, refresh);
}
return refresh;
}
function adoptPeerRotatedGrokAuth(auth: GrokTokenSet, error: unknown): GrokTokenSet {
@@ -403,6 +403,43 @@ test("Grok local provider does not adopt a different account record on refresh 4
});
});
test("Grok local provider coalesces concurrent refreshes of the same credential", async (t) => {
await withGrokHome(async (grokHome) => {
writeGrokAuth(grokHome, {
key: "expired-token",
refresh_token: "stale-refresh-token",
expires_at: "2000-01-01T00:00:00Z",
oidc_client_id: "grok-client-id",
oidc_issuer: "https://auth.x.ai"
});
writeGrokModels(grokHome);
const previousFetch = globalThis.fetch;
process.env.GROK_OIDC_TOKEN_ENDPOINT = "http://127.0.0.1/grok/oauth/token";
let refreshCalls = 0;
globalThis.fetch = async () => {
refreshCalls += 1;
// Keep the refresh pending so the concurrent callers must share it.
await new Promise((resolve) => setTimeout(resolve, 25));
return new Response(JSON.stringify({
access_token: "refreshed-grok-access-token",
expires_in: 3600,
refresh_token: "refreshed-grok-refresh-token"
}), { headers: { "content-type": "application/json" }, status: 200 });
};
t.after(() => {
globalThis.fetch = previousFetch;
});
const results = await Promise.all([resolveGrokAuth(), resolveGrokAuth(), resolveGrokAuth()]);
assert.equal(refreshCalls, 1);
for (const auth of results) {
assert.equal(auth.accessToken, "refreshed-grok-access-token");
assert.equal(auth.refreshToken, "refreshed-grok-refresh-token");
}
});
});
async function withGrokHome(run) {
const previousGrokHome = process.env.GROK_HOME;
const previousGrokAuthFile = process.env.GROK_AUTH_FILE;