diff --git a/apps/server/src/usage/usageAggregation.test.ts b/apps/server/src/usage/usageAggregation.test.ts index e7af899653aec..3819903606be2 100644 --- a/apps/server/src/usage/usageAggregation.test.ts +++ b/apps/server/src/usage/usageAggregation.test.ts @@ -292,20 +292,24 @@ describe("UsageAggregator", () => { describe("resolveModelAliases", () => { it("follows chains to the final model and drops chains that enter a loop", () => { - expect( - resolveModelAliases({ + expect([ + ...resolveModelAliases({ "preview[1m]": "preview", preview: "example-model", + shared: "preview", + intoLoop: "loop", loop: "back", back: "loop", - intoLoop: "loop", + laterIntoLoop: "back", + deepIntoLoop: "intoLoop", self: "self", + validAfterLoops: "shared", }), - ).toEqual( - new Map([ - ["preview[1m]", "example-model"], - ["preview", "example-model"], - ]), - ); + ]).toEqual([ + ["preview[1m]", "example-model"], + ["preview", "example-model"], + ["shared", "example-model"], + ["validAfterLoops", "example-model"], + ]); }); }); diff --git a/apps/server/src/usage/usageAggregation.ts b/apps/server/src/usage/usageAggregation.ts index 649133b00b67b..ce7a78a99d9c1 100644 --- a/apps/server/src/usage/usageAggregation.ts +++ b/apps/server/src/usage/usageAggregation.ts @@ -79,15 +79,26 @@ export interface AggregateOptions { export function resolveModelAliases( aliases: Readonly>, ): ReadonlyMap { - const resolved = new Map(); - for (const model of Object.keys(aliases)) { - const seen = new Set([model]); - let target = aliases[model]!; - while (Object.hasOwn(aliases, target) && !seen.has(target)) { - seen.add(target); + const models = Object.keys(aliases); + const memo = new Map(); + for (const model of models) { + if (memo.has(model)) continue; + const path: string[] = []; + let target = model; + while (Object.hasOwn(aliases, target) && !memo.has(target)) { + path.push(target); + // Null marks both the active walk and paths that end in a cycle. + memo.set(target, null); target = aliases[target]!; } - if (!Object.hasOwn(aliases, target)) resolved.set(model, target); + const terminal = Object.hasOwn(aliases, target) ? (memo.get(target) ?? null) : target; + for (const alias of path) memo.set(alias, terminal); + } + + const resolved = new Map(); + for (const model of models) { + const target = memo.get(model); + if (target !== null && target !== undefined) resolved.set(model, target); } return resolved; }