feat: log telemetry for low/depleted CCU balance notifications

Captures the severity of the notification (`SEVERITY_LOW` warning or
`SEVERITY_DEPLETED` error), the user's subscription tier, and
whether the user clicked the call-to-action (vs. dismissed the
notification).

The event is logged after the notification's promise resolves so the
`clicked_action` outcome is captured. Pairs with
`upgrade_to_pro_event` to enable funnel analysis from notification
to upgrade.

`telemetry.logLowCcuNotification` accepts the Colab-level
`SubscriptionTier` directly and converts it to the proto-shaped
top-level `SubscriptionTier` enum internally, keeping callers free
of the conversion.

Author: kevineger

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,111 @@ 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,
+      );
+    });
+
+    it('logs SEVERITY_DEPLETED with clicked=true when error action is clicked', async () => {
+      const ccuInfo = createCcuInfo(
+        { paidMinutes: 0, freeMinutes: 0 },
+        SubscriptionTier.NONE,
+      );
+      autoRespond('error', 'Sign Up for Colab');
+
+      const noOp = consumptionNotifier.nextConsumptionCalculation();
+      ccuEmitter.fire(ccuInfo);
+      await noOp;
+
+      sinon.assert.calledOnceWithExactly(
+        logStub,
+        LowBalanceSeverity.SEVERITY_DEPLETED,
+        SubscriptionTier.NONE,
+        true,
+      );
+    });
+
+    it('plumbs the Pro subscription tier through unchanged', async () => {
+      const ccuInfo = createCcuInfo(
+        { paidMinutes: 0, freeMinutes: 0 },
+        SubscriptionTier.PRO,
+      );
+      autoRespond('error', 'Upgrade to Pro+');
+
+      const noOp = consumptionNotifier.nextConsumptionCalculation();
+      ccuEmitter.fire(ccuInfo);
+      await noOp;
+
+      sinon.assert.calledOnceWithExactly(
+        logStub,
+        LowBalanceSeverity.SEVERITY_DEPLETED,
+        SubscriptionTier.PRO,
+        true,
+      );
+    });
+
+    it('plumbs the Pro+ subscription tier through unchanged', async () => {
+      const ccuInfo = createCcuInfo(
+        { paidMinutes: 0, freeMinutes: 0 },
+        SubscriptionTier.PRO_PLUS,
+      );
+      autoRespond('error', 'Purchase More CCUs');
+
+      const noOp = consumptionNotifier.nextConsumptionCalculation();
+      ccuEmitter.fire(ccuInfo);
+      await noOp;
+
+      sinon.assert.calledOnceWithExactly(
+        logStub,
+        LowBalanceSeverity.SEVERITY_DEPLETED,
+        SubscriptionTier.PRO_PLUS,
+        true,
+      );
+    });
+  });
+
   describe('snooze', () => {
     let fakeClock: SinonFakeTimers;
 

src/telemetry/api.ts

@@ -58,6 +58,10 @@ export type ColabEvent =
       /** An event representing a Colab toolbar click. */
       colab_toolbar_event: ColabToolbarEvent;
     }
+  | {
+      /** An event representing a low/depleted CCU balance notification. */
+      low_ccu_notification_event: LowCcuNotificationEvent;
+    }
   | {
       /** An event representing an error. */
       error_event: ErrorEvent;
@@ -119,6 +123,31 @@ export enum NotebookSource {
   NOTEBOOK_SOURCE_DRIVE = 1,
 }
 
+/** 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,
@@ -140,6 +169,20 @@ type AutoConnectEvent = Record<string, never>;
 /** An event representing a Colab toolbar click. */
 type ColabToolbarEvent = Record<string, never>;
 
+/** 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';
@@ -15,7 +18,9 @@ import {
   ColabEvent,
   CommandSource,
   AuthFlow,
+  LowBalanceSeverity,
   NotebookSource,
+  SubscriptionTier,
 } from './api';
 import { ClearcutClient } from './client';
 
@@ -79,6 +84,19 @@ export const telemetry = {
   logColabToolbar: () => {
     log({ colab_toolbar_event: {} });
   },
+  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({
@@ -141,3 +159,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,15 +8,20 @@ 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';
 import {
   ColabLogEventBase,
   CommandSource,
   AuthFlow,
+  LowBalanceSeverity,
   NotebookSource,
+  SubscriptionTier,
 } from './api';
 import { ClearcutClient } from './client';
 import { initializeTelemetry, telemetry } from '.';
@@ -275,6 +280,52 @@ describe('Telemetry Module', () => {
       });
     });
 
+    it('logs on low CCU notification', () => {
+      telemetry.logLowCcuNotification(
+        LowBalanceSeverity.SEVERITY_DEPLETED,
+        ColabSubscriptionTier.PRO,
+        true,
+      );
+
+      sinon.assert.calledOnceWithExactly(logStub, {
+        ...baseLog,
+        low_ccu_notification_event: {
+          severity: LowBalanceSeverity.SEVERITY_DEPLETED,
+          subscription_tier: SubscriptionTier.SUBSCRIPTION_TIER_PRO,
+          clicked_action: true,
+        },
+      });
+    });
+
+    it('converts each Colab SubscriptionTier to its proto-shaped value', () => {
+      const cases: [ColabSubscriptionTier, SubscriptionTier][] = [
+        [ColabSubscriptionTier.NONE, SubscriptionTier.SUBSCRIPTION_TIER_NONE],
+        [ColabSubscriptionTier.PRO, SubscriptionTier.SUBSCRIPTION_TIER_PRO],
+        [
+          ColabSubscriptionTier.PRO_PLUS,
+          SubscriptionTier.SUBSCRIPTION_TIER_PRO_PLUS,
+        ],
+      ];
+      for (const [tier, expected] of cases) {
+        logStub.resetHistory();
+
+        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;