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

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

* fix(telemetry): fix circuit breaker off-by-one, allSettled on shutdown

* fix(telemetry): bound client.shutdown() phase, strengthen breaker test

Author: edelauna

packages/telemetry/src/TelemetryService.ts

@@ -7,14 +7,66 @@ 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 applied separately to each phase of shutdown(): draining in-flight capture
+ * calls, then awaiting client.shutdown(). deactivate() awaits shutdown() before terminal
+ * cleanup, so an unbounded wait in either phase (e.g. a capture stuck on network I/O, or
+ * a client's own shutdown() -- posthog-node defaults to a 30s internal timeout -- never
+ * settling) would block the extension host from ever finishing deactivation. Losing an
+ * in-flight capture, or a client's graceful flush, on timeout is an acceptable tradeoff
+ * against blocking shutdown indefinitely. Worst case, shutdown() takes up to roughly
+ * 2 * SHUTDOWN_PHASE_TIMEOUT_MS.
+ */
+const SHUTDOWN_PHASE_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 +103,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 +165,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 +359,35 @@ 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_PHASE_TIMEOUT_MS) {
+			await Promise.race([
+				Promise.all(this.pendingClientCalls),
+				new Promise((resolve) => setTimeout(resolve, SHUTDOWN_PHASE_TIMEOUT_MS - (Date.now() - drainStart))),
+			])
+		}
+
+		// Bound client shutdown the same way as the drain above: posthog-node's own shutdown()
+		// defaults to a 30s internal timeout when called with no argument (as PostHogTelemetryClient
+		// does), and TelemetryClient#shutdown() takes no timeout parameter to pass one through. Racing
+		// against our own timer here, instead of just awaiting client.shutdown() directly, keeps
+		// deactivate() from blocking for up to 30s on a client that never settles.
+		// allSettled, not all: one client rejecting must not stop us from awaiting the others.
+		await Promise.race([
+			Promise.allSettled(this.clients.map((client) => client.shutdown())),
+			new Promise((resolve) => setTimeout(resolve, SHUTDOWN_PHASE_TIMEOUT_MS)),
+		])
 	}
 
 	private static _instance: TelemetryService | null = null

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

@@ -0,0 +1,148 @@
+// 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 at the 50th CODE_INDEX_ERROR capture within the window and drops further ones", () => {
+		const service = new TelemetryService([mockClient])
+
+		for (let i = 0; i < 49; i++) {
+			service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i })
+		}
+		expect(mockClient.capture).toHaveBeenCalledTimes(49)
+
+		// 50th capture trips the breaker but is itself dropped.
+		service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i: 49 })
+		expect(mockClient.capture).toHaveBeenCalledTimes(49)
+
+		// Keeps dropping while tripped.
+		service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i: 50 })
+		expect(mockClient.capture).toHaveBeenCalledTimes(49)
+	})
+
+	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 })
+		}
+		expect(mockClient.capture).toHaveBeenCalledTimes(49)
+
+		// Just under 10 minutes - still tripped.
+		vi.setSystemTime(10 * 60 * 1000 - 1)
+		service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i: 50 })
+		expect(mockClient.capture).toHaveBeenCalledTimes(49)
+
+		// Cooldown elapsed - one more error gets through.
+		vi.setSystemTime(10 * 60 * 1000)
+		service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i: 51 })
+		expect(mockClient.capture).toHaveBeenCalledTimes(50)
+	})
+
+	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 < 49; i++) {
+			service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i })
+			service.captureEvent(TelemetryEventName.TASK_CREATED, { taskId: `task-${i}` })
+		}
+		// 49 CODE_INDEX_ERROR so far - still under the threshold of 50.
+		expect(mockClient.capture).toHaveBeenCalledTimes(49 + 49)
+
+		// 50th CODE_INDEX_ERROR trips the breaker (and is itself dropped); TASK_CREATED
+		// events are never guarded.
+		service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i: 49 })
+		service.captureEvent(TelemetryEventName.TASK_CREATED, { taskId: "task-49" })
+		expect(mockClient.capture).toHaveBeenCalledTimes(49 + 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(49 + 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)
+	})
+
+	it("returns early on the not-ready (zero-client) branch without touching circuit breaker state", () => {
+		// captureEvent's !this.isReady check runs before shouldDropForCircuitBreaker, so with
+		// no clients registered, guarded-event bookkeeping should never be reached.
+		const service = new TelemetryService([])
+
+		for (let i = 0; i < 50; i++) {
+			service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i })
+		}
+
+		expect(mockClient.capture).not.toHaveBeenCalled()
+
+		// Prove the 50 zero-client calls didn't count towards the breaker: once a client is
+		// registered, the very next capture must still go through instead of being dropped.
+		service.register(mockClient)
+		service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i: 50 })
+
+		expect(mockClient.capture).toHaveBeenCalledTimes(1)
+	})
+})