From ea78a27b0a46cea3d99bed68ea58c6872aa0880a Mon Sep 17 00:00:00 2001 From: Elliott de Launay Date: Fri, 31 Jul 2026 03:22:14 +0000 Subject: [PATCH 1/3] fix(telemetry): add circuit breaker and shutdown draining to TelemetryService --- packages/telemetry/src/TelemetryService.ts | 122 ++++++++++- .../TelemetryService.circuit-breaker.test.ts | 124 +++++++++++ .../TelemetryService.shutdown.test.ts | 195 ++++++++++++++++++ src/extension.ts | 10 +- 4 files changed, 445 insertions(+), 6 deletions(-) create mode 100644 packages/telemetry/src/__tests__/TelemetryService.circuit-breaker.test.ts create mode 100644 packages/telemetry/src/__tests__/TelemetryService.shutdown.test.ts diff --git a/packages/telemetry/src/TelemetryService.ts b/packages/telemetry/src/TelemetryService.ts index 8eb1ed0ab6..d4090018b7 100644 --- a/packages/telemetry/src/TelemetryService.ts +++ b/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.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() + private trippedUntil = new Map() + + // 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>() + + // 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): 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,6 +100,44 @@ 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 @@ -58,11 +145,15 @@ export class TelemetryService { */ // eslint-disable-next-line @typescript-eslint/no-explicit-any public captureEvent(eventName: TelemetryEventName, properties?: Record): 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): 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 diff --git a/packages/telemetry/src/__tests__/TelemetryService.circuit-breaker.test.ts b/packages/telemetry/src/__tests__/TelemetryService.circuit-breaker.test.ts new file mode 100644 index 0000000000..53ebaaebb1 --- /dev/null +++ b/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) + }) +}) diff --git a/packages/telemetry/src/__tests__/TelemetryService.shutdown.test.ts b/packages/telemetry/src/__tests__/TelemetryService.shutdown.test.ts new file mode 100644 index 0000000000..3e765cadc1 --- /dev/null +++ b/packages/telemetry/src/__tests__/TelemetryService.shutdown.test.ts @@ -0,0 +1,195 @@ +// pnpm --filter @roo-code/telemetry test src/__tests__/TelemetryService.shutdown.test.ts + +import { TelemetryEventName, type TelemetryClient } from "@roo-code/types" + +import { TelemetryService } from "../TelemetryService" + +describe("TelemetryService.shutdown draining", () => { + it("awaits in-flight captureEvent calls before shutting down clients", async () => { + let resolveCapture!: () => void + const capturePromise = new Promise((resolve) => { + resolveCapture = resolve + }) + + const captureOrder: string[] = [] + + const mockClient: TelemetryClient = { + setProvider: vi.fn(), + capture: vi.fn().mockImplementation(async () => { + await capturePromise + captureOrder.push("captured") + }), + captureException: vi.fn(), + updateTelemetryState: vi.fn(), + isTelemetryEnabled: vi.fn().mockReturnValue(true), + shutdown: vi.fn().mockImplementation(async () => { + captureOrder.push("shutdown") + }), + } + + const service = new TelemetryService([mockClient]) + + // Fire a capture whose underlying client.capture() promise hasn't resolved yet. + service.captureEvent(TelemetryEventName.TASK_CREATED, { taskId: "abc" }) + + const shutdownPromise = service.shutdown() + + // Capture is still pending - shutdown must not have run yet. + expect(captureOrder).toEqual([]) + + resolveCapture() + await shutdownPromise + + // The in-flight capture must complete before the client is shut down. + expect(captureOrder).toEqual(["captured", "shutdown"]) + }) + + it("drains a call already queued before shutdown() started, even across multiple drain passes", async () => { + // Regression test: shutdown() must not take a single Promise.all snapshot of + // pendingClientCalls. A promise added to pendingClientCalls *after* Promise.all(set) + // has already been called is never awaited by that call, even if it never resolves + // -- Promise.all takes its list of promises to await synchronously, at call time. So a + // capture that was already in flight (tracked in pendingClientCalls) when shutdown() + // took its first Promise.all snapshot, but whose *own* async chain enqueues more work + // tracked via a fresh pendingClientCalls entry, must still be drained by a later pass + // of the loop rather than being silently dropped by a single-pass drain. + let resolveFirstCapture!: () => void + const firstCapturePromise = new Promise((resolve) => { + resolveFirstCapture = resolve + }) + + let resolveSecondCapture!: () => void + const secondCapturePromise = new Promise((resolve) => { + resolveSecondCapture = resolve + }) + + const captureOrder: string[] = [] + let firstCaptureStarted = false + + const mockClient: TelemetryClient = { + setProvider: vi.fn(), + capture: vi.fn().mockImplementation(async () => { + if (!firstCaptureStarted) { + firstCaptureStarted = true + await firstCapturePromise + captureOrder.push("first-captured") + return + } + + await secondCapturePromise + captureOrder.push("second-captured") + }), + captureException: vi.fn(), + updateTelemetryState: vi.fn(), + isTelemetryEnabled: vi.fn().mockReturnValue(true), + shutdown: vi.fn().mockImplementation(async () => { + captureOrder.push("shutdown") + }), + } + + const service = new TelemetryService([mockClient]) + + // Both captures are fired *before* shutdown() is called, so both are legitimately + // in flight (tracked in pendingClientCalls) at the moment shutdown() takes its first + // snapshot -- unlike a capture fired after shutdown() has started, which is expected + // to be gated out instead (see the "stops accepting new captures" test below). + service.captureEvent(TelemetryEventName.TASK_CREATED, { taskId: "first" }) + service.captureEvent(TelemetryEventName.TASK_CREATED, { taskId: "second" }) + + const shutdownPromise = service.shutdown() + + resolveFirstCapture() + // Flush several microtask ticks so a buggy single-pass drain has every opportunity + // to run client.shutdown() before we check -- a couple of ticks isn't enough since + // the mocked capture/shutdown chain itself spans a few microtask hops. + for (let i = 0; i < 4; i++) { + await Promise.resolve() + } + + // The second capture is still pending (its own resolver hasn't fired) -- shutdown + // must not have completed yet, otherwise it dropped the second capture. + expect(captureOrder).not.toContain("shutdown") + + resolveSecondCapture() + await shutdownPromise + + expect(mockClient.capture).toHaveBeenCalledTimes(2) + expect(captureOrder).toEqual(["first-captured", "second-captured", "shutdown"]) + }) + + it("stops accepting new captures once shutdown() has started", async () => { + // Finding #4: shutdown() must mark itself as shutting down before draining, so a + // steady trickle of new captures firing after shutdown() has begun (e.g. from a + // teardown-time error handler) can't keep pendingClientCalls non-empty forever and + // prevent the drain loop from ever terminating on its own. + const mockClient: TelemetryClient = { + setProvider: vi.fn(), + capture: vi.fn().mockResolvedValue(undefined), + captureException: vi.fn(), + updateTelemetryState: vi.fn(), + isTelemetryEnabled: vi.fn().mockReturnValue(true), + shutdown: vi.fn().mockResolvedValue(undefined), + } + + const service = new TelemetryService([mockClient]) + + const shutdownPromise = service.shutdown() + + // Fired after shutdown() has already started -- must be dropped, not tracked/drained. + service.captureEvent(TelemetryEventName.TASK_CREATED, { taskId: "late" }) + + await shutdownPromise + + expect(mockClient.capture).not.toHaveBeenCalled() + }) + + it("does not hang forever when a capture never resolves, bounded by the drain timeout", async () => { + vi.useFakeTimers() + try { + const mockClient: TelemetryClient = { + setProvider: vi.fn(), + // Never resolves -- simulates a capture stuck on network I/O. + capture: vi.fn().mockImplementation(() => new Promise(() => {})), + captureException: vi.fn(), + updateTelemetryState: vi.fn(), + isTelemetryEnabled: vi.fn().mockReturnValue(true), + shutdown: vi.fn().mockResolvedValue(undefined), + } + + const service = new TelemetryService([mockClient]) + + service.captureEvent(TelemetryEventName.TASK_CREATED, { taskId: "stuck" }) + + const shutdownPromise = service.shutdown() + let settled = false + void shutdownPromise.then(() => { + settled = true + }) + + await vi.advanceTimersByTimeAsync(3000) + + expect(settled).toBe(true) + expect(mockClient.shutdown).toHaveBeenCalledTimes(1) + } finally { + vi.useRealTimers() + } + }) + + it("does not let a rejected capture prevent shutdown from completing", async () => { + const mockClient: TelemetryClient = { + setProvider: vi.fn(), + capture: vi.fn().mockRejectedValue(new Error("capture failed")), + captureException: vi.fn(), + updateTelemetryState: vi.fn(), + isTelemetryEnabled: vi.fn().mockReturnValue(true), + shutdown: vi.fn().mockResolvedValue(undefined), + } + + const service = new TelemetryService([mockClient]) + + service.captureEvent(TelemetryEventName.TASK_CREATED, { taskId: "abc" }) + + await expect(service.shutdown()).resolves.toBeUndefined() + expect(mockClient.shutdown).toHaveBeenCalledTimes(1) + }) +}) diff --git a/src/extension.ts b/src/extension.ts index f006b4b52b..b880bee410 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -388,7 +388,15 @@ export async function deactivate() { } await McpServerManager.cleanup(extensionContext) - TelemetryService.instance.shutdown() + + try { + await TelemetryService.instance.shutdown() + } catch (error) { + outputChannel.appendLine( + `Failed to shut down telemetry service: ${error instanceof Error ? error.message : String(error)}`, + ) + } + Terminal.setTerminalProfile(undefined) TerminalRegistry.cleanup() } From 32bd50b393ccddb1fe4adee3699f87ff86ed4876 Mon Sep 17 00:00:00 2001 From: Elliott de Launay Date: Sat, 1 Aug 2026 02:03:33 +0000 Subject: [PATCH 2/3] fix(telemetry): fix circuit breaker off-by-one, allSettled on shutdown --- packages/telemetry/src/TelemetryService.ts | 5 +- .../TelemetryService.circuit-breaker.test.ts | 53 ++++++++++------ .../TelemetryService.shutdown.test.ts | 63 +++++++++++++++++++ 3 files changed, 101 insertions(+), 20 deletions(-) diff --git a/packages/telemetry/src/TelemetryService.ts b/packages/telemetry/src/TelemetryService.ts index d4090018b7..cabdbfd311 100644 --- a/packages/telemetry/src/TelemetryService.ts +++ b/packages/telemetry/src/TelemetryService.ts @@ -129,7 +129,7 @@ export class TelemetryService { occurrences.push(now) this.guardedEventOccurrences.set(eventName, occurrences) - if (occurrences.length > CIRCUIT_BREAKER_MAX_IN_WINDOW) { + if (occurrences.length >= CIRCUIT_BREAKER_MAX_IN_WINDOW) { this.trippedUntil.set(eventName, now + CIRCUIT_BREAKER_COOLDOWN_MS) this.guardedEventOccurrences.delete(eventName) return true @@ -375,7 +375,8 @@ export class TelemetryService { ]) } - await Promise.all(this.clients.map((client) => client.shutdown())) + // allSettled, not all: one client rejecting must not stop us from awaiting the others. + await Promise.allSettled(this.clients.map((client) => client.shutdown())) } private static _instance: TelemetryService | null = null diff --git a/packages/telemetry/src/__tests__/TelemetryService.circuit-breaker.test.ts b/packages/telemetry/src/__tests__/TelemetryService.circuit-breaker.test.ts index 53ebaaebb1..b606306ea6 100644 --- a/packages/telemetry/src/__tests__/TelemetryService.circuit-breaker.test.ts +++ b/packages/telemetry/src/__tests__/TelemetryService.circuit-breaker.test.ts @@ -35,21 +35,21 @@ describe("TelemetryService circuit breaker", () => { expect(mockClient.capture).toHaveBeenCalledTimes(49) }) - it("trips after 50 CODE_INDEX_ERROR captures within the window and drops further ones", () => { + 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 < 50; i++) { + for (let i = 0; i < 49; i++) { service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i }) } - expect(mockClient.capture).toHaveBeenCalledTimes(50) + expect(mockClient.capture).toHaveBeenCalledTimes(49) - // 51st capture should be dropped - breaker has tripped. - service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i: 50 }) - expect(mockClient.capture).toHaveBeenCalledTimes(50) + // 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: 51 }) - expect(mockClient.capture).toHaveBeenCalledTimes(50) + service.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, { i: 50 }) + expect(mockClient.capture).toHaveBeenCalledTimes(49) }) it("re-allows captures after the cooldown window elapses", () => { @@ -58,18 +58,17 @@ describe("TelemetryService circuit breaker", () => { 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) + expect(mockClient.capture).toHaveBeenCalledTimes(49) // 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) + 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: 52 }) - expect(mockClient.capture).toHaveBeenCalledTimes(51) + 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", () => { @@ -85,17 +84,23 @@ describe("TelemetryService circuit breaker", () => { // 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++) { + for (let i = 25; i < 49; 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) + // 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(50 + 51) + expect(mockClient.capture).toHaveBeenCalledTimes(49 + 51) }) it("expires old occurrences outside the counting window instead of trapping the breaker open forever", () => { @@ -121,4 +126,16 @@ describe("TelemetryService circuit breaker", () => { 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() + }) }) diff --git a/packages/telemetry/src/__tests__/TelemetryService.shutdown.test.ts b/packages/telemetry/src/__tests__/TelemetryService.shutdown.test.ts index 3e765cadc1..89ce05574d 100644 --- a/packages/telemetry/src/__tests__/TelemetryService.shutdown.test.ts +++ b/packages/telemetry/src/__tests__/TelemetryService.shutdown.test.ts @@ -44,6 +44,46 @@ describe("TelemetryService.shutdown draining", () => { expect(captureOrder).toEqual(["captured", "shutdown"]) }) + it("awaits in-flight captureException calls before shutting down clients", async () => { + let resolveCapture!: () => void + const capturePromise = new Promise((resolve) => { + resolveCapture = resolve + }) + + const captureOrder: string[] = [] + + const mockClient: TelemetryClient = { + setProvider: vi.fn(), + capture: vi.fn(), + captureException: vi.fn().mockImplementation(async () => { + await capturePromise + captureOrder.push("captured") + }), + updateTelemetryState: vi.fn(), + isTelemetryEnabled: vi.fn().mockReturnValue(true), + shutdown: vi.fn().mockImplementation(async () => { + captureOrder.push("shutdown") + }), + } + + const service = new TelemetryService([mockClient]) + + // Fire a captureException whose underlying client.captureException() promise hasn't + // resolved yet -- follows the same trackPendingClientCall path as captureEvent. + service.captureException(new Error("boom")) + + const shutdownPromise = service.shutdown() + + // Capture is still pending - shutdown must not have run yet. + expect(captureOrder).toEqual([]) + + resolveCapture() + await shutdownPromise + + // The in-flight capture must complete before the client is shut down. + expect(captureOrder).toEqual(["captured", "shutdown"]) + }) + it("drains a call already queued before shutdown() started, even across multiple drain passes", async () => { // Regression test: shutdown() must not take a single Promise.all snapshot of // pendingClientCalls. A promise added to pendingClientCalls *after* Promise.all(set) @@ -175,6 +215,29 @@ describe("TelemetryService.shutdown draining", () => { } }) + it("calling shutdown() twice shuts down clients twice (no re-entrancy guard)", async () => { + // Documents current behavior: shutdown() has no guard against being called more than + // once. deactivate() only calls it once, so this isn't a bug to fix here, but a second + // call re-running the (now-trivial, since pendingClientCalls is empty) drain loop and + // calling client.shutdown() again should be an explicit, intentional outcome rather than + // an untested one. + const mockClient: TelemetryClient = { + setProvider: vi.fn(), + capture: vi.fn().mockResolvedValue(undefined), + captureException: vi.fn(), + updateTelemetryState: vi.fn(), + isTelemetryEnabled: vi.fn().mockReturnValue(true), + shutdown: vi.fn().mockResolvedValue(undefined), + } + + const service = new TelemetryService([mockClient]) + + await service.shutdown() + await service.shutdown() + + expect(mockClient.shutdown).toHaveBeenCalledTimes(2) + }) + it("does not let a rejected capture prevent shutdown from completing", async () => { const mockClient: TelemetryClient = { setProvider: vi.fn(), From 00589993c100eb8f0173e552b1225796e9bcd3a6 Mon Sep 17 00:00:00 2001 From: Elliott de Launay Date: Sat, 1 Aug 2026 03:19:58 +0000 Subject: [PATCH 3/3] fix(telemetry): bound client.shutdown() phase, strengthen breaker test --- packages/telemetry/src/TelemetryService.ts | 29 ++++++++++----- .../TelemetryService.circuit-breaker.test.ts | 7 ++++ .../TelemetryService.shutdown.test.ts | 37 +++++++++++++++++++ 3 files changed, 64 insertions(+), 9 deletions(-) diff --git a/packages/telemetry/src/TelemetryService.ts b/packages/telemetry/src/TelemetryService.ts index cabdbfd311..a327799fc5 100644 --- a/packages/telemetry/src/TelemetryService.ts +++ b/packages/telemetry/src/TelemetryService.ts @@ -26,13 +26,16 @@ const CIRCUIT_BREAKER_WINDOW_MS = 10 * 60 * 1000 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. + * 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_DRAIN_TIMEOUT_MS = 3000 +const SHUTDOWN_PHASE_TIMEOUT_MS = 3000 /** * TelemetryService wrapper class that defers initialization. @@ -368,15 +371,23 @@ export class TelemetryService { // 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) { + 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_DRAIN_TIMEOUT_MS - (Date.now() - drainStart))), + 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.allSettled(this.clients.map((client) => client.shutdown())) + 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 diff --git a/packages/telemetry/src/__tests__/TelemetryService.circuit-breaker.test.ts b/packages/telemetry/src/__tests__/TelemetryService.circuit-breaker.test.ts index b606306ea6..6807f22bd5 100644 --- a/packages/telemetry/src/__tests__/TelemetryService.circuit-breaker.test.ts +++ b/packages/telemetry/src/__tests__/TelemetryService.circuit-breaker.test.ts @@ -137,5 +137,12 @@ describe("TelemetryService circuit breaker", () => { } 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) }) }) diff --git a/packages/telemetry/src/__tests__/TelemetryService.shutdown.test.ts b/packages/telemetry/src/__tests__/TelemetryService.shutdown.test.ts index 89ce05574d..d710abdc44 100644 --- a/packages/telemetry/src/__tests__/TelemetryService.shutdown.test.ts +++ b/packages/telemetry/src/__tests__/TelemetryService.shutdown.test.ts @@ -215,6 +215,43 @@ describe("TelemetryService.shutdown draining", () => { } }) + it("does not hang forever when client.shutdown() itself never settles", async () => { + // A capture-drain timeout alone isn't enough: posthog-node's shutdown() defaults to a + // 30s internal timeout when called with no argument, and TelemetryClient#shutdown() has + // no way to pass a shorter one through. TelemetryService.shutdown() must apply its own + // bound to this phase too, or deactivate() could block for up to 30s on a client whose + // own shutdown() never settles. + vi.useFakeTimers() + try { + const mockClient: TelemetryClient = { + setProvider: vi.fn(), + capture: vi.fn().mockResolvedValue(undefined), + captureException: vi.fn(), + updateTelemetryState: vi.fn(), + isTelemetryEnabled: vi.fn().mockReturnValue(true), + // Never resolves -- simulates posthog-node's shutdown() stuck past its own timeout. + shutdown: vi.fn().mockImplementation(() => new Promise(() => {})), + } + + const service = new TelemetryService([mockClient]) + + const shutdownPromise = service.shutdown() + let settled = false + void shutdownPromise.then(() => { + settled = true + }) + + // No pending captures, so the drain phase resolves immediately; only the + // client-shutdown phase's own timeout should be needed here. + await vi.advanceTimersByTimeAsync(3000) + + expect(settled).toBe(true) + expect(mockClient.shutdown).toHaveBeenCalledTimes(1) + } finally { + vi.useRealTimers() + } + }) + it("calling shutdown() twice shuts down clients twice (no re-entrancy guard)", async () => { // Documents current behavior: shutdown() has no guard against being called more than // once. deactivate() only calls it once, so this isn't a bug to fix here, but a second