feat(billing): implement G1 AI credits overage flow with billing telemetry (#18590)

packages/core/src/code_assist/codeAssist.test.ts

@@ -73,6 +73,8 @@ describe('codeAssist', () => {
         'session-123',
         'free-tier',
         'free-tier-name',
+        undefined,
+        mockConfig,
       );
       expect(generator).toBeInstanceOf(MockedCodeAssistServer);
     });
@@ -103,6 +105,8 @@ describe('codeAssist', () => {
         undefined, // No session ID
         'free-tier',
         'free-tier-name',
+        undefined,
+        mockConfig,
       );
       expect(generator).toBeInstanceOf(MockedCodeAssistServer);
     });

packages/core/src/code_assist/codeAssist.ts

@@ -36,6 +36,8 @@ export async function createCodeAssistContentGenerator(
       sessionId,
       userData.userTier,
       userData.userTierName,
+      userData.paidTier,
+      config,
     );
   }
 

packages/core/src/code_assist/converter.ts

@@ -28,12 +28,14 @@ import type {
 } from '@google/genai';
 import { GenerateContentResponse } from '@google/genai';
 import { debugLogger } from '../utils/debugLogger.js';
+import type { Credits } from './types.js';
 
 export interface CAGenerateContentRequest {
   model: string;
   project?: string;
   user_prompt_id?: string;
   request: VertexGenerateContentRequest;
+  enabled_credit_types?: string[];
 }
 
 interface VertexGenerateContentRequest {
@@ -75,6 +77,8 @@ interface VertexGenerationConfig {
 export interface CaGenerateContentResponse {
   response?: VertexGenerateContentResponse;
   traceId?: string;
+  consumedCredits?: Credits[];
+  remainingCredits?: Credits[];
 }
 
 interface VertexGenerateContentResponse {
@@ -127,12 +131,14 @@ export function toGenerateContentRequest(
   userPromptId: string,
   project?: string,
   sessionId?: string,
+  enabledCreditTypes?: string[],
 ): CAGenerateContentRequest {
   return {
     model: req.model,
     project,
     user_prompt_id: userPromptId,
     request: toVertexGenerateContentRequest(req, sessionId),
+    enabled_credit_types: enabledCreditTypes,
   };
 }
 
@@ -305,3 +311,16 @@ function toVertexGenerationConfig(
     thinkingConfig: config.thinkingConfig,
   };
 }
+
+export function fromGenerateContentResponseUsage(
+  metadata?: GenerateContentResponseUsageMetadata,
+): GenerateContentResponseUsageMetadata | undefined {
+  if (!metadata) {
+    return undefined;
+  }
+  return {
+    promptTokenCount: metadata.promptTokenCount,
+    candidatesTokenCount: metadata.candidatesTokenCount,
+    totalTokenCount: metadata.totalTokenCount,
+  };
+}

packages/core/src/code_assist/server.ts

@@ -21,7 +21,10 @@ import type {
   ConversationInteraction,
   StreamingLatency,
   RecordCodeAssistMetricsRequest,
+  GeminiUserTier,
+  Credits,
 } from './types.js';
+import { UserTierId } from './types.js';
 import type {
   ListExperimentsRequest,
   ListExperimentsResponse,
@@ -37,7 +40,15 @@ import type {
 import * as readline from 'node:readline';
 import { Readable } from 'node:stream';
 import type { ContentGenerator } from '../core/contentGenerator.js';
-import { UserTierId } from './types.js';
+import type { Config } from '../config/config.js';
+import {
+  G1_CREDIT_TYPE,
+  getG1CreditBalance,
+  isOverageEligibleModel,
+  shouldAutoUseCredits,
+} from '../billing/billing.js';
+import { logBillingEvent } from '../telemetry/loggers.js';
+import { CreditsUsedEvent } from '../telemetry/billingEvents.js';
 import type {
   CaCountTokenResponse,
   CaGenerateContentResponse,
@@ -72,6 +83,8 @@ export class CodeAssistServer implements ContentGenerator {
     readonly sessionId?: string,
     readonly userTier?: UserTierId,
     readonly userTierName?: string,
+    readonly paidTier?: GeminiUserTier,
+    readonly config?: Config,
   ) {}
 
   async generateContentStream(
@@ -80,6 +93,19 @@ export class CodeAssistServer implements ContentGenerator {
     // eslint-disable-next-line @typescript-eslint/no-unused-vars
     role: LlmRole,
   ): Promise<AsyncGenerator<GenerateContentResponse>> {
+    const autoUse = this.config
+      ? shouldAutoUseCredits(
+          this.config.getBillingSettings().overageStrategy,
+          getG1CreditBalance(this.paidTier),
+        )
+      : false;
+    const modelIsEligible = isOverageEligibleModel(req.model);
+    const shouldEnableCredits = modelIsEligible && autoUse;
+
+    const enabledCreditTypes = shouldEnableCredits
+      ? ([G1_CREDIT_TYPE] as string[])
+      : undefined;
+
     const responses =
       await this.requestStreamingPost<CaGenerateContentResponse>(
         'streamGenerateContent',
@@ -88,6 +114,7 @@ export class CodeAssistServer implements ContentGenerator {
           userPromptId,
           this.projectId,
           this.sessionId,
+          enabledCreditTypes,
         ),
         req.config?.abortSignal,
       );
@@ -99,6 +126,9 @@ export class CodeAssistServer implements ContentGenerator {
     return (async function* (
       server: CodeAssistServer,
     ): AsyncGenerator<GenerateContentResponse> {
+      let totalConsumed = 0;
+      let lastRemaining = 0;
+
       for await (const response of responses) {
         if (isFirst) {
           streamingLatency.firstMessageLatency = formatProtoJsonDuration(
@@ -121,8 +151,38 @@ export class CodeAssistServer implements ContentGenerator {
           req.config?.abortSignal,
         );
 
+        if (response.consumedCredits) {
+          for (const credit of response.consumedCredits) {
+            if (credit.creditType === G1_CREDIT_TYPE && credit.creditAmount) {
+              totalConsumed += parseInt(credit.creditAmount, 10) || 0;
+            }
+          }
+        }
+        if (response.remainingCredits) {
+          // Sum all G1 credit entries for consistency with getG1CreditBalance
+          lastRemaining = response.remainingCredits.reduce((sum, credit) => {
+            if (credit.creditType === G1_CREDIT_TYPE && credit.creditAmount) {
+              return sum + (parseInt(credit.creditAmount, 10) || 0);
+            }
+            return sum;
+          }, 0);
+          server.updateCredits(response.remainingCredits);
+        }
+
         yield translatedResponse;
       }
+
+      // Emit credits used telemetry after the stream completes
+      if (totalConsumed > 0 && server.config) {
+        logBillingEvent(
+          server.config,
+          new CreditsUsedEvent(
+            req.model ?? 'unknown',
+            totalConsumed,
+            lastRemaining,
+          ),
+        );
+      }
     })(this);
   }
 
@@ -140,6 +200,7 @@ export class CodeAssistServer implements ContentGenerator {
         userPromptId,
         this.projectId,
         this.sessionId,
+        undefined,
       ),
       req.config?.abortSignal,
       GENERATE_CONTENT_RETRY_DELAY_IN_MILLISECONDS,
@@ -160,9 +221,29 @@ export class CodeAssistServer implements ContentGenerator {
       req.config?.abortSignal,
     );
 
+    if (response.remainingCredits) {
+      this.updateCredits(response.remainingCredits);
+    }
+
     return translatedResponse;
   }
 
+  private updateCredits(remainingCredits: Credits[]): void {
+    if (!this.paidTier) {
+      return;
+    }
+
+    // Replace the G1 credits entries with the latest remaining amounts.
+    // Non-G1 credits are preserved as-is.
+    const nonG1Credits = (this.paidTier.availableCredits ?? []).filter(
+      (c) => c.creditType !== G1_CREDIT_TYPE,
+    );
+    const updatedG1Credits = remainingCredits.filter(
+      (c) => c.creditType === G1_CREDIT_TYPE,
+    );
+    this.paidTier.availableCredits = [...nonG1Credits, ...updatedG1Credits];
+  }
+
   async onboardUser(
     req: OnboardUserRequest,
   ): Promise<LongRunningOperationResponse> {
@@ -192,6 +273,25 @@ export class CodeAssistServer implements ContentGenerator {
     }
   }
 
+  async refreshAvailableCredits(): Promise<void> {
+    if (!this.paidTier) {
+      return;
+    }
+    const res = await this.loadCodeAssist({
+      cloudaicompanionProject: this.projectId,
+      metadata: {
+        ideType: 'IDE_UNSPECIFIED',
+        platform: 'PLATFORM_UNSPECIFIED',
+        pluginType: 'GEMINI',
+        duetProject: this.projectId,
+      },
+      mode: 'HEALTH_CHECK',
+    });
+    if (res.paidTier?.availableCredits) {
+      this.paidTier.availableCredits = res.paidTier.availableCredits;
+    }
+  }
+
   async fetchAdminControls(
     req: FetchAdminControlsRequest,
   ): Promise<FetchAdminControlsResponse> {

packages/core/src/code_assist/setup.ts

@@ -52,6 +52,7 @@ export interface UserData {
   projectId: string;
   userTier: UserTierId;
   userTierName?: string;
+  paidTier?: GeminiUserTier;
 }
 
 /**
@@ -146,6 +147,7 @@ export async function setupUser(
             loadRes.currentTier.id ??
             UserTierId.STANDARD,
           userTierName: loadRes.paidTier?.name ?? loadRes.currentTier.name,
+          paidTier: loadRes.paidTier ?? undefined,
         };
       }
 
@@ -157,6 +159,7 @@ export async function setupUser(
       userTier:
         loadRes.paidTier?.id ?? loadRes.currentTier.id ?? UserTierId.STANDARD,
       userTierName: loadRes.paidTier?.name ?? loadRes.currentTier.name,
+      paidTier: loadRes.paidTier ?? undefined,
     };
   }
 

packages/core/src/code_assist/types.ts

@@ -39,11 +39,41 @@ export type ClientMetadataPluginType =
   | 'AIPLUGIN_INTELLIJ'
   | 'AIPLUGIN_STUDIO';
 
+/**
+ * Credit types that can be used for API consumption.
+ */
+export type CreditType = 'CREDIT_TYPE_UNSPECIFIED' | 'GOOGLE_ONE_AI';
+
+/**
+ * Represents a credit amount for a specific credit type.
+ * Used in LoadCodeAssistResponse for available credits and
+ * in GenerateContentResponse for consumed/remaining credits.
+ */
+export interface Credits {
+  creditType: CreditType;
+  creditAmount: string; // int64 represented as string in JSON
+}
+
+/** Alias for Credits used in available_credits context */
+export type AvailableCredits = Credits;
+
+/** Alias for Credits used in consumedCredits context */
+export type ConsumedCredits = Credits;
+
+/** Alias for Credits used in remainingCredits context */
+export type RemainingCredits = Credits;
+
 export interface LoadCodeAssistRequest {
   cloudaicompanionProject?: string;
   metadata: ClientMetadata;
+  mode?: LoadCodeAssistMode;
 }
 
+export type LoadCodeAssistMode =
+  | 'MODE_UNSPECIFIED'
+  | 'FULL_ELIGIBILITY_CHECK'
+  | 'HEALTH_CHECK';
+
 /**
  * Represents LoadCodeAssistResponse proto json field
  * http://google3/google/internal/cloud/code/v1internal/cloudcode.proto;l=224
@@ -69,6 +99,8 @@ export interface GeminiUserTier {
   privacyNotice?: PrivacyNotice;
   hasAcceptedTos?: boolean;
   hasOnboardedPreviously?: boolean;
+  /** Available AI credits for this tier (e.g., Google One AI credits) */
+  availableCredits?: AvailableCredits[];
 }
 
 /**