Merge pull request #141 from github/asizikov/gated-native-pipeline

Gate native AI Credits pipeline processing

Author: asizikov

src/pipeline/aicIncludedCredits.ts

@@ -1,5 +1,6 @@
 import {
   getAicUsageMetrics,
+  parseNativeAiCreditsUsageRecord,
   parseTokenUsageHeader,
   parseNormalizedTokenUsageRecord,
   type TokenUsageHeader,
@@ -254,6 +255,18 @@ function resolvePolicyForContext(
   return TRANSITION_PERIOD_INCLUDED_CREDITS_POLICY
 }
 
+function parseIncludedCreditsRecord(
+  line: string,
+  header: TokenUsageHeader,
+  options: AicIncludedCreditsProgressOptions | undefined,
+): TokenUsageRecord | null {
+  if (options?.reportMetadata?.format === 'native-ai-credits') {
+    return parseNativeAiCreditsUsageRecord(line, header)
+  }
+
+  return parseNormalizedTokenUsageRecord(line, header)
+}
+
 export function calculateLicenseSummary(
   users: Array<{ totalMonthlyQuota: number } & ReportScopeUser>,
   policy: IncludedCreditsPolicy = TRANSITION_PERIOD_INCLUDED_CREDITS_POLICY,
@@ -318,7 +331,7 @@ export async function calculateAicIncludedCreditsContext(
       continue
     }
 
-    const record = parseNormalizedTokenUsageRecord(trimmed, header)
+    const record = parseIncludedCreditsRecord(trimmed, header, options)
     if (!record) continue
 
     reportPeriod = includeDateInReportPeriod(reportPeriod, record.date)

src/pipeline/reportAdapters.ts

@@ -24,6 +24,10 @@ export interface UsageReportAdapter {
   parseRecord(line: string, header: TokenUsageHeader): TokenUsageRecord | null
 }
 
+export interface UsageReportValidationOptions {
+  allowUnsupportedNativeAiCredits?: boolean
+}
+
 const TRANSITION_PERIOD_BILLING_PREVIEW_REPORT_ADAPTER: UsageReportAdapter = {
   metadata: {
     format: 'transition-period-billing-preview',
@@ -89,9 +93,12 @@ export function selectUsageReportAdapter(header: TokenUsageHeader, firstRecord:
 export function validateUsageReportFirstRecord(
   header: TokenUsageHeader,
   firstRecord: TokenUsageRecord,
+  options?: UsageReportValidationOptions,
 ): UsageReportAdapter {
   const adapter = selectUsageReportAdapter(header, firstRecord)
   adapter.validateHeader(header)
-  adapter.validateFirstRecord(header, firstRecord)
+  if (!(options?.allowUnsupportedNativeAiCredits && adapter.metadata.format === 'native-ai-credits')) {
+    adapter.validateFirstRecord(header, firstRecord)
+  }
   return adapter
 }

src/pipeline/runPipeline.test.ts

@@ -1,6 +1,8 @@
 import { describe, expect, it } from 'vitest'
 
 import type { Aggregator } from './aggregators/base'
+import { DailyUsageAggregator } from './aggregators/dailyUsageAggregator'
+import { UserUsageAggregator } from './aggregators/userUsageAggregator'
 import {
   InvalidReportError,
   UnsupportedNativeAiCreditsReportError,
@@ -55,6 +57,12 @@ const TRANSITION_PERIOD_REPORT_METADATA = {
   supported: true,
 }
 
+const NATIVE_AI_CREDITS_REPORT_METADATA = {
+  format: 'native-ai-credits',
+  label: 'Native AI Credits report',
+  supported: false,
+} as const
+
 function createCsv(rows: string[][], header = HEADER): File {
   const body = [header, ...rows.map((row) => row.join(','))].join('\n')
   return new File([body], 'usage.csv', { type: 'text/csv' })
@@ -160,6 +168,181 @@ describe('runPipeline', () => {
     expect(aggregator.result()).toEqual([])
   })
 
+  it('processes native AI Credits reports with native metadata when explicitly enabled', async () => {
+    const file = createCsv([
+      [
+        '5/29/26',
+        'mona',
+        'copilot',
+        'copilot_ai_credit',
+        'Auto: Claude Haiku 4.5',
+        '96.9990345',
+        'ai-credits',
+        '0.01',
+        '0.969990345',
+        '0',
+        '0.969990345',
+        '3900',
+        'example-org',
+        'Cost Center A',
+        '999',
+        '999',
+      ],
+    ], NATIVE_AI_CREDITS_HEADER)
+    const aggregator = new CaptureAggregator()
+
+    const result = await runPipeline(file, [aggregator], {
+      enableNativeAiCreditsProcessing: true,
+    })
+
+    expect(result).toEqual({
+      reportMetadata: NATIVE_AI_CREDITS_REPORT_METADATA,
+      reportRowCount: 1,
+      processedRowCount: 1,
+    })
+    expect(aggregator.headerCalls()).toBe(1)
+    expect(aggregator.result()).toEqual([
+      expect.objectContaining({
+        date: '2026-05-29',
+        username: 'mona',
+        quantity: 96.9990345,
+        gross_amount: 0.969990345,
+        net_amount: 0.969990345,
+        aic_quantity: 96.9990345,
+        aic_gross_amount: 0.969990345,
+        has_aic_quantity: true,
+        has_aic_gross_amount: true,
+      }),
+    ])
+  })
+
+  it('aggregates flagged native AI Credits rows with native-format aggregators', async () => {
+    const file = createCsv([
+      [
+        '5/29/26',
+        'mona',
+        'copilot',
+        'copilot_ai_credit',
+        'GPT-5.2',
+        '10',
+        'ai-credits',
+        '0.01',
+        '100',
+        '20',
+        '80',
+        '3900',
+        'example-org',
+        'Cost Center A',
+        '999',
+        '999',
+      ],
+      [
+        '05/29/2026',
+        'hubot',
+        'spark',
+        'spark_ai_credit',
+        'GPT-5.2',
+        '25',
+        'ai-credits',
+        '0.01',
+        '250',
+        '50',
+        '200',
+        '7000',
+        'octodemo',
+        'Cost Center A',
+        '',
+        '',
+      ],
+    ], NATIVE_AI_CREDITS_HEADER)
+    const daily = new DailyUsageAggregator(NATIVE_AI_CREDITS_REPORT_METADATA)
+    const users = new UserUsageAggregator(NATIVE_AI_CREDITS_REPORT_METADATA)
+
+    await runPipeline(file, [daily, users], {
+      enableNativeAiCreditsProcessing: true,
+    })
+
+    expect(daily.result().dailyData).toEqual([
+      expect.objectContaining({
+        date: '2026-05-29',
+        requests: 0,
+        grossAmount: 0,
+        discountAmount: 0,
+        netAmount: 0,
+        aicQuantity: 35,
+        aicGrossAmount: 350,
+        aicNetAmount: 280,
+      }),
+    ])
+    expect(users.result().users).toEqual([
+      expect.objectContaining({
+        username: 'hubot',
+        totals: expect.objectContaining({
+          requests: 0,
+          grossAmount: 0,
+          discountAmount: 0,
+          netAmount: 0,
+          aicQuantity: 25,
+          aicGrossAmount: 250,
+          aicNetAmount: 200,
+        }),
+      }),
+      expect.objectContaining({
+        username: 'mona',
+        totals: expect.objectContaining({
+          requests: 0,
+          grossAmount: 0,
+          discountAmount: 0,
+          netAmount: 0,
+          aicQuantity: 10,
+          aicGrossAmount: 100,
+          aicNetAmount: 80,
+        }),
+      }),
+    ])
+  })
+
+  it('selects native summer and September included-credit policies for flagged native reports', async () => {
+    const createNativePolicyCsv = (date: string) => createCsv([
+      [
+        date,
+        'mona',
+        'copilot',
+        'copilot_ai_credit',
+        'GPT-5.2',
+        '5000',
+        'ai-credits',
+        '0.01',
+        '50',
+        '0',
+        '50',
+        '3900',
+        'example-org',
+        'Cost Center A',
+        '5000',
+        '50',
+      ],
+    ], NATIVE_AI_CREDITS_HEADER)
+    const summerAggregator = new CaptureAggregator()
+    const septemberAggregator = new CaptureAggregator()
+
+    await runPipeline(createNativePolicyCsv('8/31/26'), [summerAggregator], {
+      enableNativeAiCreditsProcessing: true,
+    })
+    await runPipeline(createNativePolicyCsv('9/1/26'), [septemberAggregator], {
+      enableNativeAiCreditsProcessing: true,
+    })
+
+    expect(summerAggregator.result()[0]).toEqual(expect.objectContaining({
+      date: '2026-08-31',
+      aic_net_amount: 0,
+    }))
+    expect(septemberAggregator.result()[0]).toEqual(expect.objectContaining({
+      date: '2026-09-01',
+    }))
+    expect(septemberAggregator.result()[0].aic_net_amount).toBeCloseTo(11)
+  })
+
   it('returns transition-period metadata while processing supported reports', async () => {
     const file = createCsv([
       ['2026-04-25', 'mona', 'copilot', 'copilot_premium_request', 'GPT-5', '0', 'requests', '0.04', '0', '0', '0', 'False', '300', '', '', '0', '0'],

src/pipeline/runPipeline.ts

@@ -12,10 +12,11 @@ import {
   validateUsageReportHeader,
   type ReportFormatMetadata,
   type UsageReportAdapter,
+  type UsageReportValidationOptions,
 } from './reportAdapters'
 import { streamLines, type StreamProgress } from './streamer'
 
-async function validateFileFormat(file: File): Promise<UsageReportAdapter> {
+async function validateFileFormat(file: File, options?: UsageReportValidationOptions): Promise<UsageReportAdapter> {
   let header: TokenUsageHeader | null = null
   let selectedAdapter: UsageReportAdapter | null = null
 
@@ -31,7 +32,7 @@ async function validateFileFormat(file: File): Promise<UsageReportAdapter> {
       continue
     }
 
-    return validateUsageReportFirstRecord(header, parseTokenUsageRecord(trimmed, header))
+    return validateUsageReportFirstRecord(header, parseTokenUsageRecord(trimmed, header), options)
   }
 
   if (!selectedAdapter) {
@@ -50,6 +51,7 @@ export interface PipelineProgress {
 }
 
 export interface PipelineOptions {
+  enableNativeAiCreditsProcessing?: boolean
   includedCreditsOverrides?: AicIncludedCreditsOverrides
   progressResolution?: number
   onProgress?: (progress: PipelineProgress) => void
@@ -89,8 +91,15 @@ export async function runPipeline(
   aggregators: Aggregator<TokenUsageRecord, unknown, TokenUsageHeader>[],
   options?: PipelineOptions,
 ): Promise<PipelineResult> {
-  const { includedCreditsOverrides = {}, progressResolution = 500, onProgress } = options ?? {}
-  const reportAdapter = await validateFileFormat(file)
+  const {
+    enableNativeAiCreditsProcessing = false,
+    includedCreditsOverrides = {},
+    progressResolution = 500,
+    onProgress,
+  } = options ?? {}
+  const reportAdapter = await validateFileFormat(file, {
+    allowUnsupportedNativeAiCredits: enableNativeAiCreditsProcessing,
+  })
   const reportMetadata = reportAdapter.metadata
   let lastProgressStage: PipelineProgress['stage'] | null = null
   let lastProgressPercent = -1