feat: log telemetry for low/depleted CCU balance notifications

src/colab/consumption/notifier.ts

@@ -5,7 +5,8 @@
  */
 
 import vscode, { Disposable, Event } from 'vscode';
-import { CommandSource } from '../../telemetry/api';
+import { telemetry } from '../../telemetry';
+import { CommandSource, LowBalanceSeverity } from '../../telemetry/api';
 import { ConsumptionUserInfo, SubscriptionTier } from '../api';
 import { openColabSignup } from '../commands/external';
 
@@ -89,14 +90,20 @@ export class ConsumptionNotifier implements Disposable {
       return;
     }
 
+    const severity =
+      notification.notify === this.vs.window.showErrorMessage
+        ? LowBalanceSeverity.SEVERITY_DEPLETED
+        : LowBalanceSeverity.SEVERITY_LOW;
     const action = notification.notify(
       notification.message,
       this.getTierRelevantAction(info.subscriptionTier, paidMinutesLeft > 0),
     );
     this.setSnoozeTimeout(notification.notify);
-    if (await action) {
+    const clicked = !!(await action);
+    if (clicked) {
       openColabSignup(this.vs, CommandSource.COMMAND_SOURCE_NOTIFICATION);
     }
+    telemetry.logLowCcuNotification(severity, info.subscriptionTier, clicked);
   }
 
   private buildNotification(totalMinutesLeft: number):

src/colab/consumption/notifier.unit.test.ts

@@ -5,7 +5,9 @@
  */
 
 import { assert, expect } from 'chai';
-import sinon, { SinonFakeTimers } from 'sinon';
+import sinon, { SinonFakeTimers, SinonStubbedFunction } from 'sinon';
+import { telemetry } from '../../telemetry';
+import { LowBalanceSeverity } from '../../telemetry/api';
 import { TestEventEmitter } from '../../test/helpers/events';
 import { newVsCodeStub, VsCodeStub } from '../../test/helpers/vscode';
 import { SubscriptionTier, ConsumptionUserInfo } from '../api';
@@ -426,6 +428,89 @@ describe('ConsumptionNotifier', () => {
     });
   }
 
