Skip to content

Commit f93227c

Browse files
authored
Merge pull request #146 from github/asizikov/native-budget-simulation
Enable native budget simulation
2 parents 6dceec7 + 23516e8 commit f93227c

5 files changed

Lines changed: 190 additions & 65 deletions

File tree

src/App.tsx

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -358,11 +358,6 @@ function App() {
358358
const handleApplyBudgetSimulation = useCallback(async () => {
359359
const file = currentFileRef.current
360360
if (!file) return
361-
if (isNativeAiCreditsReport) {
362-
setBudgetSimulation(null)
363-
setBudgetSimulationError('Budget simulation is not available for native AI Credits reports yet.')
364-
return
365-
}
366361

367362
const budgetReportUsers = userUsage?.users ?? []
368363
const hasBudgetOrganizationContext = budgetReportUsers.some((user) => user.organizations.length > 0 || user.costCenters.length > 0)
@@ -427,6 +422,7 @@ function App() {
427422
},
428423
},
429424
resolveIncludedCreditOverrides(seatOverrides),
425+
{ reportMetadata: reportMetadata ?? undefined },
430426
)
431427

432428
if (simulationId !== latestSimulationIdRef.current) return
@@ -448,7 +444,7 @@ function App() {
448444
budgetValues.productCopilot,
449445
budgetValues.productSpark,
450446
budgetValues.user,
451-
isNativeAiCreditsReport,
447+
reportMetadata,
452448
resolveIncludedCreditOverrides,
453449
seatOverrides,
454450
userUsage,

src/utils/budgetSimulation.test.ts

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,35 @@ const HEADER = [
2323
'aic_gross_amount',
2424
].join(',')
2525

26+
const NATIVE_AI_CREDITS_HEADER = [
27+
'date',
28+
'username',
29+
'product',
30+
'sku',
31+
'model',
32+
'quantity',
33+
'unit_type',
34+
'applied_cost_per_quantity',
35+
'gross_amount',
36+
'discount_amount',
37+
'net_amount',
38+
'total_monthly_quota',
39+
'organization',
40+
'cost_center_name',
41+
'aic_quantity',
42+
'aic_gross_amount',
43+
].join(',')
44+
2645
function createCsv(rows: string[][]): File {
2746
const body = [HEADER, ...rows.map((row) => row.join(','))].join('\n')
2847
return new File([body], 'usage.csv', { type: 'text/csv' })
2948
}
3049

50+
function createNativeAiCreditsCsv(rows: string[][]): File {
51+
const body = [NATIVE_AI_CREDITS_HEADER, ...rows.map((row) => row.join(','))].join('\n')
52+
return new File([body], 'native-usage.csv', { type: 'text/csv' })
53+
}
54+
3155
function createRecord(overrides: Partial<TokenUsageRecord>): TokenUsageRecord {
3256
const quantity = overrides.quantity ?? 0
3357

@@ -471,4 +495,34 @@ describe('runBudgetSimulation', () => {
471495
adjustedDailyGrossCostByDate: [{ date: '2026-04-25', amount: 0.5 }],
472496
})
473497
})
498+
499+
it('uses native AI Credits report parsing and policy context when report metadata is provided', async () => {
500+
const file = createNativeAiCreditsCsv([
501+
['6/1/26', 'mona', 'copilot', 'copilot_ai_credit', 'GPT-5', '100', 'ai-credits', '0.01', '1.00', '0', '1.00', '0', 'example-org', 'Cost Center A', '', ''],
502+
])
503+
504+
await expect(runBudgetSimulation(
505+
file,
506+
{ accountBudgetUsd: 0.5 },
507+
{},
508+
{
509+
reportMetadata: {
510+
format: 'native-ai-credits',
511+
label: 'Native AI Credits report',
512+
},
513+
},
514+
)).resolves.toEqual({
515+
totalBill: 0.5,
516+
blockedUsers: 1,
517+
blockedRequests: 0,
518+
blockedIncludedCreditsAic: 0,
519+
allowedAicQuantity: 50,
520+
budgetExhausted: true,
521+
firstUserBlockedDate: null,
522+
accountBlockedDate: '2026-06-01',
523+
productBlockedDates: {},
524+
adjustedDailyNetCostByDate: [{ date: '2026-06-01', amount: 0.5 }],
525+
adjustedDailyGrossCostByDate: [{ date: '2026-06-01', amount: 0.5 }],
526+
})
527+
})
474528
})

src/utils/budgetSimulation.ts

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import { calculateAicIncludedCreditsContext, getUsageMonthKey, type AicIncludedCreditsContext, type AicIncludedCreditsOverrides } from '../pipeline/aicIncludedCredits'
2-
import { getAicUsageMetrics, getUsageMetrics, parseNormalizedTokenUsageRecord, parseTokenUsageHeader, type TokenUsageHeader, type TokenUsageRecord } from '../pipeline/parser'
2+
import { getAicUsageMetrics, getUsageMetrics, parseNativeAiCreditsUsageRecord, parseNormalizedTokenUsageRecord, parseTokenUsageHeader, type TokenUsageHeader, type TokenUsageRecord } from '../pipeline/parser'
3+
import type { ReportFormatMetadata } from '../pipeline/reportAdapters'
34
import { getProductBudgetName, isNonCopilotCodeReviewUsage, NON_COPILOT_CODE_REVIEW_USER_LABEL, type ProductBudgetName } from '../pipeline/productClassification'
45
import { streamLines } from '../pipeline/streamer'
56
import type { UserSpendSegmentId } from './userSpendSegments'
@@ -26,6 +27,10 @@ export type BudgetSimulationOptions = {
2627
productBudgetsUsd?: Partial<Record<ProductBudgetName, number>>
2728
}
2829

30+
export type BudgetSimulationRunOptions = {
31+
reportMetadata?: ReportFormatMetadata
32+
}
33+
2934
type BudgetSimulationContext = Pick<AicIncludedCreditsContext, 'reportPlanScope' | 'organizationIncludedCreditsPool' | 'individualMonthlyIncludedCredits'>
3035
type BudgetSimulationState = {
3136
remainingAccountBudget: number
@@ -378,8 +383,11 @@ export async function runBudgetSimulation(
378383
file: File,
379384
options: BudgetSimulationOptions,
380385
includedCreditsOverrides: AicIncludedCreditsOverrides = {},
386+
runOptions: BudgetSimulationRunOptions = {},
381387
): Promise<BudgetSimulationResult> {
382-
const context = await calculateAicIncludedCreditsContext(file, includedCreditsOverrides)
388+
const context = await calculateAicIncludedCreditsContext(file, includedCreditsOverrides, {
389+
reportMetadata: runOptions.reportMetadata,
390+
})
383391
const state = createBudgetSimulationState(options, context)
384392
let header: TokenUsageHeader | null = null
385393

@@ -392,7 +400,9 @@ export async function runBudgetSimulation(
392400
continue
393401
}
394402

395-
const record = parseNormalizedTokenUsageRecord(trimmed, header)
403+
const record = runOptions.reportMetadata?.format === 'native-ai-credits'
404+
? parseNativeAiCreditsUsageRecord(trimmed, header)
405+
: parseNormalizedTokenUsageRecord(trimmed, header)
396406
if (!record) continue
397407

398408
simulateBudgetRecord(state, record, context)
Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
import { createElement } from 'react'
2+
import { renderToStaticMarkup } from 'react-dom/server'
3+
import { describe, expect, it, vi } from 'vitest'
4+
5+
import type { BudgetSimulationResult } from '../utils/budgetSimulation'
6+
import { EMPTY_BUDGET_VALUES, type BudgetValues } from '../utils/costManagementBudgets'
7+
import { CostManagementView } from './CostManagementView'
8+
9+
const baseBudgetValues: BudgetValues = {
10+
...EMPTY_BUDGET_VALUES,
11+
account: '1',
12+
}
13+
14+
const baseBudgetSimulation: BudgetSimulationResult = {
15+
totalBill: 0.4,
16+
blockedUsers: 1,
17+
blockedRequests: 7,
18+
blockedIncludedCreditsAic: 0,
19+
allowedAicQuantity: 40,
20+
budgetExhausted: true,
21+
firstUserBlockedDate: null,
22+
accountBlockedDate: '2026-06-01',
23+
productBlockedDates: {},
24+
adjustedDailyNetCostByDate: [],
25+
adjustedDailyGrossCostByDate: [],
26+
}
27+
28+
function renderCostManagementView(overrides: Partial<Parameters<typeof CostManagementView>[0]> = {}): string {
29+
return renderToStaticMarkup(createElement(CostManagementView, {
30+
budgetValues: baseBudgetValues,
31+
isIndividualReport: false,
32+
currentPruBill: 0,
33+
currentPruGrossAmount: 0,
34+
currentPruDiscountAmount: 0,
35+
currentPruQuantity: 0,
36+
currentAicBill: 1,
37+
currentAicGrossAmount: 1,
38+
currentAicDiscountAmount: 0,
39+
currentAicQuantity: 100,
40+
includedAicPoolSize: 0,
41+
dailyUsageData: [],
42+
budgetSimulation: null,
43+
budgetSimulationError: null,
44+
isApplyingBudgetSimulation: false,
45+
onBudgetValueChange: vi.fn(),
46+
onApplyBudgetSimulation: vi.fn(),
47+
showOrganizationPromotionalDataDisclaimer: false,
48+
...overrides,
49+
}))
50+
}
51+
52+
describe('CostManagementView', () => {
53+
it('shows budget controls for native usage-based billing reports', () => {
54+
const html = renderCostManagementView({
55+
reportMode: 'native-ai-credits',
56+
})
57+
58+
expect(html).toContain('Set USD budgets and preview how they would affect usage-based billing for this report.')
59+
expect(html).toContain('Account-level budget')
60+
expect(html).toContain('Apply')
61+
expect(html).not.toContain('Budget simulation is not available')
62+
expect(html).not.toContain('native AI Credits reports yet')
63+
expect(html).not.toContain('PRU')
64+
})
65+
66+
it('uses AI Credits result labels instead of PRU labels for native reports', () => {
67+
const html = renderCostManagementView({
68+
reportMode: 'native-ai-credits',
69+
budgetSimulation: baseBudgetSimulation,
70+
})
71+
72+
expect(html).toContain('Blocked AI Credits')
73+
expect(html).toContain('60')
74+
expect(html).toContain('Simulated AI Credits additional usage spend')
75+
expect(html).not.toContain('Blocked PRUs')
76+
expect(html).not.toContain('later requests')
77+
})
78+
79+
it('keeps PRU blocked-usage labeling for transition-period reports', () => {
80+
const html = renderCostManagementView({
81+
reportMode: 'transition-period-billing-preview',
82+
budgetSimulation: baseBudgetSimulation,
83+
})
84+
85+
expect(html).toContain('Blocked PRUs')
86+
expect(html).toContain('later requests')
87+
expect(html).not.toContain('Blocked AI Credits')
88+
})
89+
})

0 commit comments

Comments
 (0)