mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-05 23:37:25 +00:00
fix(ai,coding-agent): persist rotated OAuth tokens when a refresh is cancelled
Cancelling a request or model refresh while an OAuth token refresh was in flight discarded the rotated credential. Providers that rotate refresh tokens (e.g. Sign in with ChatGPT) had already invalidated the old token, so the next request failed with refresh_token_invalidated. The caller's signal now cancels only the wait for the credential lock. Once a refresh starts, it completes and is persisted, bounded by the refresh timeout. Request auth and model-list refresh share this path.
This commit is contained in:
parent
200387122c
commit
bde882c747
6 changed files with 128 additions and 65 deletions
|
|
@ -2,6 +2,10 @@
|
|||
|
||||
## [Unreleased]
|
||||
|
||||
### Fixed
|
||||
|
||||
- Fixed OAuth credentials being invalidated when a request or model refresh was cancelled or superseded during a token refresh: a token refresh that has started now completes and persists the rotated refresh token
|
||||
|
||||
## [1.0.2] - 2026-10-04
|
||||
|
||||
### Added
|
||||
|
|
|
|||
|
|
@ -102,10 +102,61 @@ function overlayEnvAuthContext(base: AuthContext, env: ProviderEnv): AuthContext
|
|||
const DEFAULT_OAUTH_MINIMUM_VALIDITY_MS = 5 * 60 * 1000;
|
||||
const DEFAULT_OAUTH_REFRESH_TIMEOUT_MS = 15_000;
|
||||
|
||||
/**
|
||||
* Refresh a stored OAuth credential under the credential-store lock and persist
|
||||
* the result before the lock is released. `needsRefresh` is re-checked under
|
||||
* the lock, so concurrent callers and processes refresh only once.
|
||||
*
|
||||
* `signal` cancels only the wait for the lock. Once a refresh starts, the
|
||||
* provider may already have rotated the refresh token, so the refresh and its
|
||||
* persistence ignore `signal` and are bounded only by a timeout. Otherwise a
|
||||
* cancelled caller could discard the only valid refresh token. Callers that
|
||||
* must return promptly on cancellation race this promise with their signal.
|
||||
*
|
||||
* Resolves with the stored OAuth credential after the operation, or undefined
|
||||
* when the provider no longer has an OAuth credential.
|
||||
*/
|
||||
export async function refreshStoredOAuthCredential(
|
||||
credentials: CredentialStore,
|
||||
providerId: string,
|
||||
oauth: OAuthAuth,
|
||||
needsRefresh: (credential: OAuthCredential) => boolean,
|
||||
signal: AbortSignal,
|
||||
): Promise<OAuthCredential | undefined> {
|
||||
let post: Credential | undefined;
|
||||
const lockWait = new AbortController();
|
||||
const cancelLockWait = () => lockWait.abort(signal.reason);
|
||||
signal.addEventListener("abort", cancelLockWait, { once: true });
|
||||
if (signal.aborted) cancelLockWait();
|
||||
try {
|
||||
post = await credentials.modify(
|
||||
providerId,
|
||||
async (current) => {
|
||||
signal.removeEventListener("abort", cancelLockWait);
|
||||
signal.throwIfAborted();
|
||||
if (current?.type !== "oauth") return undefined; // logged out meanwhile
|
||||
if (!needsRefresh(current)) return undefined; // another process/request refreshed
|
||||
try {
|
||||
return await oauth.refresh(current, AbortSignal.timeout(DEFAULT_OAUTH_REFRESH_TIMEOUT_MS));
|
||||
} catch (error) {
|
||||
throw new ModelsError("oauth", `OAuth refresh failed for ${providerId}`, { cause: error });
|
||||
}
|
||||
},
|
||||
{ signal: lockWait.signal },
|
||||
);
|
||||
} catch (error) {
|
||||
if (error instanceof ModelsError) throw error;
|
||||
signal.throwIfAborted();
|
||||
throw new ModelsError("auth", `Credential store modify failed for ${providerId}`, { cause: error });
|
||||
} finally {
|
||||
signal.removeEventListener("abort", cancelLockWait);
|
||||
}
|
||||
return post?.type === "oauth" ? post : undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* OAuth resolution with double-checked locking: tokens with less than five
|
||||
* minutes remaining lock, re-check expiry under the lock, refresh once
|
||||
* globally, and persist the rotated credential before release.
|
||||
* minutes remaining are refreshed through `refreshStoredOAuthCredential`.
|
||||
*/
|
||||
async function resolveStoredOAuth(
|
||||
credentials: CredentialStore,
|
||||
|
|
@ -121,30 +172,8 @@ async function resolveStoredOAuth(
|
|||
|
||||
if (expiresSoon(credential)) {
|
||||
// Optimistic check said expired; the authoritative check runs under the lock.
|
||||
let post: Credential | undefined;
|
||||
try {
|
||||
post = await credentials.modify(
|
||||
providerId,
|
||||
async (current) => {
|
||||
if (current?.type !== "oauth") return undefined; // logged out meanwhile
|
||||
if (!expiresSoon(current)) return undefined; // another process/request refreshed
|
||||
try {
|
||||
const refreshSignal = AbortSignal.any([
|
||||
signal,
|
||||
AbortSignal.timeout(DEFAULT_OAUTH_REFRESH_TIMEOUT_MS),
|
||||
]);
|
||||
return await oauth.refresh(current, refreshSignal);
|
||||
} catch (error) {
|
||||
throw new ModelsError("oauth", `OAuth refresh failed for ${providerId}`, { cause: error });
|
||||
}
|
||||
},
|
||||
{ signal },
|
||||
);
|
||||
} catch (error) {
|
||||
if (error instanceof ModelsError) throw error;
|
||||
throw new ModelsError("auth", `Credential store modify failed for ${providerId}`, { cause: error });
|
||||
}
|
||||
if (post?.type !== "oauth") return undefined; // logged out meanwhile
|
||||
const post = await refreshStoredOAuthCredential(credentials, providerId, oauth, expiresSoon, signal);
|
||||
if (!post) return undefined; // logged out meanwhile
|
||||
credential = post;
|
||||
// The normal five-minute window triggers a refresh but does not impose a
|
||||
// provider contract. Explicit callers (such as bearer-token export) do
|
||||
|
|
|
|||
|
|
@ -1,7 +1,12 @@
|
|||
import { lazyStream } from "./api/lazy.ts";
|
||||
import { defaultProviderAuthContext as defaultAuthContext } from "./auth/context.ts";
|
||||
import { InMemoryCredentialStore } from "./auth/credential-store.ts";
|
||||
import { type AuthResolutionOverrides, ModelsError, resolveProviderAuth } from "./auth/resolve.ts";
|
||||
import {
|
||||
type AuthResolutionOverrides,
|
||||
ModelsError,
|
||||
refreshStoredOAuthCredential,
|
||||
resolveProviderAuth,
|
||||
} from "./auth/resolve.ts";
|
||||
import type {
|
||||
AuthCheck,
|
||||
AuthContext,
|
||||
|
|
@ -615,15 +620,15 @@ class ModelsImpl implements MutableModels {
|
|||
if (!oauth) return undefined;
|
||||
if (Date.now() < stored.expires) return stored;
|
||||
if (signal.aborted) return undefined;
|
||||
const post = await this.credentials.modify(
|
||||
// A refresh that has started survives cancellation or a superseding model refresh, so a
|
||||
// rotated refresh token is always persisted. A newer refresh then sees the fresh credential.
|
||||
return refreshStoredOAuthCredential(
|
||||
this.credentials,
|
||||
provider.id,
|
||||
async (current) => {
|
||||
if (current?.type !== "oauth" || Date.now() < current.expires) return undefined;
|
||||
return oauth.refresh(current, signal);
|
||||
},
|
||||
{ signal },
|
||||
oauth,
|
||||
(current) => Date.now() >= current.expires,
|
||||
signal,
|
||||
);
|
||||
return post?.type === "oauth" ? post : undefined;
|
||||
}
|
||||
|
||||
const apiKey = provider.auth.apiKey;
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import { describe, expect, it, vi } from "vitest";
|
||||
import { InMemoryCredentialStore } from "../src/auth/credential-store.ts";
|
||||
import type { ApiKeyAuth, CredentialStore, OAuthAuth, OAuthCredential, ProviderAuth } from "../src/auth/types.ts";
|
||||
import type { ApiKeyAuth, CredentialStore, OAuthAuth, ProviderAuth } from "../src/auth/types.ts";
|
||||
import { calculateCost, createModels, createProvider, hasApi, type Provider } from "../src/models.ts";
|
||||
import { InMemoryModelsStore, type ModelsStore, type ModelsStoreEntry } from "../src/models-store.ts";
|
||||
import type { Api, AssistantMessage, Context, Model, SimpleStreamOptions, StreamOptions, Usage } from "../src/types.ts";
|
||||
|
|
@ -473,6 +473,37 @@ describe("Models runtime", () => {
|
|||
expect(await credentials.read("oauth-dynamic")).toMatchObject({ access: "fresh", refresh: "rotated" });
|
||||
});
|
||||
|
||||
it("persists an OAuth refresh that started before the model refresh was cancelled", async () => {
|
||||
const credentials = new InMemoryCredentialStore();
|
||||
await credentials.modify("p1", async () => ({
|
||||
type: "oauth",
|
||||
access: "old",
|
||||
refresh: "old-refresh",
|
||||
expires: 0,
|
||||
}));
|
||||
const controller = new AbortController();
|
||||
const models = createModels({ credentials });
|
||||
models.setProvider(
|
||||
testProvider({
|
||||
id: "p1",
|
||||
auth: {
|
||||
oauth: testOAuth({
|
||||
refresh: async (credential) => {
|
||||
controller.abort();
|
||||
return { ...credential, access: "new", refresh: "new-refresh", expires: Date.now() + 60_000 };
|
||||
},
|
||||
}),
|
||||
},
|
||||
refreshModels: async () => {},
|
||||
}),
|
||||
);
|
||||
|
||||
expect((await models.refresh({ signal: controller.signal })).aborted).toBe(true);
|
||||
await vi.waitFor(async () => {
|
||||
expect(await credentials.read("p1")).toMatchObject({ refresh: "new-refresh" });
|
||||
});
|
||||
});
|
||||
|
||||
it("always gives providers a concrete signal", async () => {
|
||||
let receivedSignal: AbortSignal | undefined;
|
||||
const models = createModels();
|
||||
|
|
@ -746,46 +777,36 @@ describe("Models runtime", () => {
|
|||
expect(await credentials.read("p1")).toEqual({ type: "api_key", key: "first" });
|
||||
});
|
||||
|
||||
it("passes cancellation to OAuth refresh and preserves the previous credential", async () => {
|
||||
// Radius bug report 01a10855-43d8-7447-a710-2acf5ee0b2ee: refresh_token_invalidated after a cancelled refresh.
|
||||
it("persists an OAuth refresh that started before the request was cancelled", async () => {
|
||||
const credentials = new InMemoryCredentialStore();
|
||||
const previous: OAuthCredential = { type: "oauth", access: "old", refresh: "old-refresh", expires: 0 };
|
||||
await credentials.modify("p1", async () => previous);
|
||||
let startRefresh: (() => void) | undefined;
|
||||
let finishRefresh: ((credential: typeof previous) => void) | undefined;
|
||||
const refreshStarted = new Promise<void>((resolve) => {
|
||||
startRefresh = resolve;
|
||||
});
|
||||
const blockedRefresh = new Promise<typeof previous>((resolve) => {
|
||||
finishRefresh = resolve;
|
||||
});
|
||||
let receivedSignal: AbortSignal | undefined;
|
||||
await credentials.modify("p1", async () => ({
|
||||
type: "oauth",
|
||||
access: "old",
|
||||
refresh: "old-refresh",
|
||||
expires: 0,
|
||||
}));
|
||||
const controller = new AbortController();
|
||||
const models = createModels({ credentials });
|
||||
models.setProvider(
|
||||
testProvider({
|
||||
id: "p1",
|
||||
auth: {
|
||||
oauth: testOAuth({
|
||||
refresh: async (_credential, signal) => {
|
||||
receivedSignal = signal;
|
||||
startRefresh?.();
|
||||
return blockedRefresh;
|
||||
refresh: async (credential) => {
|
||||
// The provider has rotated old-refresh by the time the request is cancelled.
|
||||
controller.abort();
|
||||
return { ...credential, access: "new", refresh: "new-refresh", expires: Date.now() + 60_000 };
|
||||
},
|
||||
}),
|
||||
},
|
||||
}),
|
||||
);
|
||||
const controller = new AbortController();
|
||||
const auth = models.getAuth("p1", { signal: controller.signal });
|
||||
await refreshStarted;
|
||||
controller.abort();
|
||||
|
||||
await expect(auth).rejects.toMatchObject({ name: "AbortError" });
|
||||
expect(receivedSignal).toBeInstanceOf(AbortSignal);
|
||||
expect(receivedSignal?.aborted).toBe(true);
|
||||
expect(receivedSignal?.reason).toBe(controller.signal.reason);
|
||||
finishRefresh?.({ ...previous, access: "new", expires: Date.now() + 60_000 });
|
||||
await new Promise((resolve) => setTimeout(resolve, 0));
|
||||
expect(await credentials.read("p1")).toEqual(previous);
|
||||
await expect(models.getAuth("p1", { signal: controller.signal })).rejects.toMatchObject({ name: "AbortError" });
|
||||
await vi.waitFor(async () => {
|
||||
expect(await credentials.read("p1")).toMatchObject({ refresh: "new-refresh" });
|
||||
});
|
||||
});
|
||||
|
||||
it("resolves auth: stored credential owns the provider, ambient only when nothing stored", async () => {
|
||||
|
|
|
|||
|
|
@ -2,6 +2,10 @@
|
|||
|
||||
## [Unreleased]
|
||||
|
||||
### Fixed
|
||||
|
||||
- Fixed subscription logins such as Sign in with ChatGPT failing with `refresh_token_invalidated` after a request was cancelled during an OAuth token refresh
|
||||
|
||||
## [1.0.2] - 2026-10-04
|
||||
|
||||
### New Features
|
||||
|
|
|
|||
|
|
@ -263,7 +263,7 @@ describe("ModelRuntime auth options", () => {
|
|||
});
|
||||
});
|
||||
|
||||
it("forwards cancellation to extension OAuth refresh", async () => {
|
||||
it("does not cancel an extension OAuth refresh with the request", async () => {
|
||||
const credentials = AuthStorage.inMemory({
|
||||
"extension-oauth": {
|
||||
type: "oauth",
|
||||
|
|
@ -295,8 +295,8 @@ describe("ModelRuntime auth options", () => {
|
|||
expect(refreshSignal).toBeInstanceOf(AbortSignal);
|
||||
const reason = new Error("cancelled");
|
||||
controller.abort(reason);
|
||||
expect(refreshSignal?.aborted).toBe(true);
|
||||
expect(refreshSignal?.reason).toBe(reason);
|
||||
// A started refresh may already have rotated the refresh token and must be persisted.
|
||||
expect(refreshSignal?.aborted).toBe(false);
|
||||
});
|
||||
|
||||
it("does not fabricate an API key method for an extension OAuth-only provider", async () => {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue