diff --git a/.changeset/usage-budget.md b/.changeset/usage-budget.md new file mode 100644 index 00000000..61813243 --- /dev/null +++ b/.changeset/usage-budget.md @@ -0,0 +1,7 @@ +--- +"@open-codesign/desktop": minor +"@open-codesign/i18n": patch +"@open-codesign/shared": patch +--- + +Show recorded token and USD totals for the current design, including today and the local week, beside the existing per-run token count. diff --git a/apps/desktop/src/main/ipc/generate.ts b/apps/desktop/src/main/ipc/generate.ts index cb976e63..4b497b3d 100644 --- a/apps/desktop/src/main/ipc/generate.ts +++ b/apps/desktop/src/main/ipc/generate.ts @@ -29,6 +29,8 @@ import { deriveResourceStateFromChatRows, GeneratePayloadV1, ListActiveMessagesInputV1, + summarizeUsageBudget, + type UsageTotals, } from '@open-codesign/shared'; import { computeFingerprint } from '@open-codesign/shared/fingerprint'; import type { BrowserWindow as ElectronBrowserWindow, WebContents } from 'electron'; @@ -58,7 +60,6 @@ import { workspaceNameFromPath, } from '../memory-ipc'; import { getApiKeyForProvider, getCachedConfig, hasApiKeyForProvider } from '../onboarding-ipc'; - import { readPersisted as readPreferences } from '../preferences-ipc'; import { runPreview } from '../preview-runtime'; import { preparePromptContext } from '../prompt-context'; @@ -510,6 +511,32 @@ export function registerGenerateIpc({ db, getMainWindow }: RegisterGenerateIpcDe return recovered; }); + const emptyUsage = (): UsageTotals => ({ inputTokens: 0, outputTokens: 0, costUsd: 0 }); + ipcMain.handle('codesign:v1:usage-budget', async (_event, raw: unknown) => { + const input = raw as { schemaVersion?: unknown; designId?: unknown } | null; + if ( + !input || + input.schemaVersion !== 1 || + typeof input.designId !== 'string' || + input.designId.length === 0 || + input.designId.length > 512 + ) { + throw new CodesignError('Invalid usage budget request', 'IPC_BAD_INPUT'); + } + const empty = { + schemaVersion: 1 as const, + design: emptyUsage(), + today: emptyUsage(), + week: emptyUsage(), + }; + if (!journal) return empty; + const records = await journal.usageRecords(); + return { + schemaVersion: 1 as const, + ...summarizeUsageBudget(records, input.designId, Date.now()), + }; + }); + const recordFinalError = (scope: string, runId: string, err: unknown): void => { if (db === null) return; const code = err instanceof CodesignError ? (err.code as string) : 'PROVIDER_UPSTREAM_ERROR'; diff --git a/apps/desktop/src/main/run-journal.test.ts b/apps/desktop/src/main/run-journal.test.ts index 81ecce21..02b36009 100644 --- a/apps/desktop/src/main/run-journal.test.ts +++ b/apps/desktop/src/main/run-journal.test.ts @@ -325,4 +325,26 @@ describe('RunJournal', () => { expect.objectContaining({ type: 'run_settled', seq: 1, outcome: 'interrupted' }), ]); }); + + it('sums settled provider usage after the journal is reopened', async () => { + const journal = new RunJournal(directory); + await start(journal, 'run-a', 'design-a', 10); + const response: GenerateResponse = { + message: 'done', + artifacts: [], + inputTokens: 12, + outputTokens: 8, + costUsd: 0.03, + }; + await journal.append( + event('run-a', 'design-a', { type: 'run_settled', outcome: 'completed', response }), + ); + await start(journal, 'run-b', 'design-b', 11); + await journal.append(event('run-b', 'design-b', { type: 'run_settled', outcome: 'failed' })); + + const reopened = new RunJournal(directory); + await expect(reopened.usageRecords()).resolves.toEqual([ + { designId: 'design-a', startedAt: 10, inputTokens: 12, outputTokens: 8, costUsd: 0.03 }, + ]); + }); }); diff --git a/apps/desktop/src/main/run-journal.ts b/apps/desktop/src/main/run-journal.ts index ad99a895..ca689ae2 100644 --- a/apps/desktop/src/main/run-journal.ts +++ b/apps/desktop/src/main/run-journal.ts @@ -16,11 +16,18 @@ export interface RunJournalRecovery { type JournalEvent = AgentStreamEvent & { schemaVersion: 1; runId: string; seq: number }; +interface RunUsage { + inputTokens: number; + outputTokens: number; + costUsd: number; +} + interface RunState extends RunJournalStart { seq: number; bytes: number; settled: boolean; writeError?: Error; + usage?: RunUsage; } const EVENT_TYPES = new Set([ @@ -89,6 +96,20 @@ function validateEvent(value: unknown, run: RunState, seq: number): JournalEvent return record as unknown as JournalEvent; } +function finiteUsage(value: unknown): number { + return typeof value === 'number' && Number.isFinite(value) && value >= 0 ? value : 0; +} + +function usageFromResponse(response: unknown): RunUsage | undefined { + if (response === null || typeof response !== 'object') return undefined; + const record = response as Record; + return { + inputTokens: finiteUsage(record['inputTokens']), + outputTokens: finiteUsage(record['outputTokens']), + costUsd: finiteUsage(record['costUsd']), + }; +} + function encode(value: unknown): Buffer { const data = Buffer.from(`${JSON.stringify(value)}\n`, 'utf8'); if (data.length > MAX_RECORD_BYTES) throw invalid('record exceeds the 64 MiB size limit'); @@ -173,6 +194,25 @@ export class RunJournal { }); } + /** Settled runs that recorded a provider response. Totals are derived from these rows. */ + usageRecords(): Promise> { + return this.serialize(async () => { + await this.initialize(); + const records: Array<{ designId: string; startedAt: number } & RunUsage> = []; + for (const run of this.sortedRuns()) { + if (!run.usage) continue; + records.push({ + designId: run.designId, + startedAt: run.startedAt, + inputTokens: run.usage.inputTokens, + outputTokens: run.usage.outputTokens, + costUsd: run.usage.costUsd, + }); + } + return records; + }); + } + /** Startup projection repair; history is streamed from disk and never retained by this instance. */ allEvents(): Promise { return this.serialize(async () => { @@ -278,6 +318,10 @@ export class RunJournal { run.seq = event.seq; run.bytes += bytes.length; run.settled = event.type === 'run_settled'; + if (event.type === 'run_settled') { + const usage = usageFromResponse(event.response); + if (usage) run.usage = usage; + } return committed; } @@ -319,6 +363,10 @@ export class RunJournal { const event = validateEvent(record, run, run.seq + 1); run.seq = event.seq; run.settled = event.type === 'run_settled'; + if (event.type === 'run_settled') { + const usage = usageFromResponse(event.response); + if (usage) run.usage = usage; + } if (events && event.seq > after) events.push(event); } run.bytes += newline + 1; diff --git a/apps/desktop/src/preload/index.ts b/apps/desktop/src/preload/index.ts index 7feedfce..6ecfbf78 100644 --- a/apps/desktop/src/preload/index.ts +++ b/apps/desktop/src/preload/index.ts @@ -439,6 +439,19 @@ export interface RunRecoveryResult { events: AgentStreamEvent[]; } +export interface UsageBudgetTotals { + inputTokens: number; + outputTokens: number; + costUsd: number; +} + +export interface UsageBudgetResult { + schemaVersion: 1; + design: UsageBudgetTotals; + today: UsageBudgetTotals; + week: UsageBudgetTotals; +} + const api = { detectProvider: (key: string) => ipcRenderer.invoke('codesign:detect-provider', key) as Promise, @@ -471,6 +484,11 @@ const api = { schemaVersion: 1, cursors, }) as Promise, + usageBudget: (designId: string) => + ipcRenderer.invoke('codesign:v1:usage-budget', { + schemaVersion: 1, + designId, + }) as Promise, generationStatus: () => ipcRenderer.invoke('codesign:v1:generation-status') as Promise, sendActiveMessage: (payload: ActiveRunMessageInputV1) => diff --git a/apps/desktop/src/renderer/src/components/Sidebar.tsx b/apps/desktop/src/renderer/src/components/Sidebar.tsx index a079d54d..0ce9e2d6 100644 --- a/apps/desktop/src/renderer/src/components/Sidebar.tsx +++ b/apps/desktop/src/renderer/src/components/Sidebar.tsx @@ -1,5 +1,10 @@ import { getCurrentLocale, useT, useTranslation } from '@open-codesign/i18n'; -import type { LocalInputFile, OnboardingState } from '@open-codesign/shared'; +import { + formatUsageCost, + formatUsageTokens, + type LocalInputFile, + type OnboardingState, +} from '@open-codesign/shared'; import { FolderOpen, Link2, Paperclip, X } from 'lucide-react'; import { useCallback, useEffect, useRef } from 'react'; import { useCodesignStore } from '../store'; @@ -10,6 +15,7 @@ import { CommentChipBar } from './chat/CommentChipBar'; import { EmptyState } from './chat/EmptyState'; import { PromptInput, type PromptInputHandle } from './chat/PromptInput'; import { ModelSwitcher } from './ModelSwitcher'; +import { useDesignUsageBudget } from './usage-budget-refresh'; export interface SidebarProps { prefillPrompt: { id: number; text: string } | null; @@ -90,6 +96,19 @@ export function Sidebar({ prefillPrompt }: SidebarProps) { const pickDesignSystemDirectory = useCodesignStore((s) => s.pickDesignSystemDirectory); const clearDesignSystem = useCodesignStore((s) => s.clearDesignSystem); const lastUsage = useCodesignStore((s) => s.lastUsage); + const pushToast = useCodesignStore((s) => s.pushToast); + const reportUsageError = useCallback( + (error: unknown) => { + const description = error instanceof Error ? error.message : String(error); + pushToast({ + variant: 'error', + title: t('sidebar.chat.usageLoadFailed'), + description, + }); + }, + [pushToast, t], + ); + const usageBudget = useDesignUsageBudget(reportUsageError); const chatMessages = useCodesignStore((s) => s.chatMessages); const chatLoaded = useCodesignStore((s) => s.chatLoaded); @@ -304,6 +323,34 @@ export function Sidebar({ prefillPrompt }: SidebarProps) { ) : null} + {usageBudget ? ( +

+ + {t('sidebar.chat.usageDesign', { + tokens: formatUsageTokens( + usageBudget.design.inputTokens + usageBudget.design.outputTokens, + ), + cost: formatUsageCost(usageBudget.design.costUsd), + })} + + + {t('sidebar.chat.usageWindow', { + todayTokens: formatUsageTokens( + usageBudget.today.inputTokens + usageBudget.today.outputTokens, + ), + todayCost: formatUsageCost(usageBudget.today.costUsd), + weekTokens: formatUsageTokens( + usageBudget.week.inputTokens + usageBudget.week.outputTokens, + ), + weekCost: formatUsageCost(usageBudget.week.costUsd), + })} + +

+ ) : null} ); diff --git a/apps/desktop/src/renderer/src/components/usage-budget-refresh.test.tsx b/apps/desktop/src/renderer/src/components/usage-budget-refresh.test.tsx new file mode 100644 index 00000000..62e027cb --- /dev/null +++ b/apps/desktop/src/renderer/src/components/usage-budget-refresh.test.tsx @@ -0,0 +1,201 @@ +// @vitest-environment happy-dom + +import { act } from 'react'; +import { createRoot, type Root } from 'react-dom/client'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import type { UsageBudgetResult } from '../../../preload/index'; +import { useCodesignStore } from '../store'; +import { useDesignUsageBudget } from './usage-budget-refresh'; + +const initialDesignId = useCodesignStore.getState().currentDesignId; +const initialLastUsage = useCodesignStore.getState().lastUsage; +const initialSettled = useCodesignStore.getState().settledGenerationIds; + +function budget(todayTokens: number): UsageBudgetResult { + const totals = { inputTokens: todayTokens, outputTokens: 0, costUsd: 0 }; + return { schemaVersion: 1, design: totals, today: totals, week: totals }; +} + +function deferred() { + let resolve!: (value: T) => void; + const promise = new Promise((done) => { + resolve = done; + }); + return { promise, resolve }; +} + +function Probe({ onError }: { onError: (error: unknown) => void }) { + const result = useDesignUsageBudget(onError); + return {result ? String(result.today.inputTokens) : ''}; +} + +describe('useDesignUsageBudget', () => { + let container: HTMLDivElement; + let root: Root; + let visibility: DocumentVisibilityState; + + beforeEach(() => { + vi.stubGlobal('IS_REACT_ACT_ENVIRONMENT', true); + vi.useFakeTimers(); + vi.setSystemTime(new Date(2024, 0, 6, 23, 0, 0)); + visibility = 'visible'; + Object.defineProperty(document, 'visibilityState', { + configurable: true, + get: () => visibility, + }); + container = document.createElement('div'); + document.body.append(container); + root = createRoot(container); + useCodesignStore.setState({ + currentDesignId: 'design-a', + lastUsage: null, + settledGenerationIds: new Set(), + }); + }); + + afterEach(() => { + act(() => root.unmount()); + container.remove(); + useCodesignStore.setState({ + currentDesignId: initialDesignId, + lastUsage: initialLastUsage, + settledGenerationIds: initialSettled, + }); + Reflect.deleteProperty(window, 'codesign'); + vi.useRealTimers(); + vi.unstubAllGlobals(); + }); + + async function flush(): Promise { + await act(async () => { + await Promise.resolve(); + }); + } + + function shown(): string { + return container.querySelector('output')?.textContent ?? ''; + } + + it('refetches after a cancelled run settles without changing lastUsage', async () => { + const usageBudget = vi.fn().mockResolvedValueOnce(budget(1)).mockResolvedValueOnce(budget(4)); + Object.defineProperty(window, 'codesign', { configurable: true, value: { usageBudget } }); + + await act(async () => { + root.render(); + }); + await flush(); + expect(shown()).toBe('1'); + + await act(async () => { + useCodesignStore.setState((state) => ({ + settledGenerationIds: new Set([...state.settledGenerationIds, 'gen-cancel']), + })); + }); + await flush(); + + expect(useCodesignStore.getState().lastUsage).toBeNull(); + expect(usageBudget).toHaveBeenCalledTimes(2); + expect(usageBudget).toHaveBeenNthCalledWith(2, 'design-a'); + expect(shown()).toBe('4'); + }); + + it('refetches at the next local midnight and the Monday that starts the week', async () => { + const usageBudget = vi + .fn() + .mockResolvedValueOnce(budget(1)) + .mockResolvedValueOnce(budget(2)) + .mockResolvedValueOnce(budget(3)); + Object.defineProperty(window, 'codesign', { configurable: true, value: { usageBudget } }); + await act(async () => { + root.render(); + }); + await flush(); + expect(shown()).toBe('1'); + + await act(async () => { + await vi.advanceTimersByTimeAsync(60 * 60 * 1000); + }); + await flush(); + expect(new Date(Date.now()).getDay()).toBe(0); + expect(shown()).toBe('2'); + + await act(async () => { + await vi.advanceTimersByTimeAsync(24 * 60 * 60 * 1000); + }); + await flush(); + expect(new Date(Date.now()).getDay()).toBe(1); + expect(useCodesignStore.getState().lastUsage).toBeNull(); + expect(usageBudget).toHaveBeenCalledTimes(3); + expect(shown()).toBe('3'); + }); + + it('refetches when the window is focused or shown again', async () => { + const usageBudget = vi.fn(async () => budget(1)); + Object.defineProperty(window, 'codesign', { configurable: true, value: { usageBudget } }); + await act(async () => { + root.render(); + }); + await flush(); + expect(usageBudget).toHaveBeenCalledTimes(1); + + visibility = 'hidden'; + await act(async () => { + document.dispatchEvent(new Event('visibilitychange')); + }); + await flush(); + expect(usageBudget).toHaveBeenCalledTimes(1); + + visibility = 'visible'; + await act(async () => { + document.dispatchEvent(new Event('visibilitychange')); + window.dispatchEvent(new Event('focus')); + }); + await flush(); + expect(usageBudget).toHaveBeenCalledTimes(3); + }); + + it('ignores a response that resolves after cleanup or after a newer read', async () => { + const first = deferred(); + const second = deferred(); + const usageBudget = vi + .fn() + .mockReturnValueOnce(first.promise) + .mockReturnValueOnce(second.promise); + Object.defineProperty(window, 'codesign', { configurable: true, value: { usageBudget } }); + await act(async () => { + root.render(); + }); + + await act(async () => { + window.dispatchEvent(new Event('focus')); + }); + await act(async () => { + second.resolve(budget(2)); + await second.promise; + }); + await flush(); + expect(shown()).toBe('2'); + + await act(async () => { + first.resolve(budget(1)); + await first.promise; + }); + await flush(); + expect(shown()).toBe('2'); + + const late = deferred(); + usageBudget.mockReturnValueOnce(late.promise); + await act(async () => { + useCodesignStore.setState({ currentDesignId: 'design-b' }); + }); + await act(async () => { + useCodesignStore.setState({ currentDesignId: null }); + }); + await act(async () => { + late.resolve(budget(9)); + await late.promise; + }); + await flush(); + expect(shown()).toBe(''); + }); +}); diff --git a/apps/desktop/src/renderer/src/components/usage-budget-refresh.ts b/apps/desktop/src/renderer/src/components/usage-budget-refresh.ts new file mode 100644 index 00000000..d861d3ef --- /dev/null +++ b/apps/desktop/src/renderer/src/components/usage-budget-refresh.ts @@ -0,0 +1,83 @@ +import { nextLocalMidnight } from '@open-codesign/shared'; +import { useEffect, useRef, useState } from 'react'; +import type { UsageBudgetResult } from '../../../preload/index'; +import { useCodesignStore } from '../store'; + +/** + * Reloads the current design's journal totals. + * A cancelled run is removed before lastUsage changes, so settlement is a + * separate refresh signal. The same subscription also reloads at the next + * local midnight and when the window is focused or shown again. + */ +export function useDesignUsageBudget(onError: (error: unknown) => void): UsageBudgetResult | null { + const designId = useCodesignStore((state) => state.currentDesignId); + const lastUsage = useCodesignStore((state) => state.lastUsage); + const settledGenerationIds = useCodesignStore((state) => state.settledGenerationIds); + const [budget, setBudget] = useState(null); + const reportedError = useRef(null); + const shownDesignId = useRef(null); + + // Cancel records the settlement before lastUsage can change, and midnight + // does not change either value. Both are refresh signals for this effect. + // biome-ignore lint/correctness/useExhaustiveDependencies: refresh signals are not read by the request + useEffect(() => { + const readUsage = window.codesign?.usageBudget; + if (!designId || typeof readUsage !== 'function') { + setBudget(null); + return; + } + if (shownDesignId.current !== designId) { + shownDesignId.current = designId; + setBudget(null); + } + + let active = true; + let request = 0; + const load = () => { + const current = ++request; + void readUsage(designId).then( + (next) => { + if (!active || current !== request) return; + reportedError.current = null; + setBudget(next); + }, + (error: unknown) => { + if (!active || current !== request) return; + setBudget(null); + const description = error instanceof Error ? error.message : String(error); + if (reportedError.current === description) return; + reportedError.current = description; + onError(error); + }, + ); + }; + + let timer: ReturnType | undefined; + const arm = () => { + const now = Date.now(); + timer = setTimeout(() => { + load(); + arm(); + }, nextLocalMidnight(now) - now); + }; + + const onFocus = () => load(); + const onVisibility = () => { + if (document.visibilityState === 'hidden') return; + load(); + }; + + load(); + arm(); + window.addEventListener('focus', onFocus); + document.addEventListener('visibilitychange', onVisibility); + return () => { + active = false; + clearTimeout(timer); + window.removeEventListener('focus', onFocus); + document.removeEventListener('visibilitychange', onVisibility); + }; + }, [designId, lastUsage, settledGenerationIds, onError]); + + return budget; +} diff --git a/packages/i18n/src/locales/en.json b/packages/i18n/src/locales/en.json index 2ee9630d..9dd24826 100644 --- a/packages/i18n/src/locales/en.json +++ b/packages/i18n/src/locales/en.json @@ -374,6 +374,10 @@ "thinking": "Thinking", "streamingLabel": "Assistant is typing", "tokensLine": "~{{count}} tokens", + "usageLabel": "Usage for this design", + "usageLoadFailed": "Could not read usage totals", + "usageDesign": "Design {{tokens}} tokens · {{cost}}", + "usageWindow": "Today {{todayTokens}} · {{todayCost}} · Week {{weekTokens}} · {{weekCost}}", "artifactDelivered": "delivered", "artifactDefaultLabel": "design.html", "tool": { diff --git a/packages/i18n/src/locales/es.json b/packages/i18n/src/locales/es.json index b9140674..2c9d51a0 100644 --- a/packages/i18n/src/locales/es.json +++ b/packages/i18n/src/locales/es.json @@ -294,6 +294,10 @@ "thinking": "Pensando", "streamingLabel": "El asistente está escribiendo", "tokensLine": "~{{count}} tokens", + "usageLabel": "Uso de este diseño", + "usageLoadFailed": "No se pudieron leer los totales de uso", + "usageDesign": "Diseño {{tokens}} tokens · {{cost}}", + "usageWindow": "Hoy {{todayTokens}} · {{todayCost}} · Semana {{weekTokens}} · {{weekCost}}", "artifactDelivered": "entregado", "artifactDefaultLabel": "diseño.html", "tool": { diff --git a/packages/i18n/src/locales/pt-BR.json b/packages/i18n/src/locales/pt-BR.json index 9beab78f..1f943059 100644 --- a/packages/i18n/src/locales/pt-BR.json +++ b/packages/i18n/src/locales/pt-BR.json @@ -290,6 +290,10 @@ "noModel": "Nenhum modelo selecionado", "thinking": "Pensando", "tokensLine": "~{{count}} tokens", + "usageLabel": "Uso deste design", + "usageLoadFailed": "Não foi possível ler os totais de uso", + "usageDesign": "Design {{tokens}} tokens · {{cost}}", + "usageWindow": "Hoje {{todayTokens}} · {{todayCost}} · Semana {{weekTokens}} · {{weekCost}}", "artifactDelivered": "entregue", "artifactDefaultLabel": "design.html", "tool": { diff --git a/packages/i18n/src/locales/zh-CN.json b/packages/i18n/src/locales/zh-CN.json index 05ab4582..fddc4149 100644 --- a/packages/i18n/src/locales/zh-CN.json +++ b/packages/i18n/src/locales/zh-CN.json @@ -374,6 +374,10 @@ "thinking": "正在思考", "streamingLabel": "助手正在回复", "tokensLine": "约 {{count}} tokens", + "usageLabel": "当前设计用量", + "usageLoadFailed": "无法读取用量统计", + "usageDesign": "本设计 {{tokens}} tokens · {{cost}}", + "usageWindow": "今日 {{todayTokens}} · {{todayCost}} · 本周 {{weekTokens}} · {{weekCost}}", "artifactDelivered": "已生成", "artifactDefaultLabel": "design.html", "tool": { diff --git a/packages/shared/src/index.ts b/packages/shared/src/index.ts index da8c5e13..ff2afad3 100644 --- a/packages/shared/src/index.ts +++ b/packages/shared/src/index.ts @@ -626,4 +626,14 @@ export { replaceTweakSchema, } from './editmode'; +export { + formatUsageCost, + formatUsageTokens, + nextLocalMidnight, + summarizeUsageBudget, + type UsageBudget, + type UsageRecord, + type UsageTotals, +} from './usage-budget'; + export * from './web-research'; diff --git a/packages/shared/src/usage-budget.test.ts b/packages/shared/src/usage-budget.test.ts new file mode 100644 index 00000000..eb6cbe49 --- /dev/null +++ b/packages/shared/src/usage-budget.test.ts @@ -0,0 +1,99 @@ +import { describe, expect, it } from 'vitest'; +import { + formatUsageCost, + formatUsageTokens, + nextLocalMidnight, + summarizeUsageBudget, +} from './usage-budget'; + +const now = new Date(2024, 0, 7, 12).getTime(); + +describe('summarizeUsageBudget', () => { + it('keeps another design out of the current design totals', () => { + const budget = summarizeUsageBudget( + [ + { + designId: 'a', + startedAt: new Date(2024, 0, 7, 9).getTime(), + inputTokens: 10, + outputTokens: 5, + costUsd: 0.25, + }, + { + designId: 'b', + startedAt: new Date(2024, 0, 7, 9).getTime(), + inputTokens: 100, + outputTokens: 50, + costUsd: 9, + }, + { + designId: 'a', + startedAt: new Date(2024, 0, 2, 9).getTime(), + inputTokens: 4, + outputTokens: 1, + costUsd: 0.5, + }, + { + designId: 'a', + startedAt: new Date(2023, 11, 31, 23).getTime(), + inputTokens: 7, + outputTokens: 3, + costUsd: 0.125, + }, + { + designId: 'a', + startedAt: new Date(2024, 0, 8, 1).getTime(), + inputTokens: 1, + outputTokens: 1, + costUsd: 0.25, + }, + ], + 'a', + now, + ); + + expect(budget.design).toEqual({ inputTokens: 22, outputTokens: 10, costUsd: 1.125 }); + expect(budget.today).toEqual({ inputTokens: 10, outputTokens: 5, costUsd: 0.25 }); + expect(budget.week).toEqual({ inputTokens: 14, outputTokens: 6, costUsd: 0.75 }); + }); + + it('counts a record that starts at the local week boundary', () => { + const budget = summarizeUsageBudget( + [ + { + designId: 'a', + startedAt: new Date(2024, 0, 1).getTime(), + inputTokens: 2, + outputTokens: 3, + costUsd: 0.5, + }, + ], + 'a', + now, + ); + expect(budget.week).toEqual({ inputTokens: 2, outputTokens: 3, costUsd: 0.5 }); + expect(budget.today).toEqual({ inputTokens: 0, outputTokens: 0, costUsd: 0 }); + }); +}); + +describe('usage formatting', () => { + it('formats token counts and small USD amounts', () => { + expect(formatUsageTokens(1234)).toBe('1,234'); + expect(formatUsageCost(0)).toBe('$0'); + expect(formatUsageCost(1.2)).toBe('$1.2'); + expect(formatUsageCost(0.0125)).toBe('$0.0125'); + expect(formatUsageCost(0.000012)).toBe('$0.000012'); + }); +}); + +describe('nextLocalMidnight', () => { + it('lands on the next local midnight and on Monday when the week turns', () => { + expect(nextLocalMidnight(new Date(2024, 0, 6, 23).getTime())).toBe( + new Date(2024, 0, 7).getTime(), + ); + expect(nextLocalMidnight(new Date(2024, 0, 7, 12).getTime())).toBe( + new Date(2024, 0, 8).getTime(), + ); + expect(nextLocalMidnight(new Date(2024, 0, 8).getTime())).toBe(new Date(2024, 0, 9).getTime()); + }); +}); diff --git a/packages/shared/src/usage-budget.ts b/packages/shared/src/usage-budget.ts new file mode 100644 index 00000000..5c02bc97 --- /dev/null +++ b/packages/shared/src/usage-budget.ts @@ -0,0 +1,77 @@ +export interface UsageRecord { + designId: string; + startedAt: number; + inputTokens: number; + outputTokens: number; + costUsd: number; +} + +export interface UsageTotals { + inputTokens: number; + outputTokens: number; + costUsd: number; +} + +export interface UsageBudget { + design: UsageTotals; + today: UsageTotals; + week: UsageTotals; +} + +function startOfLocalDay(now: number): number { + const date = new Date(now); + return new Date(date.getFullYear(), date.getMonth(), date.getDate()).getTime(); +} + +/** The next local midnight after now. Monday 00:00 is the week boundary too. */ +export function nextLocalMidnight(now: number): number { + const date = new Date(now); + return new Date(date.getFullYear(), date.getMonth(), date.getDate() + 1).getTime(); +} + +/** Local weeks start on Monday. */ +function startOfLocalWeek(now: number): number { + const date = new Date(now); + const daysSinceMonday = (date.getDay() + 6) % 7; + return new Date(date.getFullYear(), date.getMonth(), date.getDate() - daysSinceMonday).getTime(); +} + +function add(total: UsageTotals, record: UsageRecord): void { + total.inputTokens += record.inputTokens; + total.outputTokens += record.outputTokens; + total.costUsd += record.costUsd; +} + +export function summarizeUsageBudget( + records: readonly UsageRecord[], + designId: string, + now: number, +): UsageBudget { + const design: UsageTotals = { inputTokens: 0, outputTokens: 0, costUsd: 0 }; + const today: UsageTotals = { inputTokens: 0, outputTokens: 0, costUsd: 0 }; + const week: UsageTotals = { inputTokens: 0, outputTokens: 0, costUsd: 0 }; + const dayStart = startOfLocalDay(now); + const weekStart = startOfLocalWeek(now); + for (const record of records) { + if (record.designId !== designId || !Number.isFinite(record.startedAt)) continue; + add(design, record); + if (record.startedAt > now) continue; + if (record.startedAt >= weekStart) add(week, record); + if (record.startedAt >= dayStart) add(today, record); + } + return { design, today, week }; +} + +export function formatUsageTokens(value: number): string { + return new Intl.NumberFormat('en-US', { maximumFractionDigits: 0 }).format(value); +} + +export function formatUsageCost(value: number): string { + if (!Number.isFinite(value) || value <= 0) return '$0'; + const digits = value >= 1 ? 2 : value >= 0.01 ? 4 : 6; + const text = value + .toFixed(digits) + .replace(/(\.\d*?)0+$/, '$1') + .replace(/\.$/, ''); + return `$${text}`; +}