+  describe('telemetry', () => {
+    let logStub: SinonStubbedFunction<typeof telemetry.logLowCcuNotification>;
+
+    beforeEach(() => {
+      logStub = sinon.stub(telemetry, 'logLowCcuNotification');
+    });
+
+    /**
+     * Stubs out the given severity's notification API to immediately resolve
+     * with the given action (or `undefined` to simulate a dismiss).
+     *
+     * @param severity - The severity of the notification to stub ("warn" or
+     * "error").
+     * @param respondWith - The action to resolve with, or `undefined` to
+     * simulate a dismiss.
+     */
+    function autoRespond(
+      severity: NotificationSeverity,
+      respondWith: string | undefined,
+    ): void {
+      const stub =
+        severity === 'warn'
+          ? vs.window.showWarningMessage
+          : vs.window.showErrorMessage;
+      // Type assertion needed due to overloading.
+      (stub as sinon.SinonStub).callsFake(() => Promise.resolve(respondWith));
+    }
+
+    it('logs SEVERITY_LOW with clicked=false when warning is dismissed', async () => {
+      const ccuInfo = createCcuInfo(
+        { paidMinutes: 0, freeMinutes: 1 },
+        SubscriptionTier.NONE,
+      );
+      autoRespond('warn', undefined);
+
+      const noOp = consumptionNotifier.nextConsumptionCalculation();
+      ccuEmitter.fire(ccuInfo);
+      await noOp;
+
+      sinon.assert.calledOnceWithExactly(
+        logStub,
+        LowBalanceSeverity.SEVERITY_LOW,
+        SubscriptionTier.NONE,
+        false,
+      );
+    });
+
+    const depletedTierCases = [
+      {
+        tierLabel: 'NONE (free)',
+        tier: SubscriptionTier.NONE,
+        actionLabel: 'Sign Up for Colab',
+      },
+      {
+        tierLabel: 'PRO',
+        tier: SubscriptionTier.PRO,
+        actionLabel: 'Upgrade to Pro+',
+      },
+      {
+        tierLabel: 'PRO_PLUS',
+        tier: SubscriptionTier.PRO_PLUS,
+        actionLabel: 'Purchase More CCUs',
+      },
+    ];
+    for (const { tierLabel, tier, actionLabel } of depletedTierCases) {
+      it(`logs SEVERITY_DEPLETED with clicked=true and tier=${tierLabel} when error action is clicked`, async () => {
+        const ccuInfo = createCcuInfo({ paidMinutes: 0, freeMinutes: 0 }, tier);
+        autoRespond('error', actionLabel);
+
+        const noOp = consumptionNotifier.nextConsumptionCalculation();
+        ccuEmitter.fire(ccuInfo);
+        await noOp;
+
+        sinon.assert.calledOnceWithExactly(
+          logStub,
+          LowBalanceSeverity.SEVERITY_DEPLETED,
+          tier,
+          true,
+        );
+      });
+    }
+  });
+
   describe('snooze', () => {
     let fakeClock: SinonFakeTimers;
 

src/telemetry/api.ts

@@ -66,6 +66,10 @@ export type ColabEvent =
       /** An event representing a file download from a Colab server. */
       download_event: DownloadEvent;
     }
+  | {
+      /** An event representing a low/depleted CCU balance notification. */
+      low_ccu_notification_event: LowCcuNotificationEvent;
+    }
   | {
       /** An event representing an error. */
       error_event: ErrorEvent;
@@ -176,6 +180,31 @@ export enum ContentBrowserTarget {
   TARGET_DIRECTORY = 2,
 }
 
+/** The severity of a Colab Compute Units (CCU) low balance notification. */
+export enum LowBalanceSeverity {
+  SEVERITY_UNSPECIFIED = 0,
+  /**
+   * Balance is low (less than 30 minutes of compute remaining at the current
+   * consumption rate); shown as a warning.
+   */
+  SEVERITY_LOW = 1,
+  /** Balance is fully depleted; shown as an error. */
+  SEVERITY_DEPLETED = 2,
+}
+
+/**
+ * The user's Colab subscription tier as recorded in telemetry. Mirrors the
+ * top-level `SubscriptionTier` proto enum. Distinct from
+ * `colab/api.SubscriptionTier`, which uses unprefixed values (`NONE`,
+ * `PRO`, `PRO_PLUS`).
+ */
+export enum SubscriptionTier {
+  SUBSCRIPTION_TIER_UNSPECIFIED = 0,
+  SUBSCRIPTION_TIER_NONE = 1,
+  SUBSCRIPTION_TIER_PRO = 2,
+  SUBSCRIPTION_TIER_PRO_PLUS = 3,
+}
+
 // The authentication flow used for sign in.
 export enum AuthFlow {
   AUTH_FLOW_UNSPECIFIED = 0,
@@ -273,6 +302,20 @@ interface DownloadEvent {
   downloaded_bytes: number;
 }
 
+/** An event representing a low/depleted CCU balance notification. */
+interface LowCcuNotificationEvent {
+  /** The severity level of the notification. */
+  severity: LowBalanceSeverity;
+  /** The user's subscription tier when the notification was shown. */
+  subscription_tier: SubscriptionTier;
+  /**
+   * Whether the user clicked the call-to-action (e.g., "Sign Up for Colab",
+   * "Upgrade to Pro+", "Purchase More CCUs"). False if the notification was
+   * dismissed without taking action.
+   */
+  clicked_action: boolean;
+}
+
 /** An event representing an error. */
 interface ErrorEvent {
   /** The name of the error. */

src/telemetry/index.ts

@@ -6,7 +6,10 @@
 import assert from 'assert';
 import vscode from 'vscode';
 import { Disposable } from 'vscode';
-import { AuthType } from '../colab/api';
+import {
+  AuthType,
+  SubscriptionTier as ColabSubscriptionTier,
+} from '../colab/api';
 import { COLAB_EXT_IDENTIFIER } from '../config/constants';
 import { getPackageInfo } from '../config/package-info';
 import { JUPYTER_EXT_IDENTIFIER } from '../jupyter/jupyter-extension';
@@ -18,8 +21,10 @@ import {
   AuthFlow,
   ContentBrowserOperation,
   ContentBrowserTarget,
+  LowBalanceSeverity,
   NotebookSource,
   Outcome,
+  SubscriptionTier,
 } from './api';
 import { ClearcutClient } from './client';
 
@@ -115,6 +120,19 @@ export const telemetry = {
       download_event: { outcome, downloaded_bytes: downloadedBytes },
     });
   },
+  logLowCcuNotification: (
+    severity: LowBalanceSeverity,
+    subscriptionTier: ColabSubscriptionTier,
+    clickedAction: boolean,
+  ) => {
+    log({
+      low_ccu_notification_event: {
+        severity,
+        subscription_tier: toTelemetrySubscriptionTier(subscriptionTier),
+        clicked_action: clickedAction,
+      },
+    });
+  },
   logError: (e: unknown) => {
     if (e instanceof Error) {
       log({
@@ -203,3 +221,16 @@ function log(event: ColabEvent) {
     timestamp: new Date().toISOString(),
   });
 }
+
+function toTelemetrySubscriptionTier(
+  tier: ColabSubscriptionTier,
+): SubscriptionTier {
+  switch (tier) {
+    case ColabSubscriptionTier.NONE:
+      return SubscriptionTier.SUBSCRIPTION_TIER_NONE;
+    case ColabSubscriptionTier.PRO:
+      return SubscriptionTier.SUBSCRIPTION_TIER_PRO;
+    case ColabSubscriptionTier.PRO_PLUS:
+      return SubscriptionTier.SUBSCRIPTION_TIER_PRO_PLUS;
+  }
+}

src/telemetry/telemetry.unit.test.ts

@@ -8,7 +8,10 @@ import { expect } from 'chai';
 import sinon, { SinonSpy, SinonFakeTimers } from 'sinon';
 import vscode from 'vscode';
 import { Disposable } from 'vscode';
-import { AuthType } from '../colab/api';
+import {
+  AuthType,
+  SubscriptionTier as ColabSubscriptionTier,
+} from '../colab/api';
 import { COLAB_EXT_IDENTIFIER } from '../config/constants';
 import { JUPYTER_EXT_IDENTIFIER } from '../jupyter/jupyter-extension';
 import { newVsCodeStub, VsCodeStub } from '../test/helpers/vscode';
@@ -19,8 +22,10 @@ import {
   AuthFlow,
   ContentBrowserOperation,
   ContentBrowserTarget,
+  LowBalanceSeverity,
   NotebookSource,
   Outcome,
+  SubscriptionTier,
 } from './api';
 import { ClearcutClient } from './client';
 import { initializeTelemetry, telemetry } from '.';
@@ -324,6 +329,42 @@ describe('Telemetry Module', () => {
       });
     });
 
+    const subscriptionTierCases: {
+      tier: ColabSubscriptionTier;
+      expected: SubscriptionTier;
+    }[] = [
+      {
+        tier: ColabSubscriptionTier.NONE,
+        expected: SubscriptionTier.SUBSCRIPTION_TIER_NONE,
+      },
+      {
+        tier: ColabSubscriptionTier.PRO,
+        expected: SubscriptionTier.SUBSCRIPTION_TIER_PRO,
+      },
+      {
+        tier: ColabSubscriptionTier.PRO_PLUS,
+        expected: SubscriptionTier.SUBSCRIPTION_TIER_PRO_PLUS,
+      },
+    ];
+    for (const { tier, expected } of subscriptionTierCases) {
+      it(`logs on low CCU notification for SubscriptionTier.${ColabSubscriptionTier[tier]}`, () => {
+        telemetry.logLowCcuNotification(
+          LowBalanceSeverity.SEVERITY_LOW,
+          tier,
+          false,
+        );
+
+        sinon.assert.calledOnceWithExactly(logStub, {
+          ...baseLog,
+          low_ccu_notification_event: {
+            severity: LowBalanceSeverity.SEVERITY_LOW,
+            subscription_tier: expected,
+            clicked_action: false,
+          },
+        });
+      });
+    }
+
     it('logs on mount Drive snippet', () => {
       const source = CommandSource.COMMAND_SOURCE_COMMAND_PALETTE;