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 cbfde562b5..bfa9c62de5 100644 --- a/packages/server/src/server/plugins/usage-sources/index.test.ts +++ b/packages/server/src/server/plugins/usage-sources/index.test.ts @@ -98,7 +98,7 @@ test("invalid keys become source errors and missing identities produce no report }), ); const reports = await registry.listReports(); - expect(reports.map((entry) => entry.id)).toEqual(["source:error"]); + expect(reports.map((entry) => entry.id)).toEqual(["source:!error"]); expect(reports[0]?.report.status).toBe("error"); expect( await registry.resolveReference({ source: "source", input: { account: "none" } }), @@ -125,13 +125,22 @@ test("two Codex homes and a token route resolve to their vendor account IDs", as const personal = await mkdtemp(join(tmpdir(), "usage-personal-")); const work = await mkdtemp(join(tmpdir(), "usage-work-")); try { + const token = (accountId: string, email: string) => + `header.${Buffer.from(JSON.stringify({ "https://api.openai.com/auth": { chatgpt_account_id: accountId }, email })).toString("base64url")}.signature`; await writeFile( join(personal, "auth.json"), - JSON.stringify({ tokens: { account_id: "personal-id", access_token: "personal-token" } }), + JSON.stringify({ + tokens: { + account_id: "personal-id", + access_token: token("personal-id", "personal@example.test"), + }, + }), ); await writeFile( join(work, "auth.json"), - JSON.stringify({ tokens: { account_id: "work-id", access_token: "work-token" } }), + JSON.stringify({ + tokens: { account_id: "work-id", access_token: token("work-id", "work@example.test") }, + }), ); const registry = new UsageSourceRegistry(); registry.register( @@ -141,13 +150,16 @@ test("two Codex homes and a token route resolve to their vendor account IDs", as identify: async (input) => identify(input as Parameters[0]), }), ); - const token = `header.${Buffer.from(JSON.stringify({ "https://api.openai.com/auth": { chatgpt_account_id: "work-id" } })).toString("base64url")}.signature`; + expect( - await registry.resolveReference({ source: "codex", input: { accessToken: token } }), + await registry.resolveReference({ + source: "codex", + input: { accessToken: token("work-id", "work@example.test") }, + }), ).toBe("codex:work-id"); - expect((await registry.listReports()).map((entry) => entry.id)).toEqual([ - "codex:personal-id", - "codex:work-id", + expect((await registry.listReports()).map((entry) => [entry.id, entry.account.label])).toEqual([ + ["codex:personal-id", "personal@example.test"], + ["codex:work-id", "work@example.test"], ]); } finally { await rm(personal, { recursive: true, force: true }); @@ -179,3 +191,51 @@ test("concurrent requests for the same ID share one vendor fetch", async () => { expect(one[0]).toBe(two[0]); expect(fetches).toBe(1); }); + +test("source failure IDs cannot collide with an account named error", async () => { + const registry = new UsageSourceRegistry(); + registry.register( + source({ + id: "source", + discover: async () => [{ account: "error" }, { account: "bad:key" }], + }), + ); + const reports = await registry.listReports(); + expect(reports).toHaveLength(2); + expect(reports.find((entry) => entry.id === "source:error")?.report.status).toBe("available"); + expect(reports.find((entry) => entry.report.status === "error")?.id).toMatch( + /^source:[^A-Za-z0-9._-]/, + ); +}); + +test("expired cached entries are pruned when a new report is written", async () => { + let now = 0; + const registry = new UsageSourceRegistry(() => now, 100); + registry.register(source({ id: "source" })); + await registry.resolveReference({ source: "source", input: { account: "old" } }); + await registry.listReports({ reportIds: ["source:old"] }); + now = 101; + await registry.resolveReference({ source: "source", input: { account: "new" } }); + await registry.listReports({ reportIds: ["source:new"] }); + const cache = Reflect.get(registry, "cache") as Map; + expect([...cache.keys()]).toEqual(["source:new"]); +}); + +test("discovery failures have an ID outside the account namespace", async () => { + const registry = new UsageSourceRegistry(); + registry.register( + source({ + id: "source", + discover: async () => { + throw new Error("discovery failed"); + }, + }), + ); + expect(await registry.resolveReference({ source: "source", input: { account: "error" } })).toBe( + "source:error", + ); + const discovered = await registry.listReports(); + expect(discovered.map((entry) => entry.id)).toEqual(["source:!error"]); + const both = await registry.listReports({ reportIds: ["source:error", "source:!error"] }); + expect(both.map((entry) => entry.report.status)).toEqual(["available", "error"]); +}); diff --git a/packages/server/src/server/plugins/usage-sources/index.ts b/packages/server/src/server/plugins/usage-sources/index.ts index 1734817104..a6705114de 100644 --- a/packages/server/src/server/plugins/usage-sources/index.ts +++ b/packages/server/src/server/plugins/usage-sources/index.ts @@ -64,8 +64,8 @@ export class UsageSourceRegistry { this.known.set(id, { source, input, label: account.label }); return id; } catch (error) { - const id = `${source.id}:error`; - this.cache.set(id, { at: this.now(), entry: this.errorEntry(source, id, error) }); + const id = `${source.id}:!error`; + this.writeCache(id, this.errorEntry(source, id, error)); return id; } } @@ -91,8 +91,8 @@ export class UsageSourceRegistry { (id): id is string => id !== null, ); } catch (error) { - const id = `${source.id}:error`; - this.cache.set(id, { at: this.now(), entry: this.errorEntry(source, id, error) }); + const id = `${source.id}:!error`; + this.writeCache(id, this.errorEntry(source, id, error)); return [id]; } }), @@ -148,7 +148,7 @@ export class UsageSourceRegistry { } catch (error) { entry = this.errorEntry(known.source, id, error, known.label); } - this.cache.set(id, { at: this.now(), entry }); + this.writeCache(id, entry); return entry; })(); this.pending.set(id, request); @@ -158,6 +158,14 @@ export class UsageSourceRegistry { return request; } + private writeCache(id: string, entry: UsageReportEntry): void { + const at = this.now(); + for (const [cachedId, cached] of this.cache) { + if (at - cached.at >= this.ttlMs) this.cache.delete(cachedId); + } + this.cache.set(id, { at, entry }); + } + private errorEntry( source: UsageSource, id: string,