fix(telemetry): address second review pass

Author: edelauna

PRIVACY.md

@@ -40,6 +40,14 @@ go—and, importantly, where they don't.
   We retain telemetry only as long as needed for product analytics and
   debugging. This PostHog-based telemetry does **not** collect your code or AI
   prompts, and you can opt out at any time through the settings.
+- **Zoo Code Cloud Task Sync**: If you sign in to Zoo Code Cloud, your task
+  history — including your prompts and the assistant's responses — is uploaded
+  to and stored on Zoo Code's servers so you can view and resume tasks across
+  devices. This is separate from, and more than, the PostHog usage telemetry
+  described above. Task sync is on by default once you're signed in (or follows
+  your organization's policy if your account belongs to one); you can turn it
+  off in Zoo Code Cloud account settings, and signing out of Zoo Code Cloud
+  stops it entirely.
 - **Marketplace Requests**: When you browse or search the Marketplace for Model
   Configuration Profiles (MCPs) or Custom Modes, Zoo Code makes a secure API
   call to Zoo Code's backend servers to retrieve listing information. These

packages/telemetry/src/TelemetryService.ts

@@ -26,6 +26,15 @@ 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
@@ -43,6 +52,10 @@ export class TelemetryService {
 	// 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 {
@@ -133,7 +146,7 @@ export class TelemetryService {
 	 */
 	// 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
 		}
 
@@ -150,7 +163,7 @@ 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
 		}
 
