diff --git a/packages/coding-agent/src/core/model-runtime.ts b/packages/coding-agent/src/core/model-runtime.ts index 7edd44048..368fc9c0c 100644 --- a/packages/coding-agent/src/core/model-runtime.ts +++ b/packages/coding-agent/src/core/model-runtime.ts @@ -97,13 +97,24 @@ interface RegisteredVirtualModel { } interface ModelRuntimeSnapshot { - all: readonly Model[]; + /** Catalog chat models whose provider has auth. */ available: readonly Model[]; - configuredProviders: ReadonlySet; + /** Auth check result of every provider that has auth. */ + auth: ReadonlyMap; + /** Provider ids with a stored or runtime credential, including providers that are not registered. */ storedProviders: ReadonlySet; - auth: ReadonlyMap; } +/** Availability check that a later request can still widen because it has not started. */ +interface WaitingAvailabilityPass { + /** Provider ids to check, or undefined for every provider. */ + providerIds: Set | undefined; + result: Promise>; +} + +/** Upper bound for one availability pass, so a hung auth check or credential store cannot block callers. */ +const AVAILABILITY_PASS_TIMEOUT_MS = 10_000; + export interface CreateModelRuntimeOptions { /** Credential storage. Defaults to the file at authPath. */ credentials?: CredentialStore; @@ -182,15 +193,13 @@ export class ModelRuntime implements Models { private readonly modelNetworkEnabled: boolean; private config: ModelConfig; private snapshot: ModelRuntimeSnapshot = { - all: [], available: [], - configuredProviders: new Set(), - storedProviders: new Set(), auth: new Map(), + storedProviders: new Set(), }; - private availabilityRefreshSeq = 0; - private availabilityErrorSeq = 0; - private readonly providerAvailabilitySeq = new Map(); + /** Settles when the last requested availability pass finishes. Passes run one at a time. */ + private availabilityTail: Promise = Promise.resolve(); + private waitingAvailabilityPass: WaitingAvailabilityPass | undefined; private availabilityError: string | undefined; private readonly credentialOperations = new Map>(); @@ -323,110 +332,105 @@ export class ModelRuntime implements Models { } private updateModelSnapshot(): void { - const all = [...this.models.getModels()]; this.snapshot = { ...this.snapshot, - all, - available: all.filter((model) => this.snapshot.configuredProviders.has(model.provider)), + available: this.models.getModels().filter((model) => this.snapshot.auth.has(model.provider)), }; } - private async runAvailabilityRefresh(seq: number, errorSeq: number, signal: AbortSignal): Promise { - const providers = this.models.getProviders(); - const [available, checks, credentials] = await Promise.all([ - this.models.getAvailable(undefined, { signal }), - Promise.all( - providers.map( - async (provider): Promise<[string, AuthCheck | undefined]> => [ - provider.id, - await this.models.checkAuth(provider.id, { signal }), - ], - ), - ), - this.credentials.list({ signal }), - ]); - if (seq !== this.availabilityRefreshSeq) return; - const auth = new Map(checks); - const configuredProviders = new Set( - checks - .filter((entry): entry is [string, AuthCheck] => entry[1] !== undefined) - .map(([providerId]) => providerId), - ); - this.snapshot = { - all: [...this.models.getModels()], - available: [...available], - configuredProviders, - storedProviders: new Set(credentials.map((entry) => entry.providerId)), - auth, - }; - if (errorSeq === this.availabilityErrorSeq) this.availabilityError = undefined; - } - - private queueAvailabilityRefresh(signal?: AbortSignal): Promise { - const seq = ++this.availabilityRefreshSeq; - for (const [providerId, providerSeq] of this.providerAvailabilitySeq) { - this.providerAvailabilitySeq.set(providerId, providerSeq + 1); + /** + * Check auth for `providerIds`, or for every provider when undefined, and resolve with the checks that + * failed. Passes run one at a time, so the last pass to save is always the newest. A request joins the + * waiting pass when there is one: it has not read anything yet, so it will see the requester's changes. + * The caller's signal only stops its own wait; the shared pass keeps running for the other callers. + */ + private requestAvailabilityPass( + providerIds: readonly string[] | undefined, + signal?: AbortSignal, + ): Promise> { + const waiting = this.waitingAvailabilityPass; + if (waiting) { + if (!providerIds) waiting.providerIds = undefined; + else for (const id of providerIds) waiting.providerIds?.add(id); + return raceWithAbortSignal(waiting.result, signal); } - const errorSeq = ++this.availabilityErrorSeq; - const effectiveSignal = operationSignal(signal); - return this.runAvailabilityRefresh(seq, errorSeq, effectiveSignal).catch((error) => { - if (errorSeq === this.availabilityErrorSeq && !effectiveSignal.aborted) { - this.availabilityError = error instanceof Error ? error.message : String(error); - } - throw error; - }); + const pass: WaitingAvailabilityPass = { + providerIds: providerIds && new Set(providerIds), + result: this.availabilityTail.then(() => { + this.waitingAvailabilityPass = undefined; + return this.runAvailabilityPass(pass.providerIds); + }), + }; + this.waitingAvailabilityPass = pass; + this.availabilityTail = pass.result.catch(() => {}); + return raceWithAbortSignal(pass.result, signal); } - private async refreshProviderAvailability(providerId: string, signal: AbortSignal): Promise { - // Invalidate any full availability pass that started before this credential change. - ++this.availabilityRefreshSeq; - const providerSeq = (this.providerAvailabilitySeq.get(providerId) ?? 0) + 1; - this.providerAvailabilitySeq.set(providerId, providerSeq); - const errorSeq = ++this.availabilityErrorSeq; + /** + * A provider whose check fails or times out keeps its previous entry and is reported in the result. + * A failed credential list fails the whole pass without saving anything. + */ + private async runAvailabilityPass( + providerIds: ReadonlySet | undefined, + ): Promise> { + const controller = new AbortController(); + const timeout = setTimeout( + () => controller.abort(new Error(`Availability check timed out after ${AVAILABILITY_PASS_TIMEOUT_MS} ms`)), + AVAILABILITY_PASS_TIMEOUT_MS, + ); + const signal = controller.signal; try { - const [available, auth, credential] = await Promise.all([ - this.models.getAvailable(providerId, { signal }), - this.models.checkAuth(providerId, { signal }), - this.credentials.read(providerId, { signal }), + const ids = providerIds ? [...providerIds] : this.models.getProviders().map((provider) => provider.id); + const [credentials, results] = await Promise.all([ + raceWithAbortSignal(this.credentials.list({ signal }), signal), + Promise.all( + ids.map(async (id) => { + try { + // Provider checks run extension code that may ignore the signal. + const [check, available] = await raceWithAbortSignal( + Promise.all([this.models.checkAuth(id, { signal }), this.models.getAvailable(id, { signal })]), + signal, + ); + return { id, check, available }; + } catch (error) { + return { id, error: error instanceof Error ? error : new Error(String(error)) }; + } + }), + ), ]); - signal.throwIfAborted(); - if (this.providerAvailabilitySeq.get(providerId) !== providerSeq) return; - const configuredProviders = new Set(this.snapshot.configuredProviders); - const storedProviders = new Set(this.snapshot.storedProviders); - const authByProvider = new Map(this.snapshot.auth); - if (auth) { - configuredProviders.add(providerId); - authByProvider.set(providerId, auth); - } else { - configuredProviders.delete(providerId); - authByProvider.delete(providerId); + const auth = new Map(this.snapshot.auth); + const checkedAvailable = new Map[]>(); + const errors = new Map(); + for (const result of results) { + if (result.error) { + errors.set(result.id, result.error); + continue; + } + checkedAvailable.set(result.id, result.available); + if (result.check) auth.set(result.id, result.check); + else auth.delete(result.id); } - if (credential) storedProviders.add(providerId); - else storedProviders.delete(providerId); - const all = [...this.models.getModels()]; - const availableById = new Map( - [...this.snapshot.available.filter((model) => model.provider !== providerId), ...available].map((model) => [ - `${model.provider}\0${model.id}`, - model, - ]), - ); + for (const id of auth.keys()) { + if (!this.models.getProvider(id)) auth.delete(id); + } + const previousAvailable = Map.groupBy(this.snapshot.available, (model) => model.provider); this.snapshot = { - all, - available: all.flatMap((model) => availableById.get(`${model.provider}\0${model.id}`) ?? []), - configuredProviders, - storedProviders, - auth: authByProvider, + available: this.models + .getProviders() + .flatMap(({ id }) => + auth.has(id) ? (checkedAvailable.get(id) ?? previousAvailable.get(id) ?? []) : [], + ), + auth, + storedProviders: new Set(credentials.map((entry) => entry.providerId)), }; - if (errorSeq === this.availabilityErrorSeq) this.availabilityError = undefined; + this.availabilityError = + errors.size > 0 ? [...errors.values()].map((error) => error.message).join("\n") : undefined; + return errors; } catch (error) { - if ( - this.providerAvailabilitySeq.get(providerId) === providerSeq && - errorSeq === this.availabilityErrorSeq && - !signal.aborted - ) { - this.availabilityError = error instanceof Error ? error.message : String(error); - } + this.availabilityError = error instanceof Error ? error.message : String(error); throw error; + } finally { + clearTimeout(timeout); } } @@ -479,21 +483,11 @@ export class ModelRuntime implements Models { } async getAvailable(providerId?: string, options?: AuthOperationOptions): Promise[]> { - if (providerId) { - const errorSeq = ++this.availabilityErrorSeq; - try { - const available = await this.models.getAvailable(providerId, options); - if (errorSeq === this.availabilityErrorSeq) this.availabilityError = undefined; - return available; - } catch (error) { - if (errorSeq === this.availabilityErrorSeq && !options?.signal?.aborted) { - this.availabilityError = error instanceof Error ? error.message : String(error); - } - throw error; - } - } - await this.queueAvailabilityRefresh(options?.signal); - return this.snapshot.available; + const errors = await this.requestAvailabilityPass(providerId ? [providerId] : undefined, options?.signal); + if (!providerId) return this.snapshot.available; + const error = errors.get(providerId); + if (error) throw error; + return this.snapshot.available.filter((model) => model.provider === providerId); } getAvailableSnapshot(): readonly Model[] { @@ -541,7 +535,7 @@ export class ModelRuntime implements Models { } hasConfiguredAuth(providerId: string): boolean { - return this.snapshot.configuredProviders.has(providerId); + return this.snapshot.auth.has(providerId); } getAuth(providerId: string, overrides?: ModelRuntimeAuthOverrides): Promise; @@ -604,7 +598,8 @@ export class ModelRuntime implements Models { const refreshError = result.errors.get(providerId); if (refreshError) throw refreshError; this.updateModelSnapshot(); - await this.refreshProviderAvailability(providerId, signal); + const availabilityError = (await this.requestAvailabilityPass([providerId], signal)).get(providerId); + if (availabilityError) throw availabilityError; } catch (cause) { throw new CredentialSynchronizationError(providerId, operation, credential, { cause }); } @@ -857,23 +852,18 @@ export class ModelRuntime implements Models { }; const errors = new Map(result.errors); this.updateModelSnapshot(); - if (options.providers) { - await Promise.all( - [...new Set(options.providers)].map(async (providerId) => { - try { - await this.refreshProviderAvailability(providerId, operationSignal(options.signal)); - } catch (error) { - if (!options.signal?.aborted) { - errors.set(providerId, error instanceof Error ? error : new Error(String(error))); - } - } - }), - ); - } else { - try { - await this.queueAvailabilityRefresh(options.signal); - } catch { - // Availability errors are recorded by the latest pass; refreshed models remain usable. + try { + const availabilityErrors = await this.requestAvailabilityPass(options.providers, options.signal); + for (const providerId of options.providers ?? []) { + const error = availabilityErrors.get(providerId); + if (error) errors.set(providerId, error); + } + } catch (error) { + // A failed pass is recorded in availabilityError; refreshed models remain usable. + if (options.providers && !options.signal?.aborted) { + for (const providerId of options.providers) { + errors.set(providerId, error instanceof Error ? error : new Error(String(error))); + } } } return { aborted: result.aborted || (options.signal?.aborted ?? false), errors }; @@ -904,16 +894,13 @@ export class ModelRuntime implements Models { type: AuthType, ): void { if (!this.snapshot.storedProviders.has(providerId) && !configuredStatus?.configured) return; - const configuredProviders = new Set(this.snapshot.configuredProviders).add(providerId); - const auth = new Map(this.snapshot.auth); // Never clobber a real check result. - if (!auth.get(providerId)) auth.set(providerId, { type, source: "configured provider" }); + if (this.snapshot.auth.has(providerId)) return; this.snapshot = { ...this.snapshot, - auth, - configuredProviders, - available: this.snapshot.all.filter((model) => configuredProviders.has(model.provider)), + auth: new Map(this.snapshot.auth).set(providerId, { type, source: "configured provider" }), }; + this.updateModelSnapshot(); } registerProvider(providerId: string, config: ProviderConfigInput): void { @@ -962,12 +949,11 @@ export class ModelRuntime implements Models { const models = this.virtualModels.get(providerId) ?? new Map(); models.set(id, { model: createVirtualModel(definition), route: (request) => definition.route(request) }); this.virtualModels.set(providerId, models); - if (!this.recomposeProvider(providerId) && !this.snapshot.configuredProviders.has(providerId)) { + if (!this.recomposeProvider(providerId) && !this.snapshot.auth.has(providerId)) { // A provider of only virtual models needs no credentials. Mark it configured now: session // restore checks auth before the refresh below lands. const auth = new Map(this.snapshot.auth).set(providerId, { type: "api_key", source: "virtual" }); - const configuredProviders = new Set(this.snapshot.configuredProviders).add(providerId); - this.snapshot = { ...this.snapshot, auth, configuredProviders }; + this.snapshot = { ...this.snapshot, auth }; } this.updateModelSnapshot(); void this.refresh({ allowNetwork: false }); diff --git a/packages/coding-agent/test/model-runtime-availability-order.test.ts b/packages/coding-agent/test/model-runtime-availability-order.test.ts new file mode 100644 index 000000000..2f0463dd9 --- /dev/null +++ b/packages/coding-agent/test/model-runtime-availability-order.test.ts @@ -0,0 +1,156 @@ +import type { Model, Provider } from "@earendil-works/pi-ai"; +import { describe, expect, it, vi } from "vitest"; +import { AuthStorage } from "../src/core/auth-storage.ts"; +import { ModelRuntime } from "../src/core/model-runtime.ts"; + +function testProvider(id: string, isConfigured: () => boolean, onCheck: () => void = () => {}): Provider { + const model: Model<"openai-completions"> = { + id: `${id}-model`, + name: `${id} model`, + api: "openai-completions", + provider: id, + baseUrl: "https://example.test/v1", + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 1000, + maxTokens: 100, + }; + const unused = () => { + throw new Error("unused"); + }; + return { + id, + name: id, + auth: { + apiKey: { + name: "API key", + check: async () => { + onCheck(); + return isConfigured() ? { type: "api_key", source: "test" } : undefined; + }, + resolve: async () => (isConfigured() ? { auth: { apiKey: "key" }, source: "test" } : undefined), + }, + }, + getModels: () => [model], + stream: unused, + streamSimple: unused, + }; +} + +/** Hold the `n`th credentials.list() call from now; each availability pass makes exactly one. */ +function holdCredentialListCall( + credentials: AuthStorage, + n: number, +): { calls: () => number; started: Promise; release: () => void } { + const list = credentials.list.bind(credentials); + let calls = 0; + let markStarted = () => {}; + const started = new Promise((resolve) => { + markStarted = resolve; + }); + let release = () => {}; + const gate = new Promise((resolve) => { + release = resolve; + }); + credentials.list = async (options) => { + calls++; + if (calls !== n) return list(options); + markStarted(); + await gate; + return list(options); + }; + return { calls: () => calls, started, release }; +} + +async function createRuntime(credentials: AuthStorage): Promise { + return ModelRuntime.create({ credentials, modelsPath: null, allowModelNetwork: false }); +} + +describe("availability passes", () => { + // Regression for #8810 and #9884: startup awaits a pass while a registration requests another one. + it("resolves an awaited request from its own pass while a later request waits", async () => { + const credentials = AuthStorage.inMemory(); + const runtime = await createRuntime(credentials); + let configured = false; + runtime.registerNativeProvider(testProvider("dynamic", () => configured)); + await runtime.refresh({ allowNetwork: false }); + expect(runtime.hasConfiguredAuth("dynamic")).toBe(false); + + configured = true; + const later = holdCredentialListCall(credentials, 2); + const awaited = runtime.getAvailable(); + await vi.waitFor(() => expect(later.calls()).toBe(1)); + const laterPass = runtime.getAvailable(); + + await awaited; + expect(runtime.hasConfiguredAuth("dynamic")).toBe(true); + + later.release(); + await laterPass; + expect(runtime.hasConfiguredAuth("dynamic")).toBe(true); + }); + + // Regression for #8810: a login during a full pass must not lose the login or the rest of the pass. + it("runs a login check after a running full pass and keeps both results", async () => { + const credentials = AuthStorage.inMemory(); + const runtime = await createRuntime(credentials); + let dynamicConfigured = false; + let dynamicChecks = 0; + let otherConfigured = false; + runtime.registerNativeProvider( + testProvider( + "dynamic", + () => dynamicConfigured, + () => dynamicChecks++, + ), + ); + runtime.registerNativeProvider(testProvider("other", () => otherConfigured)); + await runtime.refresh({ allowNetwork: false }); + + otherConfigured = true; + const fullPass = holdCredentialListCall(credentials, 1); + const checksBefore = dynamicChecks; + const running = runtime.getAvailable(); + await fullPass.started; + // The full pass has read "dynamic" as not configured. + await vi.waitFor(() => expect(dynamicChecks).toBeGreaterThan(checksBefore)); + + dynamicConfigured = true; + const login = runtime.refresh({ allowNetwork: false, providers: ["dynamic"] }); + fullPass.release(); + await running; + await login; + expect(runtime.hasConfiguredAuth("dynamic")).toBe(true); + expect(runtime.hasConfiguredAuth("other")).toBe(true); + }); + + it("shares one waiting pass between requests made while another pass runs", async () => { + const credentials = AuthStorage.inMemory(); + const runtime = await createRuntime(credentials); + const running = holdCredentialListCall(credentials, 1); + const first = runtime.getAvailable(); + await running.started; + + const waiting = [runtime.getAvailable(), runtime.getAvailable(), runtime.getAvailable()]; + running.release(); + await Promise.all([first, ...waiting]); + expect(running.calls()).toBe(2); + }); + + it("keeps a registration's provisional entry when a pass that started earlier finishes", async () => { + const credentials = AuthStorage.inMemory({ dynamic: { type: "api_key", key: "stored" } }); + const runtime = await createRuntime(credentials); + + const olderPass = holdCredentialListCall(credentials, 1); + const olderFullPass = runtime.getAvailable(); + await olderPass.started; + runtime.registerNativeProvider(testProvider("dynamic", () => true)); + expect(runtime.hasConfiguredAuth("dynamic")).toBe(true); + + // The older pass did not check "dynamic"; the registration's own pass is still queued. + olderPass.release(); + await olderFullPass; + expect(runtime.hasConfiguredAuth("dynamic")).toBe(true); + }); +}); diff --git a/packages/coding-agent/test/suite/regressions/7301-stalled-availability-refresh.test.ts b/packages/coding-agent/test/suite/regressions/7301-stalled-availability-refresh.test.ts index fc424d9ff..9afd09cf4 100644 --- a/packages/coding-agent/test/suite/regressions/7301-stalled-availability-refresh.test.ts +++ b/packages/coding-agent/test/suite/regressions/7301-stalled-availability-refresh.test.ts @@ -50,42 +50,50 @@ function stallNextCredentialList(harness: Harness): StalledCredentialList { }; } -async function waitForRecoveryRefresh(stalledList: StalledCredentialList, harness: Harness): Promise { +/** Queue a refresh behind the stalled pass, then let the stalled pass time out. */ +async function timeOutStalledPass(harness: Harness): Promise { const refresh = harness.session.modelRuntime.refresh({ allowNetwork: false }); - await vi.waitFor(() => expect(stalledList.getCallCount()).toBe(2)); + await vi.advanceTimersByTimeAsync(AVAILABILITY_PASS_TIMEOUT_MS); await refresh; } +const AVAILABILITY_PASS_TIMEOUT_MS = 10_000; +const TIMEOUT_MESSAGE = `Availability check timed out after ${AVAILABILITY_PASS_TIMEOUT_MS} ms`; + describe("issue #7301 stalled availability refresh", () => { let harness: Harness | undefined; let stalledList: StalledCredentialList | undefined; afterEach(() => { + vi.useRealTimers(); stalledList?.release(); stalledList = undefined; harness?.cleanup(); harness = undefined; }); - it("recovers without letting the stalled refresh overwrite the newer snapshot", async () => { + it("times out a stalled pass so later passes save, and ignores its late result", async () => { harness = await createHarness({ withConfiguredAuth: false }); const runtime = harness.session.modelRuntime; await harness.authStorage.modify("stale-provider", async () => ({ type: "api_key", key: "stale-key" })); await runtime.refresh({ allowNetwork: false }); expect(runtime.getProviderAuthStatus("stale-provider")).toEqual({ configured: true, source: "stored" }); + vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout"] }); stalledList = stallNextCredentialList(harness); - const staleRefresh = runtime.getAvailable(); + const staleRefresh = expect(runtime.getAvailable()).rejects.toThrow(TIMEOUT_MESSAGE); await stalledList.started; await harness.authStorage.delete("stale-provider"); await harness.authStorage.modify("current-provider", async () => ({ type: "api_key", key: "current-key" })); - await waitForRecoveryRefresh(stalledList, harness); + await timeOutStalledPass(harness); + await staleRefresh; + expect(stalledList.getCallCount()).toBe(2); expect(runtime.getProviderAuthStatus("stale-provider")).toEqual({ configured: false }); expect(runtime.getProviderAuthStatus("current-provider")).toEqual({ configured: true, source: "stored" }); stalledList.release(); - await staleRefresh; + await vi.waitFor(() => expect(stalledList?.getCallCount()).toBe(2)); expect(runtime.getProviderAuthStatus("stale-provider")).toEqual({ configured: false }); expect(runtime.getProviderAuthStatus("current-provider")).toEqual({ configured: true, source: "stored" }); }); @@ -93,41 +101,47 @@ describe("issue #7301 stalled availability refresh", () => { it("does not let a stale failure overwrite newer availability error state", async () => { harness = await createHarness({ withConfiguredAuth: false }); const runtime = harness.session.modelRuntime; + vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout"] }); stalledList = stallNextCredentialList(harness); - const staleRefresh = runtime.getAvailable(); + const staleRefresh = expect(runtime.getAvailable()).rejects.toThrow(TIMEOUT_MESSAGE); await stalledList.started; - await waitForRecoveryRefresh(stalledList, harness); + await timeOutStalledPass(harness); + await staleRefresh; expect(runtime.getError()).toBeUndefined(); stalledList.fail(new Error("stale credential list failure")); - await expect(staleRefresh).rejects.toThrow("stale credential list failure"); + await Promise.resolve(); expect(runtime.getError()).toBeUndefined(); }); - it("does not let a stale provider-scoped failure overwrite a newer availability pass", async () => { + it("times out a provider check that never settles without blocking later passes", async () => { harness = await createHarness(); const runtime = harness.session.modelRuntime; + const provider = harness.getModel().provider; const models = Reflect.get(runtime, "models") as Models; const originalGetAvailable = models.getAvailable.bind(models); const started = createDeferred(); const gate = createDeferred(); let stall = true; models.getAvailable = async (providerId, options) => { - if (!providerId || !stall) return originalGetAvailable(providerId, options); + if (providerId !== provider || !stall) return originalGetAvailable(providerId, options); stall = false; started.resolve(); await gate.promise; throw new Error("stale provider availability failure"); }; - const staleRefresh = runtime.getAvailable(harness.getModel().provider); + vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout"] }); + const staleRefresh = expect(runtime.getAvailable(provider)).rejects.toThrow(TIMEOUT_MESSAGE); await started.promise; - await runtime.getAvailable(); + await timeOutStalledPass(harness); + await staleRefresh; + expect(runtime.hasConfiguredAuth(provider)).toBe(true); expect(runtime.getError()).toBeUndefined(); gate.resolve(); - await expect(staleRefresh).rejects.toThrow("stale provider availability failure"); + await Promise.resolve(); expect(runtime.getError()).toBeUndefined(); }); });