fix(coding-agent): track every pending credential write, drop the shutdown timeout

Simpler and matches the MCP track()/settled() idiom: tracking also covers a
write still waiting for the auth.json lock, so one snapshot is enough. The
refresh (15 s) and lock wait (30 s) already bound the wait.
This commit is contained in:
Petr Houška 2026-11-04 23:10:16 +01:00
parent 94fdba50cc
commit d315415697
4 changed files with 33 additions and 102 deletions

View file

@ -30,7 +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))
- 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

View file

@ -10,7 +10,6 @@ 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";
@ -403,9 +402,9 @@ export class AgentSessionRuntime {
}
/**
* 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).
* 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<void> {
await this.session.abort();
@ -415,22 +414,7 @@ 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<void> {
let timer: NodeJS.Timeout | undefined;
const timeout = new Promise<void>((resolve) => {
timer = setTimeout(resolve, CREDENTIAL_WRITE_SETTLE_TIMEOUT_MS);
});
try {
await Promise.race([modelRuntime.settleCredentialWrites(), timeout]);
} finally {
clearTimeout(timer);
await this.services.modelRuntime.settleCredentialWrites();
}
}

View file

@ -4,8 +4,8 @@ import type { AuthOperationOptions, Credential, CredentialInfo, CredentialStore
export class RuntimeCredentials implements CredentialStore {
private readonly store: CredentialStore;
private readonly overrides = new Map<string, string>();
/** Modifications whose update function has started, for example an OAuth token request. */
private readonly startedModifications = new Set<Promise<void>>();
/** Modifications not yet written or failed, for example an OAuth refresh. */
private readonly pendingModifications = new Set<Promise<void>>();
constructor(store: CredentialStore) {
this.store = store;
@ -43,37 +43,23 @@ export class RuntimeCredentials implements CredentialStore {
fn: (current: Credential | undefined) => Promise<Credential | undefined>,
options?: AuthOperationOptions,
): Promise<Credential | undefined> {
let markDone!: () => void;
const done = new Promise<void>((resolve) => {
markDone = resolve;
});
const modification = this.store.modify(
providerId,
(current) => {
this.startedModifications.add(done);
return fn(current);
},
options,
const modification = this.store.modify(providerId, fn, options);
const pending = modification.then(
() => undefined,
() => undefined,
);
const finish = () => {
this.startedModifications.delete(done);
markDone();
};
modification.then(finish, finish);
this.pendingModifications.add(pending);
void pending.finally(() => this.pendingModifications.delete(pending));
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.
* 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<void> {
while (this.startedModifications.size > 0) {
await Promise.all(this.startedModifications);
}
await Promise.all(this.pendingModifications);
}
async delete(providerId: string, options?: AuthOperationOptions): Promise<void> {

View file

@ -1,4 +1,4 @@
import type { Credential, CredentialStore } from "@earendil-works/pi-ai";
import type { 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,38 +73,18 @@ 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<void>((resolve) => {
releaseLock = resolve;
});
let lockHeld = false;
const stored = new Map<string, Credential>();
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 () => {},
};
// #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 finishRefresh!: () => void;
const refreshing = new Promise<void>((resolve) => {
finishRefresh = resolve;
let respond!: () => void;
const response = new Promise<void>((resolve) => {
respond = resolve;
});
const started = credentials.modify("anthropic", async () => {
await refreshing;
const modification = credentials.modify("anthropic", async () => {
await response;
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(() => {
@ -113,39 +93,20 @@ describe("RuntimeCredentials", () => {
await new Promise((resolve) => setTimeout(resolve, 10));
expect(settled).toBe(false);
finishRefresh();
respond();
await settling;
expect(stored.get("anthropic")).toEqual({ type: "api_key", key: "rotated" });
expect(stored.has("openai")).toBe(false);
await started;
releaseLock();
await waiting;
expect(await storage.read("anthropic")).toEqual({ type: "api_key", key: "rotated" });
await modification;
});
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<void>((resolve) => {
markStarted = resolve;
});
let failRefresh!: () => void;
const refreshFailed = new Promise<void>((resolve) => {
failRefresh = resolve;
});
test("settled resolves after a pending modification fails", async () => {
const credentials = new RuntimeCredentials(AuthStorage.inMemory());
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;
await credentials.settled();
await expect(failed).rejects.toThrow("invalid_grant");
});
test("delete clears both the override and persisted credential", async () => {