fix(telemetry): add circuit breaker and shutdown draining to TelemetryService

Author: edelauna

packages/telemetry/src/TelemetryService.ts

@@ -7,14 +7,63 @@ import {
 	type TelemetrySetting,
 } from "@roo-code/types"
 
+/**
+ * Events prone to retry-storm-style repetition (e.g. a broken embedder config
+ * re-triggering on every file-system event). Guarded by a circuit breaker in
+ * `captureEvent` so a single broken install can't flood the Product Analytics
+ * quota. Tracked per-event via a sliding time window, independent of any
+ * other telemetry the same install may also be sending.
+ */
+const CIRCUIT_BREAKER_GUARDED_EVENTS = new Set<TelemetryEventName>([TelemetryEventName.CODE_INDEX_ERROR])
+
+/** Captures of a guarded event within the counting window allowed before the breaker trips. */
+const CIRCUIT_BREAKER_MAX_IN_WINDOW = 50
+
+/** Rolling window over which guarded-event occurrences are counted. */
+const CIRCUIT_BREAKER_WINDOW_MS = 10 * 60 * 1000
+
+/** How long a tripped breaker stays tripped before allowing captures again. */
+const CIRCUIT_BREAKER_COOLDOWN_MS = 10 * 60 * 1000
+
+/**
+ * Upper bound on how long shutdown() will wait for in-flight capture calls to drain.
+ * deactivate() awaits shutdown() before terminal cleanup, so an unbounded wait here
+ * (e.g. a capture stuck on network I/O that never resolves/rejects) would block the
+ * extension host from ever finishing deactivation. Losing an in-flight capture on
+ * timeout is an acceptable tradeoff against blocking shutdown indefinitely.
+ */
+const SHUTDOWN_DRAIN_TIMEOUT_MS = 3000
+
 /**
  * TelemetryService wrapper class that defers initialization.
  * This ensures that we only create the various clients after environment
  * variables are loaded.
  */
 export class TelemetryService {
+	// Timestamps of recent guarded-event occurrences, per event name, oldest first.
+	private guardedEventOccurrences = new Map<TelemetryEventName, number[]>()
+	private trippedUntil = new Map<TelemetryEventName, number>()
+
+	// In-flight client.capture()/captureException() promises. captureEvent/captureException are
+	// synchronous (void-returning) for callers, but the underlying client calls are async (e.g.
+	// PostHogTelemetryClient awaits property enrichment before enqueueing). Tracked here so
+	// shutdown() can drain them before flushing/closing the clients -- otherwise a capture that's
+	// still mid-flight when shutdown() runs could be lost entirely.
+	private pendingClientCalls = new Set<Promise<unknown>>()
+
+	// Set at the start of shutdown() so new captureEvent/captureException calls stop being
+	// tracked (and, once clients are closing, stop being sent) instead of racing the drain.
+	private isShuttingDown = false
+
 	constructor(private clients: TelemetryClient[]) {}
 
+	private trackPendingClientCall(promise: Promise<unknown>): void {
+		// Never let a rejected client call surface as an unhandled rejection or block shutdown.
+		const tracked = promise.catch(() => undefined)
+		this.pendingClientCalls.add(tracked)
+		void tracked.finally(() => this.pendingClientCalls.delete(tracked))
+	}
+
 	public register(client: TelemetryClient): void {
 		this.clients.push(client)
 	}
@@ -51,18 +100,60 @@ export class TelemetryService {
 		this.clients.forEach((client) => client.updateTelemetryState(isOptedIn))
 	}
 
+	/**
+	 * Checks whether a guarded event should be dropped by the circuit breaker,
+	 * updating the breaker's internal state as a side effect. Tracked entirely
+	 * independently of other event names -- unrelated telemetry from the same
+	 * install must never mask (or count towards) a guarded-event burst.
+	 */
+	private shouldDropForCircuitBreaker(eventName: TelemetryEventName): boolean {
+		if (!CIRCUIT_BREAKER_GUARDED_EVENTS.has(eventName)) {
+			return false
+		}
+
+		const now = Date.now()
+
+		const trippedUntil = this.trippedUntil.get(eventName)
+		if (trippedUntil !== undefined) {
+			if (now < trippedUntil) {
+				return true
+			}
+
+			// Cooldown elapsed - reset and allow this capture through.
+			this.trippedUntil.delete(eventName)
+			this.guardedEventOccurrences.delete(eventName)
+		}
+
+		const windowStart = now - CIRCUIT_BREAKER_WINDOW_MS
+		const occurrences = (this.guardedEventOccurrences.get(eventName) ?? []).filter((ts) => ts > windowStart)
+		occurrences.push(now)
+		this.guardedEventOccurrences.set(eventName, occurrences)
+
+		if (occurrences.length > CIRCUIT_BREAKER_MAX_IN_WINDOW) {
+			this.trippedUntil.set(eventName, now + CIRCUIT_BREAKER_COOLDOWN_MS)
+			this.guardedEventOccurrences.delete(eventName)
+			return true
+		}
+
+		return false
+	}
+
 	/**
 	 * Generic method to capture any type of event with specified properties
 	 * @param eventName The event name to capture
 	 * @param properties The event properties
 	 */
 	// eslint-disable-next-line @typescript-eslint/no-explicit-any
 	public captureEvent(eventName: TelemetryEventName, properties?: Record<string, any>): void {
-		if (!this.isReady) {
+		if (!this.isReady || this.isShuttingDown) {
+			return
+		}
+
+		if (this.shouldDropForCircuitBreaker(eventName)) {
 			return
 		}
 
-		this.clients.forEach((client) => client.capture({ event: eventName, properties }))
+		this.clients.forEach((client) => this.trackPendingClientCall(client.capture({ event: eventName, properties })))
 	}
 
 	/**
@@ -71,11 +162,13 @@ export class TelemetryService {
 	 * @param additionalProperties Additional properties to include with the exception
 	 */
 	public captureException(error: Error, additionalProperties?: Record<string, unknown>): void {
-		if (!this.isReady) {
+		if (!this.isReady || this.isShuttingDown) {
 			return
 		}
 
-		this.clients.forEach((client) => client.captureException(error, additionalProperties))
+		this.clients.forEach((client) =>
+			this.trackPendingClientCall(client.captureException(error, additionalProperties)),
+		)
 	}
 
 	public captureTaskCreated(taskId: string): void {
@@ -263,7 +356,26 @@ export class TelemetryService {
 			return
 		}
 
-		this.clients.forEach((client) => client.shutdown())
+		// Stop accepting new captures immediately, before draining -- otherwise a steady trickle
+		// of new calls (e.g. from a teardown-time error handler) could keep pendingClientCalls
+		// non-empty indefinitely and the drain loop below would never terminate on its own.
+		this.isShuttingDown = true
+
+		// Drain any in-flight capture/captureException calls first, so a client's shutdown()
+		// (which flushes its queue) can't run ahead of a capture that hasn't been enqueued yet.
+		// Loop rather than a single snapshot: a call already in flight when draining started may
+		// itself still be tracked by the time we check again. Bounded by a timeout so a capture
+		// stuck on network I/O that never resolves/rejects can't block deactivate() forever --
+		// losing that one capture is an acceptable tradeoff against hanging terminal cleanup.
+		const drainStart = Date.now()
+		while (this.pendingClientCalls.size > 0 && Date.now() - drainStart < SHUTDOWN_DRAIN_TIMEOUT_MS) {
+			await Promise.race([
+				Promise.all(this.pendingClientCalls),
+				new Promise((resolve) => setTimeout(resolve, SHUTDOWN_DRAIN_TIMEOUT_MS - (Date.now() - drainStart))),
+			])
+		}
+
+		await Promise.all(this.clients.map((client) => client.shutdown()))
 	}
 
 	private static _instance: TelemetryService | null = null

packages/telemetry/src/__tests__/TelemetryService.circuit-breaker.test.ts

@@ -0,0 +1,124 @@
+// pnpm --filter @roo-code/telemetry test src/__tests__/TelemetryService.circuit-breaker.test.ts
+
+import { TelemetryEventName, type TelemetryClient } from "@roo-code/types"
+
+import { TelemetryService } from "../TelemetryService"
+
+describe("TelemetryService circuit breaker", () => {
+	let mockClient: TelemetryClient
+
+	beforeEach(() => {
+		vi.useFakeTimers()
+		vi.setSystemTime(0)
+
+		mockClient = {
+			setProvider: vi.fn(),
+			capture: vi.fn().mockResolvedValue(undefined),
+			captureException: vi.fn().mockResolvedValue(undefined),
+			updateTelemetryState: vi.fn(),
+			isTelemetryEnabled: vi.fn().mockReturnValue(true),
+			shutdown: vi.fn().mockResolvedValue(undefined),
+		}
+	})
+
+	afterEach(() => {
+		vi.useRealTimers()
+	})
+
+	it("passes through captures under the trip threshold", () => {
+		const service = new TelemetryService([mockClient])
+
+		for (let i = 0; i < 49; i++) {
+			service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i })
+		}
+
+		expect(mockClient.capture).toHaveBeenCalledTimes(49)
+	})
+
+	it("trips after 50 CODE_INDEX_ERROR captures within the window and drops further ones", () => {
+		const service = new TelemetryService([mockClient])
+
+		for (let i = 0; i < 50; i++) {
+			service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i })
+		}
+		expect(mockClient.capture).toHaveBeenCalledTimes(50)
+
+		// 51st capture should be dropped - breaker has tripped.
+		service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i: 50 })
+		expect(mockClient.capture).toHaveBeenCalledTimes(50)
+
+		// Keeps dropping while tripped.
+		service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i: 51 })
+		expect(mockClient.capture).toHaveBeenCalledTimes(50)
+	})
+
+	it("re-allows captures after the cooldown window elapses", () => {
+		const service = new TelemetryService([mockClient])
+
+		for (let i = 0; i < 50; i++) {
+			service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i })
+		}
+		service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i: 50 })
+		expect(mockClient.capture).toHaveBeenCalledTimes(50)
+
+		// Just under 10 minutes - still tripped.
+		vi.setSystemTime(10 * 60 * 1000 - 1)
+		service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i: 51 })
+		expect(mockClient.capture).toHaveBeenCalledTimes(50)
+
+		// Cooldown elapsed - one more error gets through.
+		vi.setSystemTime(10 * 60 * 1000)
+		service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i: 52 })
+		expect(mockClient.capture).toHaveBeenCalledTimes(51)
+	})
+
+	it("does not reset the guarded count when unrelated events are interleaved", () => {
+		// A real broken install still does normal things (creates/completes other tasks)
+		// while a subsystem like code-index is stuck in a retry loop. Unrelated telemetry
+		// must not mask the CODE_INDEX_ERROR burst by resetting its count.
+		const service = new TelemetryService([mockClient])
+
+		for (let i = 0; i < 25; i++) {
+			service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i })
+			service.captureEvent(TelemetryEventName.TASK_CREATED, { taskId: `task-${i}` })
+		}
+		// 25 CODE_INDEX_ERROR so far - still under the threshold of 50.
+		expect(mockClient.capture).toHaveBeenCalledTimes(25 + 25)
+
+		for (let i = 25; i < 50; i++) {
+			service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i })
+			service.captureEvent(TelemetryEventName.TASK_CREATED, { taskId: `task-${i}` })
+		}
+		// 50th CODE_INDEX_ERROR trips the breaker; TASK_CREATED events are never guarded.
+		expect(mockClient.capture).toHaveBeenCalledTimes(50 + 50)
+
+		// Further CODE_INDEX_ERROR captures are dropped even though unrelated events keep flowing.
+		service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i: 50 })
+		service.captureEvent(TelemetryEventName.TASK_CREATED, { taskId: "task-50" })
+		expect(mockClient.capture).toHaveBeenCalledTimes(50 + 51)
+	})
+
+	it("expires old occurrences outside the counting window instead of trapping the breaker open forever", () => {
+		// A slow trickle of CODE_INDEX_ERROR (below the burst rate) should never trip the
+		// breaker, since old occurrences age out of the window rather than accumulating forever.
+		const service = new TelemetryService([mockClient])
+
+		for (let i = 0; i < 60; i++) {
+			service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i })
+			// Advance well past the counting window between each one.
+			vi.setSystemTime(Date.now() + 60 * 1000)
+		}
+
+		expect(mockClient.capture).toHaveBeenCalledTimes(60)
+	})
+
+	it("does not guard other event names", () => {
+		const service = new TelemetryService([mockClient])
+
+		for (let i = 0; i < 200; i++) {
+			service.captureEvent(TelemetryEventName.TOOL_USED, { tool: "read_file" })
+		}
+
+		expect(mockClient.capture).toHaveBeenCalledTimes(200)
+	})
+})