@@ -182,6 +195,11 @@ export class TelemetryService {
 	 * Note "attempt_completion" means the model called that tool, not that the
 	 * user accepted the result -- it fires the same way whether the user goes
 	 * on to accept, decline, or give feedback instead.
+	 *
+	 * IMPORTANT for anyone querying this event (e.g. a PostHog dashboard/funnel):
+	 * "one row" no longer means "one finished task". Group by taskId and sum
+	 * toolsUsed/messageCount across completionReason installments -- do not
+	 * treat `count()` of raw events as a count of completed tasks.
 	 */
 	public captureTaskCompleted(
 		taskId: string,
@@ -370,12 +388,23 @@ export class TelemetryService {
 			return
 		}
 
+		// 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 queued while draining (e.g. from a
-		// teardown-time error handler) would otherwise be missed by one Promise.all pass.
-		while (this.pendingClientCalls.size > 0) {
-			await Promise.all(this.pendingClientCalls)
+		// 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()))

packages/telemetry/src/__tests__/TelemetryService.shutdown.test.ts

@@ -44,14 +44,15 @@ describe("TelemetryService.shutdown draining", () => {
 		expect(captureOrder).toEqual(["captured", "shutdown"])
 	})
 
-	it("drains a second capture that gets queued after shutdown() has already started draining", async () => {
+	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 second capture queued while the first Promise.all pass is still pending (e.g.
-		// a teardown-time error handler reacting to something unrelated) would be silently
-		// dropped by a single-pass drain, letting client.shutdown() run without it.
+		// -- 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<void>((resolve) => {
 			resolveFirstCapture = resolve
@@ -88,15 +89,15 @@ describe("TelemetryService.shutdown draining", () => {
 
 		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()
 
-		// Queue the second capture synchronously, immediately after shutdown() has started
-		// (and thus after its first Promise.all(pendingClientCalls) pass has already taken
-		// its snapshot). This is the scenario the loop fix protects against.
-		service.captureEvent(TelemetryEventName.TASK_CREATED, { taskId: "second" })
-
 		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
@@ -116,6 +117,64 @@ describe("TelemetryService.shutdown draining", () => {
 		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(),

packages/types/src/vscode-extension-host.ts

@@ -359,6 +359,11 @@ export type ExtensionState = Pick<
 	telemetrySetting: TelemetrySetting
 	telemetryKey?: string
 	machineId?: string
+	// Live vscode.env.isTelemetryEnabled, so the webview's own PostHog client can respect
+	// the VS Code global telemetry toggle the same way the extension-side gate does --
+	// without this, an explicit user Accept can still send events while VS Code's global
+	// telemetry is disabled.
+	vscodeTelemetryEnabled?: boolean
 
 	renderContext: "sidebar" | "editor"
 	settingsImportedAt?: number

src/__tests__/extension.spec.ts

@@ -406,6 +406,22 @@ describe("extension.ts", () => {
 
 			expect(updateTelemetryStateMock).toHaveBeenCalledWith(false)
 		})
+
+		test("pushes a state update to the webview so its own PostHog client picks up the new vscode.env.isTelemetryEnabled value", async () => {
+			const vscode = await import("vscode")
+			const { ClineProvider } = await import("../core/webview/ClineProvider")
+
+			const { activate } = await import("../extension")
+			await activate(mockContext)
+
+			const visibleInstance = (ClineProvider as any).getVisibleInstance()
+			vi.mocked(visibleInstance.postStateToWebviewWithoutClineMessages).mockClear()
+
+			const onDidChangeHandler = vi.mocked(vscode.env.onDidChangeTelemetryEnabled).mock.calls[0][0]
+			onDidChangeHandler(undefined as any)
+
+			expect(visibleInstance.postStateToWebviewWithoutClineMessages).toHaveBeenCalled()
+		})
 	})
 
 	describe("deactivate", () => {

src/api/providers/fetchers/__tests__/modelCache.spec.ts

@@ -300,6 +300,73 @@ describe("empty cache protection", () => {
 			expect(mockSet).toHaveBeenCalledWith("openrouter", mockModels)
 		})
 
+		it("reuses an in-flight fetch for concurrent getModels() calls to the same provider", async () => {
+			// Finding #11: getModels() previously had no de-duplication at all -- two concurrent
+			// cache-miss calls would each independently fire their own provider fetch. It now
+			// shares the same inFlightRefresh single-flight coordinator refreshModels() uses.
+			const mockModels = {
+				"openrouter/model": {
+					maxTokens: 8192,
+					contextWindow: 128000,
+					supportsPromptCache: false,
+					description: "OpenRouter model",
+				},
+			}
+
+			let resolvePromise: (value: typeof mockModels) => void
+			const delayedPromise = new Promise<typeof mockModels>((resolve) => {
+				resolvePromise = resolve
+			})
+			mockGetOpenRouterModels.mockReturnValue(delayedPromise)
+			mockGet.mockReturnValue(undefined)
+
+			const promise1 = getModels({ provider: "openrouter" })
+			const promise2 = getModels({ provider: "openrouter" })
+
+			expect(mockGetOpenRouterModels).toHaveBeenCalledTimes(1)
+
+			resolvePromise!(mockModels)
+
+			const [result1, result2] = await Promise.all([promise1, promise2])
+			expect(result1).toEqual(mockModels)
+			expect(result2).toEqual(mockModels)
+		})
+
+		it("shares a single in-flight fetch between getModels() and refreshModels() for the same key", async () => {
+			// The two entry points must converge on the same coordinator so a getModels() cache
+			// miss racing a concurrent refreshModels() call can't produce two unordered writes to
+			// the same cache key -- whichever fetch happened to finish last previously won,
+			// regardless of which one was actually more current.
+			const mockModels = {
+				"openrouter/model": {
+					maxTokens: 8192,
+					contextWindow: 128000,
+					supportsPromptCache: false,
+					description: "OpenRouter model",
+				},
+			}
+
+			let resolvePromise: (value: typeof mockModels) => void
+			const delayedPromise = new Promise<typeof mockModels>((resolve) => {
+				resolvePromise = resolve
+			})
+			mockGetOpenRouterModels.mockReturnValue(delayedPromise)
+			mockGet.mockReturnValue(undefined)
+
+			const { refreshModels } = await import("../modelCache")
+
+			const getPromise = getModels({ provider: "openrouter" })
+			const refreshPromise = refreshModels({ provider: "openrouter" })
+
+			expect(mockGetOpenRouterModels).toHaveBeenCalledTimes(1)
+
+			resolvePromise!(mockModels)
+
+			const [getResult, refreshResult] = await Promise.all([getPromise, refreshPromise])
+			expect(getResult).toEqual(mockModels)
+			expect(refreshResult).toEqual(mockModels)
+		})
+
 		it("re-arms the empty-response throttle after a non-empty response from an auth-scoped provider", async () => {
 			// zoo-gateway is auth-scoped and skips caching entirely, but a non-empty response
 			// must still clear the empty-response throttle so a later empty response is reported again.
@@ -525,6 +592,7 @@ describe("MODEL_CACHE_EMPTY_RESPONSE throttling", () => {
 	let freshRefreshModels: ModelCacheModule["refreshModels"]
 	let freshMockGetOpenRouterModels: Mock<typeof getOpenRouterModels>
 	let freshMockGetLiteLLMModels: Mock<typeof getLiteLLMModels>
+	let freshMockGetZooGatewayModels: Mock<typeof getZooGatewayModels>
 
 	beforeEach(async () => {
 		// The empty-response throttle is deliberately module-level, persistent state (once per
@@ -535,11 +603,13 @@ describe("MODEL_CACHE_EMPTY_RESPONSE throttling", () => {
 		const modelCacheModule: ModelCacheModule = await import("../modelCache")
 		const openRouterModule = await import("../openrouter")
 		const liteLLMModule = await import("../litellm")
+		const zooGatewayModule = await import("../zoo-gateway")
 
 		freshGetModels = modelCacheModule.getModels
 		freshRefreshModels = modelCacheModule.refreshModels
 		freshMockGetOpenRouterModels = openRouterModule.getOpenRouterModels as Mock<typeof getOpenRouterModels>
 		freshMockGetLiteLLMModels = liteLLMModule.getLiteLLMModels as Mock<typeof getLiteLLMModels>
+		freshMockGetZooGatewayModels = zooGatewayModule.getZooGatewayModels as Mock<typeof getZooGatewayModels>
 
 		const NodeCacheModule = await import("node-cache")
 		const MockedNodeCache = vi.mocked(NodeCacheModule.default)
@@ -632,6 +702,37 @@ describe("MODEL_CACHE_EMPTY_RESPONSE throttling", () => {
 
 		expect(FreshTelemetryService.instance.captureEvent).toHaveBeenCalledTimes(2)
 	})
+
+	it("throttles zoo-gateway independently per session token, even though caching itself is skipped", async () => {
+		// zoo-gateway is auth-scoped (see AUTH_SCOPED_PROVIDERS) and never persists to the
+		// memory/disk cache, but the empty-response throttle must still discriminate by
+		// identity: a sign-out/sign-in cycle to a different account carries a different
+		// session token (apiKey) on the same gateway URL, and must not have its empty-response
+		// signal suppressed by the previous account's throttle entry.
+		const { TelemetryService: FreshTelemetryService } = await import("@roo-code/telemetry")
+
+		freshMockGetZooGatewayModels.mockResolvedValue({})
+
+		await freshGetModels({ provider: "zoo-gateway", apiKey: "account-a-token" })
+		await freshGetModels({ provider: "zoo-gateway", apiKey: "account-a-token" })
+		expect(FreshTelemetryService.instance.captureEvent).toHaveBeenCalledTimes(1)
+
+		await freshGetModels({ provider: "zoo-gateway", apiKey: "account-b-token" })
+		expect(FreshTelemetryService.instance.captureEvent).toHaveBeenCalledTimes(2)
+	})
+
+	it("throttles zoo-gateway independently per gateway baseUrl", async () => {
+		// Same session token, different gateway endpoint (e.g. staging vs. production) --
+		// must also be treated as a distinct identity for throttle purposes.
+		const { TelemetryService: FreshTelemetryService } = await import("@roo-code/telemetry")
+
+		freshMockGetZooGatewayModels.mockResolvedValue({})
+
+		await freshGetModels({ provider: "zoo-gateway", apiKey: "token", baseUrl: "https://gateway-a.example.com" })
+		await freshGetModels({ provider: "zoo-gateway", apiKey: "token", baseUrl: "https://gateway-b.example.com" })
+
+		expect(FreshTelemetryService.instance.captureEvent).toHaveBeenCalledTimes(2)
+	})
 })
 
 describe("key-scoped cache key derivation", () => {