From 9a2d2f8a29ff12cb3d0e6265282546f917097e3e Mon Sep 17 00:00:00 2001 From: Theo Browne Date: Sat, 3 Oct 2026 14:31:14 -0700 Subject: [PATCH 1/2] feat(usage): map a model ID to another model Co-Authored-By: Claude Opus 5.5 (1M context) --- .../src/environment/ServerEnvironment.ts | 1 + apps/server/src/usage/UsageService.ts | 18 +- .../server/src/usage/usageAggregation.test.ts | 40 ++- apps/server/src/usage/usageAggregation.ts | 34 +- .../components/usage/UsagePriceOverrides.tsx | 331 +++++++++++------- .../components/usage/usagePriceTable.test.ts | 59 +++- .../src/components/usage/usagePriceTable.ts | 38 ++ .../usage/usagePriceTargets.test.ts | 32 ++ .../src/components/usage/usagePriceTargets.ts | 34 +- docs/user/usage.md | 5 + packages/client-runtime/src/state/server.ts | 15 +- packages/contracts/src/environment.ts | 2 + packages/contracts/src/settings.ts | 11 + packages/shared/src/serverSettings.ts | 9 + 14 files changed, 469 insertions(+), 160 deletions(-) diff --git a/apps/server/src/environment/ServerEnvironment.ts b/apps/server/src/environment/ServerEnvironment.ts index 5381645e9fcf..16a4a93b3f20 100644 --- a/apps/server/src/environment/ServerEnvironment.ts +++ b/apps/server/src/environment/ServerEnvironment.ts @@ -233,6 +233,7 @@ export const make = Effect.gen(function* () { environmentThemes: true, usageLimitSources: true, usagePriceOverrides: true, + usageModelAliases: true, threadPinning: true, threadPinReorder: true, threadActiveReorder: true, diff --git a/apps/server/src/usage/UsageService.ts b/apps/server/src/usage/UsageService.ts index dc3121fc0024..ca56c36028bb 100644 --- a/apps/server/src/usage/UsageService.ts +++ b/apps/server/src/usage/UsageService.ts @@ -53,7 +53,7 @@ import { mergeProviderInstanceEnvironment } from "../provider/ProviderInstanceEn import { readOpenCodeUsage } from "./opencodeUsageReader.ts"; import { readAntigravityUsage } from "./antigravityUsageReader.ts"; import { readCursorAccountUsage } from "./cursorUsageReader.ts"; -import { UsageAggregator } from "./usageAggregation.ts"; +import { resolveModelAliases, UsageAggregator } from "./usageAggregation.ts"; import { createOverrideRateTable, parseRateTable, type RateTable } from "./usagePricing.ts"; import { listTranscriptFiles, @@ -821,6 +821,7 @@ export const make = Effect.gen(function* () { ...hourlyWindow, rates, priceOverrides: createOverrideRateTable(settings.usagePriceOverrides), + modelAliases: resolveModelAliases(settings.usageModelAliases), }); const sources: UsageSource[] = []; @@ -925,17 +926,13 @@ export const make = Effect.gen(function* () { }); /** - * In-flight scans by window and custom prices, so concurrent identical requests (the usage + * In-flight scans by window and usage settings, so concurrent identical requests (the usage * page open on two clients at once) share one scan instead of racing over * the same corpus twice. */ const inflightScans = new Map>(); - const scanKey = ( - input: UsageSummaryInput, - priceOverrides: ServerSettingsValue["usagePriceOverrides"], - cursorKeychainUsageEnabled: boolean, - ): string => + const scanKey = (input: UsageSummaryInput, settings: ServerSettingsValue): string => JSON.stringify([ input.timeZone, input.sinceDay, @@ -943,13 +940,14 @@ export const make = Effect.gen(function* () { input.resolution ?? "day", input.sinceTime ?? null, input.untilTime ?? null, - priceOverrides, - cursorKeychainUsageEnabled, + settings.usagePriceOverrides, + settings.usageModelAliases, + settings.cursorKeychainUsageEnabled, ]); const readSummary = Effect.fn("UsageService.readSummary")(function* (input: UsageSummaryInput) { const settings = yield* readSettings; - const key = scanKey(input, settings.usagePriceOverrides, settings.cursorKeychainUsageEnabled); + const key = scanKey(input, settings); const deferred = yield* Effect.uninterruptible( Effect.gen(function* () { const existing = inflightScans.get(key); diff --git a/apps/server/src/usage/usageAggregation.test.ts b/apps/server/src/usage/usageAggregation.test.ts index 9fea9cd96644..375ca85f95b4 100644 --- a/apps/server/src/usage/usageAggregation.test.ts +++ b/apps/server/src/usage/usageAggregation.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it } from "@effect/vitest"; -import { UsageAggregator } from "./usageAggregation.ts"; +import { resolveModelAliases, UsageAggregator } from "./usageAggregation.ts"; import type { RateTable } from "./usagePricing.ts"; import type { UsageRecord } from "./usageTranscripts.ts"; @@ -253,6 +253,25 @@ describe("UsageAggregator", () => { expect(aggregator.add(record({ timestampMs: Date.parse("2026-07-01T12:00:00Z") }))).toBe(false); }); + it("folds a mapped model into its target and prices it there", () => { + const aggregator = new UsageAggregator({ + timeZone: "UTC", + sinceDay: "2026-08-01", + untilDay: "2026-08-31", + rates, + modelAliases: resolveModelAliases({ "example-preview": "claude-fable-5" }), + }); + aggregator.add(record()); + aggregator.add(record({ model: "example-preview", rateModel: "example-preview-high" })); + const result = aggregator.finish(); + + expect(result.buckets).toHaveLength(1); + expect(result.buckets[0]?.model).toBe("claude-fable-5"); + expect(result.buckets[0]?.records).toBe(2); + expect(result.buckets[0]?.costUsd).toBeCloseTo(0.00925, 9); + expect(result.buckets[0]?.unpricedRecords).toBe(0); + }); + it("separates providers and models into their own buckets", () => { const result = aggregate([ record(), @@ -263,3 +282,22 @@ describe("UsageAggregator", () => { expect(result.buckets).toHaveLength(3); }); }); + +describe("resolveModelAliases", () => { + it("follows chains to the final model and drops loops", () => { + expect( + resolveModelAliases({ + "preview[1m]": "preview", + preview: "example-model", + loop: "back", + back: "loop", + self: "self", + }), + ).toEqual( + new Map([ + ["preview[1m]", "example-model"], + ["preview", "example-model"], + ]), + ); + }); +}); diff --git a/apps/server/src/usage/usageAggregation.ts b/apps/server/src/usage/usageAggregation.ts index 0652ca4b7dcf..6a3cca3e9af6 100644 --- a/apps/server/src/usage/usageAggregation.ts +++ b/apps/server/src/usage/usageAggregation.ts @@ -89,11 +89,34 @@ export interface AggregateOptions { readonly untilDay: string; readonly rates: RateTable; readonly priceOverrides?: RateTable; + /** From {@link resolveModelAliases}. Mapped records bucket and price as their target. */ + readonly modelAliases?: ReadonlyMap; readonly resolution?: UsageResolution; readonly sinceTimeMs?: number; readonly untilTimeMs?: number; } +/** + * Resolves user model mappings to their final target, so `a -> b` and + * `b -> c` both land on `c`. A model whose chain loops back to it is left + * unmapped. + */ +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); + target = aliases[target]!; + } + if (target !== model) resolved.set(model, target); + } + return resolved; +} + export interface AggregateResult { readonly buckets: readonly UsageBucket[]; /** Records dropped because an earlier record carried the same dedupe key. */ @@ -147,7 +170,8 @@ export class UsageAggregator { * can derive per-window facts (distinct sessions, for one) from the records * that landed rather than everything the mtime prefilter happened to admit. */ - add(record: UsageRecord, sourcePath?: string): boolean { + add(input: UsageRecord, sourcePath?: string): boolean { + const record = this.#mapModel(input); if (record.dedupeKey !== null) { if (this.#seen.has(record.dedupeKey)) { this.#duplicatesDropped += 1; @@ -217,6 +241,14 @@ export class UsageAggregator { return true; } + /** The target's own rate applies, so a provider-specific `rateModel` is dropped. */ + #mapModel(record: UsageRecord): UsageRecord { + const model = this.#options.modelAliases?.get(record.model); + if (model === undefined) return record; + const { rateModel: _rateModel, ...rest } = record; + return { ...rest, model }; + } + /** * Records arrive file by file in time order, so most land in the bucket the * previous record used. Checking that first skips building and hashing a diff --git a/apps/web/src/components/usage/UsagePriceOverrides.tsx b/apps/web/src/components/usage/UsagePriceOverrides.tsx index c9af1c0e1770..0b52687c81b0 100644 --- a/apps/web/src/components/usage/UsagePriceOverrides.tsx +++ b/apps/web/src/components/usage/UsagePriceOverrides.tsx @@ -34,6 +34,7 @@ import { Tooltip, TooltipTrigger, TooltipPopup } from "../ui/tooltip"; import { USAGE_PRICE_FIELDS } from "./usagePriceForm"; import { isEmptyUsagePriceDraft, + usageAliasCell, usagePriceCell, usagePriceTableChanges, usagePriceTableErrors, @@ -63,6 +64,10 @@ const priceTargetsAtom = Atom.make((get): readonly UsagePriceTarget[] => environmentId, label: environment.entry.target.label, prices: settings?.usagePriceOverrides ?? null, + aliases: + environment.serverConfig?.environment.capabilities.usageModelAliases === true + ? (settings?.usageModelAliases ?? null) + : null, unavailable: environment.connection.phase !== "connected" ? "Offline" @@ -91,7 +96,10 @@ type SaveAttempt = { >; }; -/** Edits custom model prices. `initialModel` opens with a new row for that model. */ +/** + * Edits custom model prices and mappings. A mapped model's usage counts as its + * target model. `initialModel` opens with a new row for that model. + */ export function UsagePriceOverrides({ usage, initialSelectedEnvironmentIds, @@ -108,12 +116,13 @@ export function UsagePriceOverrides({ const selected = environments.filter( (environment) => selectedIds === null || selectedIds.has(environment.environmentId), ); - // A model that already has a custom price somewhere is edited in its existing row. + // A model that already has a custom price or mapping is edited in its existing row. const [drafts, setDrafts] = useState(() => initialModel === undefined || selected.some( (environment) => - environment.prices !== null && Object.hasOwn(environment.prices, initialModel), + Object.hasOwn(environment.prices ?? {}, initialModel) || + Object.hasOwn(environment.aliases ?? {}, initialModel), ) ? [] : [{ id: "new:initial", model: initialModel, isNew: true, values: {} }], @@ -124,7 +133,12 @@ export function UsagePriceOverrides({ const nextRowId = useRef(0); const updateSettings = useAtomCommand(serverEnvironment.updateSettings, { reportFailure: false }); const customModels = [ - ...new Set(selected.flatMap((environment) => Object.keys(environment.prices ?? {}))), + ...new Set( + selected.flatMap((environment) => [ + ...Object.keys(environment.prices ?? {}), + ...Object.keys(environment.aliases ?? {}), + ]), + ), ].sort(); const models = [ ...new Set([ @@ -191,19 +205,31 @@ export function UsagePriceOverrides({ ); setAttempt(null); }; + // An existing row edited back to its saved values has nothing left to save. + const matchesSaved = ( + row: UsagePriceDraft, + cell: { value: string; placeholder: string }, + value: string, + ) => + !row.isNew && + cell.placeholder !== "Mixed" && + cell.placeholder !== "Unavailable" && + value === cell.value; + const updateOrDropDraft = (draft: UsagePriceDraft) => { + if (!draft.isNew && Object.keys(draft.values).length === 0 && draft.alias === undefined) + setDrafts((previous) => previous.filter((entry) => entry.id !== draft.id)); + else updateDraft(draft); + }; const editCell = (row: UsagePriceDraft, field: UsagePriceField, value: string) => { const values = { ...row.values, [field]: value }; const original = usagePriceCell(selected, row.model, field); - if ( - !row.isNew && - original.placeholder !== "Mixed" && - original.placeholder !== "Unavailable" && - value === original.value - ) - delete values[field]; - if (!row.isNew && Object.keys(values).length === 0) - setDrafts((previous) => previous.filter((entry) => entry.id !== row.id)); - else updateDraft({ ...row, values }); + if (matchesSaved(row, original, value)) delete values[field]; + updateOrDropDraft({ ...row, values }); + }; + const editAlias = (row: UsagePriceDraft, value: string) => { + const { alias: _alias, ...rest } = row; + const original = usageAliasCell(selected, row.model); + updateOrDropDraft(matchesSaved(row, original, value) ? rest : { ...rest, alias: value }); }; const save = async (retry = false) => { if (pending) return; @@ -214,7 +240,7 @@ export function UsagePriceOverrides({ (destination) => environments.find( (environment) => environment.environmentId === destination.environmentId, - ) ?? { ...destination, prices: null, unavailable: "Environment removed" }, + ) ?? { ...destination, prices: null, aliases: null, unavailable: "Environment removed" }, ); const changes = new Map( targets.map((target) => [target.environmentId, usagePriceTableChanges(target, edits)]), @@ -257,11 +283,11 @@ export function UsagePriceOverrides({ if (!pending) onOpenChange(open); }} > - + Custom model prices - Prices apply to all past and future usage on the environments you select. + Prices and mappings apply to all past and future usage on the environments you select. @@ -326,9 +352,10 @@ export function UsagePriceOverrides({ ) : ( <>
- +
- + + {USAGE_PRICE_FIELDS.map((field) => ( ))} @@ -337,6 +364,7 @@ export function UsagePriceOverrides({ Model ID + Map to {USAGE_PRICE_FIELDS.map((field) => ( {field.label} ))} @@ -360,132 +388,157 @@ export function UsagePriceOverrides({ {rows.length === 0 ? ( - +

{selected.some((environment) => environment.prices === null) ? "Some environment prices are unavailable." - : "No custom prices. Add a row to override automatic pricing."} + : "No custom prices or mappings. Add a row to set one."}

) : ( - rows.map((row) => ( - - - {row.isNew ? ( - { - if (node && focusRowRef.current === row.id) { - node.focus(); - focusRowRef.current = null; + rows.map((row) => { + const aliasCell = usageAliasCell(selected, row.model); + const alias = (row.alias ?? aliasCell.value).trim(); + return ( + + + {row.isNew ? ( + { + if (node && focusRowRef.current === row.id) { + node.focus(); + focusRowRef.current = null; + } + }} + aria-label="New model ID" + aria-invalid={ + (row.model.trim() !== "" && errors.has(row.id)) || undefined + } + list="usage-price-models" + placeholder="Model ID" + autoComplete="off" + spellCheck={false} + disabled={locked} + onChange={(event) => + updateDraft({ ...row, model: event.target.value }) } - }} - aria-label="New model ID" - aria-invalid={ - (row.model.trim() !== "" && errors.has(row.id)) || undefined - } - list="usage-price-models" - placeholder="Model ID" - autoComplete="off" - spellCheck={false} - disabled={locked} - onChange={(event) => - updateDraft({ ...row, model: event.target.value }) - } - /> + /> + ) : ( + + {row.model} + + )} + {errors.has(row.id) && + (row.model.trim() !== "" || + Object.values(row.values).some((value) => value !== "")) ? ( +

+ {errors.get(row.id)} +

+ ) : null} + + {row.removed ? ( + + + Automatic pricing after saving + + ) : ( - - {row.model} - + + editAlias(row, event.target.value)} + /> + )} - {errors.has(row.id) && - (row.model.trim() !== "" || - Object.values(row.values).some((value) => value !== "")) ? ( -

- {errors.get(row.id)} -

- ) : null} - - {row.removed ? ( - - - Automatic pricing after saving - - - ) : ( - USAGE_PRICE_FIELDS.map((field) => { - const cell = usagePriceCell(selected, row.model, field.key); - return ( - - - editCell(row, field.key, event.target.value) - } - /> - - ); - }) - )} - - - } - disabled={locked} - aria-label={ - row.removed - ? `Undo reset for ${row.model}` - : row.isNew - ? "Remove new model" - : `Reset price for ${row.model} to automatic` - } - onClick={() => { - if (row.isNew) - setDrafts((previous) => - previous.filter((entry) => entry.id !== row.id), - ); - else if (row.removed) { - if (Object.keys(row.values).length === 0) + {row.removed ? null : alias !== "" ? ( + + + Counted as {alias} + + + ) : ( + USAGE_PRICE_FIELDS.map((field) => { + const cell = usagePriceCell(selected, row.model, field.key); + return ( + + + editCell(row, field.key, event.target.value) + } + /> + + ); + }) + )} + + + } + disabled={locked} + aria-label={ + row.removed + ? `Undo reset for ${row.model}` + : row.isNew + ? "Remove new model" + : `Reset price for ${row.model} to automatic` + } + onClick={() => { + if (row.isNew) setDrafts((previous) => previous.filter((entry) => entry.id !== row.id), ); - else updateDraft({ ...row, removed: false }); - } else updateDraft({ ...row, removed: true }); - }} - > - {row.isNew ? : } - - - {row.removed - ? "Undo reset" - : row.isNew - ? "Remove row" - : "Reset to automatic"} - - - - - )) + else if (row.removed) + updateOrDropDraft({ ...row, removed: false }); + else updateDraft({ ...row, removed: true }); + }} + > + {row.isNew ? ( + + ) : ( + + )} + + + {row.removed + ? "Undo reset" + : row.isNew + ? "Remove row" + : "Reset to automatic"} + + + + + ); + }) )}
@@ -497,8 +550,20 @@ export function UsagePriceOverrides({