From 96cb327b8e35f33bdb99ed5b9258487cdeab7ac8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Petr=20Hou=C5=A1ka?= Date: Wed, 4 Nov 2026 22:44:22 +0100 Subject: [PATCH] fix(coding-agent): wait for a started OAuth refresh to be saved before exit Problem: the RPC startup catalog refresh also refreshes an expired OAuth credential. When the process exited after the token request went out but before the new token was written, the provider had already spent the old refresh token and auth.json kept only it, so the next request failed with invalid_grant. RuntimeCredentials now tracks modifications whose update function has started (a modification still waiting for the auth.json lock has sent nothing). AgentSessionRuntime.dispose() waits up to 20 s for them, so RPC, print, and interactive exits keep the rotated token. closes #11692 --- packages/coding-agent/CHANGELOG.md | 1 + packages/coding-agent/docs/rpc.md | 2 +- .../src/core/agent-session-runtime.ts | 22 ++++ .../coding-agent/src/core/model-runtime.ts | 8 ++ .../src/core/runtime-credentials.ts | 34 ++++- .../test/runtime-credentials.test.ts | 77 +++++++++++- ...92-dispose-waits-for-oauth-refresh.test.ts | 117 ++++++++++++++++++ 7 files changed, 258 insertions(+), 3 deletions(-) create mode 100644 packages/coding-agent/test/suite/regressions/11692-dispose-waits-for-oauth-refresh.test.ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index b9b6ea383..19ed155a6 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -30,6 +30,7 @@ - Fixed the find tool reporting the result limit when exactly `limit` files match, and `limit: 0` in find and ls returning every result or an empty listing; limits below 1 now mean 1 ([#7121](https://github.com/earendil-works/pi/issues/7121)) - Fixed grep match lines, codemode output, compaction summaries of tool results, and MCP error and server summaries cutting an emoji or other non-BMP character in half when truncating, which left an invalid lone surrogate in the text ([#7121](https://github.com/earendil-works/pi/issues/7121)) - Fixed the read tool loading whole files into memory, which ran out of memory on large files even with `offset`/`limit`; it now streams the file and keeps only the shown lines. Custom `ReadOperations` can stream too by providing the new optional `createReadStream` ([#9654](https://github.com/earendil-works/pi/issues/9654)) +- Fixed OAuth logins such as `openai` and `anthropic` failing with `invalid_grant` after pi was stopped while it refreshed an expired token, for example an RPC client that stops pi right after startup while the background catalog refresh runs. The token endpoint had already replaced the refresh token, but pi exited before saving the new one. On exit, pi now waits up to 20 seconds for a token refresh that has started to be saved ([#11692](https://github.com/earendil-works/pi/issues/11692)) ## [1.5.0] - 2026-11-02 diff --git a/packages/coding-agent/docs/rpc.md b/packages/coding-agent/docs/rpc.md index 61d8422c0..fbbc522eb 100644 --- a/packages/coding-agent/docs/rpc.md +++ b/packages/coding-agent/docs/rpc.md @@ -90,7 +90,7 @@ Clients must also handle child-process startup failures, unexpected exits, stder ## Shutdown -Close the child's stdin to request an orderly shutdown. Pi disposes the active runtime before exiting. Clients should still handle process signals and unexpected exits. +Close the child's stdin to request an orderly shutdown. This also works on Windows, where signals cannot be handled and killing the process ends it at once. Pi disposes the active runtime before exiting. If an OAuth token refresh has started, for example from the catalog refresh at startup, pi first waits up to 20 seconds for the new token to be saved: the provider may already have replaced the old refresh token. Give pi time to exit before killing it. Clients should still handle process signals and unexpected exits. SIGTERM and SIGHUP also start an orderly shutdown. Once a shutdown has started, from either stdin EOF or a signal, a stdin EOF or a single signal does not interrupt it; a second signal exits at once. Closing stdin and sending SIGTERM together is safe. diff --git a/packages/coding-agent/src/core/agent-session-runtime.ts b/packages/coding-agent/src/core/agent-session-runtime.ts index cf06d5fce..36e017e3b 100644 --- a/packages/coding-agent/src/core/agent-session-runtime.ts +++ b/packages/coding-agent/src/core/agent-session-runtime.ts @@ -10,6 +10,7 @@ import type { SessionStartEvent, } from "./extensions/index.ts"; import { emitSessionShutdownEvent } from "./extensions/runner.ts"; +import type { ModelRuntime } from "./model-runtime.ts"; import type { CreateAgentSessionResult } from "./sdk.ts"; import { assertSessionCwdExists } from "./session-cwd.ts"; import { SessionManager } from "./session-manager.ts"; @@ -401,6 +402,11 @@ export class AgentSessionRuntime { return { cancelled: false }; } + /** + * Dispose the session for process exit. Also waits, up to a limit, for credential writes in + * progress: an OAuth refresh may already have spent the old single-use refresh token, so exiting + * before the new one is saved would lose the login (#11692). + */ async dispose(): Promise { await this.session.abort(); await emitSessionShutdownEvent(this.session.extensionRunner, { @@ -409,6 +415,22 @@ export class AgentSessionRuntime { }); this.beforeSessionInvalidate?.(); this.session.dispose(); + await settleCredentialWrites(this.services.modelRuntime); + } +} + +// Built-in OAuth refreshes time out after 15 s; the margin covers writing the result. +const CREDENTIAL_WRITE_SETTLE_TIMEOUT_MS = 20_000; + +async function settleCredentialWrites(modelRuntime: ModelRuntime): Promise { + let timer: NodeJS.Timeout | undefined; + const timeout = new Promise((resolve) => { + timer = setTimeout(resolve, CREDENTIAL_WRITE_SETTLE_TIMEOUT_MS); + }); + try { + await Promise.race([modelRuntime.settleCredentialWrites(), timeout]); + } finally { + clearTimeout(timer); } } diff --git a/packages/coding-agent/src/core/model-runtime.ts b/packages/coding-agent/src/core/model-runtime.ts index 6a5520d68..ff5f42b39 100644 --- a/packages/coding-agent/src/core/model-runtime.ts +++ b/packages/coding-agent/src/core/model-runtime.ts @@ -640,6 +640,14 @@ export class ModelRuntime implements Models { }); } + /** + * Wait until credential writes already in progress finish, such as an OAuth refresh started by a + * background catalog refresh. Call before the process exits so a rotated token is not lost. + */ + settleCredentialWrites(): Promise { + return this.credentials.settled(); + } + listCredentials(options?: AuthOperationOptions): Promise { return this.credentials.list(options); } diff --git a/packages/coding-agent/src/core/runtime-credentials.ts b/packages/coding-agent/src/core/runtime-credentials.ts index 65f238c83..f0779bdcc 100644 --- a/packages/coding-agent/src/core/runtime-credentials.ts +++ b/packages/coding-agent/src/core/runtime-credentials.ts @@ -4,6 +4,8 @@ import type { AuthOperationOptions, Credential, CredentialInfo, CredentialStore export class RuntimeCredentials implements CredentialStore { private readonly store: CredentialStore; private readonly overrides = new Map(); + /** Modifications whose update function has started, for example an OAuth token request. */ + private readonly startedModifications = new Set>(); constructor(store: CredentialStore) { this.store = store; @@ -41,7 +43,37 @@ export class RuntimeCredentials implements CredentialStore { fn: (current: Credential | undefined) => Promise, options?: AuthOperationOptions, ): Promise { - return this.store.modify(providerId, fn, options); + let markDone!: () => void; + const done = new Promise((resolve) => { + markDone = resolve; + }); + const modification = this.store.modify( + providerId, + (current) => { + this.startedModifications.add(done); + return fn(current); + }, + options, + ); + const finish = () => { + this.startedModifications.delete(done); + markDone(); + }; + modification.then(finish, finish); + return modification; + } + + /** + * Wait until every modification whose update function has started is written or failed. + * + * A modification still waiting for the store lock is not awaited: it has not sent anything yet. + * One that has started may have spent a single-use OAuth refresh token, so exiting before its + * result is written would leave only the spent token in auth.json. + */ + async settled(): Promise { + while (this.startedModifications.size > 0) { + await Promise.all(this.startedModifications); + } } async delete(providerId: string, options?: AuthOperationOptions): Promise { diff --git a/packages/coding-agent/test/runtime-credentials.test.ts b/packages/coding-agent/test/runtime-credentials.test.ts index fe50503cb..6c9b0a3fc 100644 --- a/packages/coding-agent/test/runtime-credentials.test.ts +++ b/packages/coding-agent/test/runtime-credentials.test.ts @@ -1,4 +1,4 @@ -import type { CredentialStore } from "@earendil-works/pi-ai"; +import type { Credential, CredentialStore } from "@earendil-works/pi-ai"; import { describe, expect, test, vi } from "vitest"; import { AuthStorage } from "../src/core/auth-storage.ts"; import { RuntimeCredentials } from "../src/core/runtime-credentials.ts"; @@ -73,6 +73,81 @@ describe("RuntimeCredentials", () => { expect(await credentials.read("anthropic")).toEqual({ type: "api_key", key: "runtime-key" }); }); + // #11692: shutdown waits for a started OAuth refresh so the rotated token is not lost. + test("settled waits for started modifications, not for ones waiting on the lock", async () => { + let releaseLock!: () => void; + const lock = new Promise((resolve) => { + releaseLock = resolve; + }); + let lockHeld = false; + const stored = new Map(); + const storage: CredentialStore = { + read: async (providerId) => stored.get(providerId), + list: async () => [], + modify: async (providerId, fn) => { + // The first modification holds the lock until released; later ones wait for it. + if (lockHeld) await lock; + lockHeld = true; + const next = await fn(stored.get(providerId)); + if (next) stored.set(providerId, next); + return next; + }, + delete: async () => {}, + }; + const credentials = new RuntimeCredentials(storage); + + let finishRefresh!: () => void; + const refreshing = new Promise((resolve) => { + finishRefresh = resolve; + }); + const started = credentials.modify("anthropic", async () => { + await refreshing; + return { type: "api_key", key: "rotated" }; + }); + const waiting = credentials.modify("openai", async () => ({ type: "api_key", key: "never-started" })); + + let settled = false; + const settling = credentials.settled().then(() => { + settled = true; + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + expect(settled).toBe(false); + + finishRefresh(); + await settling; + expect(stored.get("anthropic")).toEqual({ type: "api_key", key: "rotated" }); + expect(stored.has("openai")).toBe(false); + + await started; + releaseLock(); + await waiting; + }); + + test("settled resolves after a started modification fails", async () => { + const storage = AuthStorage.inMemory(); + const credentials = new RuntimeCredentials(storage); + let markStarted!: () => void; + const started = new Promise((resolve) => { + markStarted = resolve; + }); + let failRefresh!: () => void; + const refreshFailed = new Promise((resolve) => { + failRefresh = resolve; + }); + const failed = credentials.modify("anthropic", async () => { + markStarted(); + await refreshFailed; + throw new Error("invalid_grant"); + }); + const rejected = expect(failed).rejects.toThrow("invalid_grant"); + + await started; + const settling = credentials.settled(); + failRefresh(); + await settling; + await rejected; + }); + test("delete clears both the override and persisted credential", async () => { const storage = AuthStorage.inMemory({ anthropic: { type: "api_key", key: "stored-key" } }); const credentials = new RuntimeCredentials(storage); diff --git a/packages/coding-agent/test/suite/regressions/11692-dispose-waits-for-oauth-refresh.test.ts b/packages/coding-agent/test/suite/regressions/11692-dispose-waits-for-oauth-refresh.test.ts new file mode 100644 index 000000000..25f7a7fa9 --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/11692-dispose-waits-for-oauth-refresh.test.ts @@ -0,0 +1,117 @@ +import type { Model, OAuthCredential, Provider } from "@earendil-works/pi-ai"; +import { afterEach, describe, expect, it } from "vitest"; +import { + AgentSessionRuntime, + type AgentSessionServices, + type CreateAgentSessionRuntimeFactory, +} from "../../../src/core/agent-session-runtime.ts"; +import { AuthStorage } from "../../../src/core/auth-storage.ts"; +import { ModelRuntime } from "../../../src/core/model-runtime.ts"; +import { createHarness, type Harness } from "../harness.ts"; + +const providerId = "rotating-oauth"; + +function oauthProvider( + refresh: (credential: OAuthCredential) => Promise, +): Provider<"openai-completions"> { + const model: Model<"openai-completions"> = { + id: "model", + name: "Model", + api: "openai-completions", + provider: providerId, + baseUrl: "https://example.test/v1", + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 1000, + maxTokens: 100, + }; + return { + id: providerId, + name: providerId, + auth: { + oauth: { + name: "OAuth", + login: async () => { + throw new Error("unused"); + }, + refresh, + toAuth: async (credential) => ({ apiKey: credential.access }), + }, + }, + getModels: () => [model], + // Like the built-in catalogs, the model list does not need the credential, but the catalog + // refresh still resolves (and so refreshes) an expired one. + refreshModels: async () => {}, + stream: () => { + throw new Error("unused"); + }, + streamSimple: () => { + throw new Error("unused"); + }, + }; +} + +describe("regression #11692: OAuth refresh in flight at shutdown", () => { + const harnesses: Harness[] = []; + + afterEach(() => { + while (harnesses.length > 0) { + harnesses.pop()?.cleanup(); + } + }); + + it("dispose waits until the rotated token from a background catalog refresh is saved", async () => { + const credentials = AuthStorage.inMemory({ + [providerId]: { type: "oauth", access: "old-access", refresh: "old-refresh", expires: Date.now() - 1000 }, + }); + let markRefreshStarted!: () => void; + const refreshStarted = new Promise((resolve) => { + markRefreshStarted = resolve; + }); + let respond!: () => void; + const response = new Promise((resolve) => { + respond = resolve; + }); + const modelRuntime = await ModelRuntime.create({ credentials, modelsPath: null, allowModelNetwork: false }); + modelRuntime.registerNativeProvider( + oauthProvider(async () => { + // The token endpoint has already spent "old-refresh" when this request is in flight. + markRefreshStarted(); + await response; + return { type: "oauth", access: "new-access", refresh: "new-refresh", expires: Date.now() + 3_600_000 }; + }), + ); + + const harness = await createHarness(); + harnesses.push(harness); + const services: AgentSessionServices = { + cwd: harness.tempDir, + agentDir: harness.tempDir, + modelRuntime, + settingsManager: harness.settingsManager, + resourceLoader: harness.session.resourceLoader, + diagnostics: [], + }; + const unusedFactory: CreateAgentSessionRuntimeFactory = async () => { + throw new Error("runtime replacement is not used in this test"); + }; + const runtime = new AgentSessionRuntime(harness.session, services, unusedFactory); + + // The RPC startup catalog refresh runs in the background and is not awaited by shutdown. + const catalogRefresh = modelRuntime.refresh({ allowNetwork: true, providers: [providerId] }); + await refreshStarted; + + let disposed = false; + const disposing = runtime.dispose().then(() => { + disposed = true; + }); + await new Promise((resolve) => setTimeout(resolve, 20)); + expect(disposed).toBe(false); + + respond(); + await disposing; + expect(await credentials.read(providerId)).toMatchObject({ access: "new-access", refresh: "new-refresh" }); + await catalogRefresh; + }); +});