From b2cb629237496524e745ff668ec5bd5d47130164 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Petr=20Hou=C5=A1ka?= Date: Thu, 5 Nov 2026 00:32:41 +0100 Subject: [PATCH] fix(coding-agent): wait for a started OAuth refresh to be saved before exit (#11812) Fixes #11692. ## Problem In RPC mode the startup catalog refresh runs in the background and also refreshes an expired OAuth credential. If pi exits after the token request went out but before the new token is written, the provider has already spent the old single-use refresh token and auth.json keeps only that one. The next request fails with `invalid_grant`. Trace: `pi --mode rpc` starts, `get_state` is answered, catalog refresh sends the token request, the client sends SIGTERM (or closes stdin), `shutdown()` runs `runtimeHost.dispose()`, which does not know about the refresh, and the process exits. Same for print and interactive mode, and for a prompt-time refresh cut off by an exit. ## Fix - `RuntimeCredentials.modify()` tracks every pending modification (same idiom as the MCP `track()` / `McpAuthProvider.settled()`). - `RuntimeCredentials.settled()` / `ModelRuntime.settleCredentialWrites()` wait for them. - `AgentSessionRuntime.dispose()` waits for them. All modes exit through `dispose()`. No extra timeout: the OAuth refresh (15 s) and the auth.json lock wait (30 s) are already bounded, and a second signal still exits at once. Not changed: the catalog refresh still refreshes expired OAuth credentials, also for providers whose catalog does not use them (the pi.dev catalog for `openai`, `anthropic`, ...). Skipping that would need a per-provider "catalog needs auth" signal in pi-ai; it is optional now that exits wait. `docs/rpc.md` now says that pi waits for pending work before exiting. ## Testing - `test/runtime-credentials.test.ts`: `settled()` waits until a pending modification is written, and resolves after a failed one. - `test/suite/regressions/11692-dispose-waits-for-oauth-refresh.test.ts`: `dispose()` waits until the token from a background catalog refresh is saved. Fails without the fix. - Real process: `pi --mode rpc` with an extension OAuth provider (expired credential, 400 ms refresh), SIGTERM about 100 ms after the token request started. Without the wait: 3 of 3 runs kept the spent refresh token. With the fix: 3 of 3 saved the new one (exit 143). Closing stdin instead of SIGTERM: new token saved, exit 0. - `npm run check`; `test/rpc*`, `test/print-mode`, `test/model-runtime*`, runtime and regression suites pass. --- packages/coding-agent/CHANGELOG.md | 1 + packages/coding-agent/docs/rpc.md | 2 +- .../src/core/agent-session-runtime.ts | 6 + .../coding-agent/src/core/model-runtime.ts | 8 ++ .../src/core/runtime-credentials.ts | 20 ++- .../test/runtime-credentials.test.ts | 36 ++++++ ...92-dispose-waits-for-oauth-refresh.test.ts | 117 ++++++++++++++++++ 7 files changed, 188 insertions(+), 2 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..e161043b0 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 for a token refresh in progress 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..211bfb81c 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. Pi disposes the active runtime and waits for pending work to finish before exiting, so give it 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..854ba3714 100644 --- a/packages/coding-agent/src/core/agent-session-runtime.ts +++ b/packages/coding-agent/src/core/agent-session-runtime.ts @@ -401,6 +401,11 @@ export class AgentSessionRuntime { return { cancelled: false }; } + /** + * Dispose the session for process exit. Also waits 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 +414,7 @@ export class AgentSessionRuntime { }); this.beforeSessionInvalidate?.(); this.session.dispose(); + await this.services.modelRuntime.settleCredentialWrites(); } } 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..a4b123b32 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 not yet written or failed, for example an OAuth refresh. */ + private readonly pendingModifications = new Set>(); constructor(store: CredentialStore) { this.store = store; @@ -41,7 +43,23 @@ export class RuntimeCredentials implements CredentialStore { fn: (current: Credential | undefined) => Promise, options?: AuthOperationOptions, ): Promise { - return this.store.modify(providerId, fn, options); + const modification = this.store.modify(providerId, fn, options); + const pending = modification.then( + () => undefined, + () => undefined, + ); + this.pendingModifications.add(pending); + void pending.finally(() => this.pendingModifications.delete(pending)); + return modification; + } + + /** + * Resolves when every modification started so far is written or failed. An OAuth refresh may + * already have spent the single-use refresh token, so exiting before its result is written would + * leave only the spent token in auth.json. + */ + async settled(): Promise { + await Promise.all(this.pendingModifications); } 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..4052fabdd 100644 --- a/packages/coding-agent/test/runtime-credentials.test.ts +++ b/packages/coding-agent/test/runtime-credentials.test.ts @@ -73,6 +73,42 @@ describe("RuntimeCredentials", () => { expect(await credentials.read("anthropic")).toEqual({ type: "api_key", key: "runtime-key" }); }); + // #11692: shutdown waits for an OAuth refresh in progress so the rotated token is not lost. + test("settled waits until pending modifications are written", async () => { + const storage = AuthStorage.inMemory({ anthropic: { type: "api_key", key: "old" } }); + const credentials = new RuntimeCredentials(storage); + let respond!: () => void; + const response = new Promise((resolve) => { + respond = resolve; + }); + const modification = credentials.modify("anthropic", async () => { + await response; + return { type: "api_key", key: "rotated" }; + }); + + let settled = false; + const settling = credentials.settled().then(() => { + settled = true; + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + expect(settled).toBe(false); + + respond(); + await settling; + expect(await storage.read("anthropic")).toEqual({ type: "api_key", key: "rotated" }); + await modification; + }); + + test("settled resolves after a pending modification fails", async () => { + const credentials = new RuntimeCredentials(AuthStorage.inMemory()); + const failed = credentials.modify("anthropic", async () => { + throw new Error("invalid_grant"); + }); + + await credentials.settled(); + await expect(failed).rejects.toThrow("invalid_grant"); + }); + 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; + }); +});