diff --git a/packages/server/src/server/plugins/usage-sources/index.test.ts b/packages/server/src/server/plugins/usage-sources/index.test.ts index 03155fe880..f0173691c8 100644 --- a/packages/server/src/server/plugins/usage-sources/index.test.ts +++ b/packages/server/src/server/plugins/usage-sources/index.test.ts @@ -725,3 +725,39 @@ test("attempt ordering keeps its timestamp while retained freshness records succ expect(failure?.fetchedAt).toBe(new Date(3000).toISOString()); expect(failure?.lastSuccessful?.fetchedAt).toBe(success?.successfulAt); }); + +test("failed discovery preserves history until discovery confirms account removal", async () => { + let discoveryFails = false; + let present = true; + let available = true; + const registry = new UsageSourceRegistry(() => 1000, 300_000, { warn: () => {} }); + registry.register( + source({ + id: "source", + discover: async () => { + if (discoveryFails) throw new Error("Login store temporarily unavailable"); + return present ? [{ key: "account", input: {} }] : []; + }, + fetch: async () => + available + ? { status: "available", windows: [{ id: "weekly", label: "Weekly", usedPct: 42 }] } + : { status: "error", error: "Rate limited" }, + }), + ); + const [success] = await registry.listReports(); + available = false; + discoveryFails = true; + const [duringDiscoveryFailure] = await registry.listReports({ forceRefresh: true }); + expect(duringDiscoveryFailure?.lastSuccessful).toEqual({ + fetchedAt: success?.successfulAt, + report: success?.report, + }); + discoveryFails = false; + const [afterDiscoveryRecovery] = await registry.listReports({ forceRefresh: true }); + expect(afterDiscoveryRecovery?.lastSuccessful).toEqual(duringDiscoveryFailure?.lastSuccessful); + expect(afterDiscoveryRecovery?.report).toEqual({ status: "error", error: "Rate limited" }); + present = false; + expect(await registry.listReports()).toEqual([]); + present = true; + expect((await registry.listReports())[0]?.lastSuccessful).toBeUndefined(); +}); diff --git a/packages/server/src/server/plugins/usage-sources/index.ts b/packages/server/src/server/plugins/usage-sources/index.ts index b0a4da3b66..e9e2e9c8af 100644 --- a/packages/server/src/server/plugins/usage-sources/index.ts +++ b/packages/server/src/server/plugins/usage-sources/index.ts @@ -99,7 +99,7 @@ export class UsageSourceRegistry { } else if (options.reportIds !== undefined) { ids = options.reportIds; } else { - this.defaults = await this.discover({ kind: "global" }); + this.defaults = await this.discover({ kind: "global" }, this.defaults); this.mergeKnown(); ids = [...this.known.keys()]; } @@ -191,7 +191,10 @@ export class UsageSourceRegistry { ); } - private async discover(scope: UsageScope): Promise> { + private async discover( + scope: UsageScope, + previous: ReadonlyMap = new Map(), + ): Promise> { const discovered = await Promise.all( [...this.sources.values()].map(async (source) => { const reports = new Map(); @@ -218,6 +221,11 @@ export class UsageSourceRegistry { } } catch (error) { this.logger.warn({ sourceId: source.id, err: error }, "Usage source discovery failed"); + // A failed store read does not confirm removal. Session discovery has no previous + // mapping because a resumed agent may use a different launch environment. + for (const [id, report] of previous) { + if (report.source === source) reports.set(id, report); + } } return reports; }),