fix(vscode): honor VS Code global telemetry setting

vscode/extension/src/auth/auth.test.ts

@@ -0,0 +1,82 @@
+// SPDX-License-Identifier: Apache-2.0
+
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { ok } from '@bus/result'
+
+const mocks = vi.hoisted(() => ({
+  execAsync: vi.fn(),
+  getProjectRoot: vi.fn(),
+  getTcloudBin: vi.fn(),
+  showInformationMessage: vi.fn(),
+}))
+
+vi.mock('vscode', () => ({
+  env: { openExternal: vi.fn() },
+  Uri: { parse: vi.fn((value: string) => value) },
+  EventEmitter: class {
+    event = vi.fn()
+    fire = vi.fn()
+  },
+  window: { showInformationMessage: mocks.showInformationMessage },
+}))
+
+vi.mock('../utilities/exec', () => ({ execAsync: mocks.execAsync }))
+vi.mock('../utilities/common/utilities', () => ({
+  getProjectRoot: mocks.getProjectRoot,
+}))
+vi.mock('../utilities/sqlmesh/sqlmesh', () => ({
+  getTcloudBin: mocks.getTcloudBin,
+}))
+vi.mock('../utilities/common/log', () => ({ traceError: vi.fn() }))
+
+import { AuthenticationProviderTobikoCloud } from './auth'
+
+describe('AuthenticationProviderTobikoCloud telemetry environment', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mocks.getProjectRoot.mockResolvedValue({ uri: { fsPath: '/workspace' } })
+    mocks.getTcloudBin.mockResolvedValue(
+      ok({
+        bin: '/venv/bin/tcloud',
+        workspacePath: '/workspace',
+        env: { SQLMESH__DISABLE_ANONYMIZED_ANALYTICS: 'true' },
+        args: [],
+      }),
+    )
+    mocks.execAsync
+      .mockResolvedValueOnce({
+        exitCode: 0,
+        stdout: JSON.stringify({
+          url: 'https://example.com/login',
+          verifier_code: 'verifier',
+        }),
+        stderr: '',
+      })
+      .mockResolvedValueOnce({ exitCode: 0, stdout: '', stderr: '' })
+    // Dismissing the prompt exits without waiting for the mocked login server.
+    mocks.showInformationMessage.mockResolvedValue(undefined)
+  })
+
+  it('passes the telemetry-aware environment to every OAuth tcloud subprocess', async () => {
+    const provider = new AuthenticationProviderTobikoCloud()
+
+    await provider.sign_in_oauth_flow()
+
+    const expectedOptions = expect.objectContaining({
+      cwd: '/workspace',
+      env: { SQLMESH__DISABLE_ANONYMIZED_ANALYTICS: 'true' },
+    })
+    expect(mocks.execAsync).toHaveBeenNthCalledWith(
+      1,
+      '/venv/bin/tcloud',
+      ['auth', 'vscode', 'login-url'],
+      expectedOptions,
+    )
+    expect(mocks.execAsync).toHaveBeenNthCalledWith(
+      2,
+      '/venv/bin/tcloud',
+      ['auth', 'vscode', 'start-server', 'verifier'],
+      expectedOptions,
+    )
+  })
+})

vscode/extension/src/auth/auth.ts

@@ -193,6 +193,7 @@ export class AuthenticationProviderTobikoCloud
       ['auth', 'vscode', 'login-url'],
       {
         cwd: workspacePath.uri.fsPath,
+        env: tcloudBinPath.env,
       },
     )
     if (result.exitCode !== 0) {

vscode/extension/src/extension.ts

@@ -24,7 +24,10 @@ import {
   traceError,
 } from './utilities/common/log'
 import { onDidChangePythonInterpreter } from './utilities/common/python'
-import { requiresLspRestart } from './utilities/common/configurationChange'
+import {
+  requiresLspRestart,
+  restartLspOnTelemetryChange,
+} from './utilities/common/configurationChange'
 import { coalesceAsync } from './utilities/coalesceAsync'
 import { sleep } from './utilities/sleep'
 import { ErrorType, handleError } from './utilities/errors'
@@ -128,6 +131,15 @@ export async function activate(context: vscode.ExtensionContext) {
     }
   }
 
+  // Subscribe before starting the first server so a consent change during a
+  // slow startup is not lost. The subscription defers that restart until the
+  // first client and its test controller are fully initialized.
+  const telemetryRestartSubscription = restartLspOnTelemetryChange(
+    vscode.env.onDidChangeTelemetryEnabled,
+    restartLsp,
+  )
+  context.subscriptions.push(telemetryRestartSubscription)
+
   // commands needing the restart helper
   context.subscriptions.push(
     vscode.commands.registerCommand(
@@ -141,24 +153,34 @@ export async function activate(context: vscode.ExtensionContext) {
     vscode.commands.registerCommand('sqlmesh.signout', signOut(authProvider)),
   )
 
-  //  Instantiate the LSP client (once)
-  lspClient = new LSPClient()
-  const startResult = await lspClient.start()
-  if (isErr(startResult)) {
-    await handleError(
-      authProvider,
-      restartLsp,
-      startResult.error,
-      'Failed to start LSP',
-    )
-    return // abort activation – nothing else to do
-  }
+  // Instantiate the LSP client once. The telemetry subscription is completed
+  // in a finally block so even a failed initial start cannot leave every later
+  // telemetry change deferred forever.
+  const initialLspClient = new LSPClient()
+  lspClient = initialLspClient
+  let initialStartSucceeded = false
+  await telemetryRestartSubscription.runDuringInitialStart(async () => {
+    const startResult = await initialLspClient.start()
+    if (isErr(startResult)) {
+      await handleError(
+        authProvider,
+        restartLsp,
+        startResult.error,
+        'Failed to start LSP',
+      )
+      return
+    }
 
-  context.subscriptions.push(lspClient)
+    context.subscriptions.push(initialLspClient)
 
-  // Initialize the test controller
-  testControllerDisposable = setupTestController(lspClient)
-  context.subscriptions.push(testControllerDisposable, testController)
+    // Initialize the test controller
+    testControllerDisposable = setupTestController(initialLspClient)
+    context.subscriptions.push(testControllerDisposable, testController)
+    initialStartSucceeded = true
+  })
+  if (!initialStartSucceeded) {
+    return // abort activation – nothing else to do
+  }
 
   // Register the rendered model provider
   const renderedModelProvider = new RenderedModelProvider()

vscode/extension/src/utilities/common/configurationChange.test.ts

@@ -1,7 +1,10 @@
 // SPDX-License-Identifier: Apache-2.0
 
-import { describe, expect, it } from 'vitest'
-import { requiresLspRestart } from './configurationChange'
+import { describe, expect, it, vi } from 'vitest'
+import {
+  requiresLspRestart,
+  restartLspOnTelemetryChange,
+} from './configurationChange'
 
 /**
  * Build a stand-in for `vscode.ConfigurationChangeEvent` from the settings that
@@ -58,3 +61,73 @@ describe('requiresLspRestart', () => {
     ).toBe(true)
   })
 })
+
+describe('restartLspOnTelemetryChange', () => {
+  const setup = () => {
+    let listener: ((enabled: boolean) => void) | undefined
+    const disposable = { dispose: vi.fn() }
+    const onDidChangeTelemetryEnabled = vi.fn(
+      (registeredListener: (enabled: boolean) => void) => {
+        listener = registeredListener
+        return disposable
+      },
+    )
+    const restartLsp = vi.fn(() => Promise.resolve())
+
+    const subscription = restartLspOnTelemetryChange(
+      onDidChangeTelemetryEnabled,
+      restartLsp,
+    )
+
+    return { disposable, listener: () => listener, restartLsp, subscription }
+  }
+
+  it('defers and collapses telemetry changes during initial LSP startup', async () => {
+    const { listener, restartLsp, subscription } = setup()
+
+    await subscription.runDuringInitialStart(() => {
+      listener()?.(false)
+      listener()?.(true)
+      expect(restartLsp).not.toHaveBeenCalled()
+      return Promise.resolve()
+    })
+
+    expect(restartLsp).toHaveBeenCalledOnce()
+  })
+
+  it('restarts for every telemetry preference change after initial startup', async () => {
+    const { listener, restartLsp, subscription } = setup()
+    await subscription.runDuringInitialStart(() => Promise.resolve())
+
+    expect(restartLsp).not.toHaveBeenCalled()
+
+    listener()?.(false)
+    await vi.waitFor(() => expect(restartLsp).toHaveBeenCalledTimes(1))
+
+    listener()?.(true)
+    await vi.waitFor(() => expect(restartLsp).toHaveBeenCalledTimes(2))
+  })
+
+  it('releases deferred restarts when initial startup fails', async () => {
+    const { listener, restartLsp, subscription } = setup()
+
+    await expect(
+      subscription.runDuringInitialStart(() => {
+        listener()?.(false)
+        return Promise.reject(new Error('initial start failed'))
+      }),
+    ).rejects.toThrow('initial start failed')
+
+    expect(restartLsp).toHaveBeenCalledOnce()
+    listener()?.(true)
+    await vi.waitFor(() => expect(restartLsp).toHaveBeenCalledTimes(2))
+  })
+
+  it('disposes the underlying VS Code event listener', () => {
+    const { disposable, subscription } = setup()
+
+    subscription.dispose()
+
+    expect(disposable.dispose).toHaveBeenCalledOnce()
+  })
+})

vscode/extension/src/utilities/common/configurationChange.ts

@@ -21,6 +21,18 @@ export interface ConfigurationChange {
   affectsConfiguration(section: string): boolean
 }
 
+interface Disposable {
+  dispose(): unknown
+}
+
+export interface TelemetryRestartSubscription extends Disposable {
+  runDuringInitialStart<T>(task: () => Promise<T>): Promise<T>
+}
+
+type TelemetryChangeEvent<TDisposable extends Disposable> = (
+  listener: (enabled: boolean) => unknown,
+) => TDisposable
+
 /**
  * Whether a configuration change affects a setting the language server reads.
  *
@@ -33,3 +45,34 @@ export function requiresLspRestart(event: ConfigurationChange): boolean {
     event.affectsConfiguration(section),
   )
 }
+
+/** Restart the language server whenever VS Code's effective telemetry preference changes. */
+export function restartLspOnTelemetryChange<TDisposable extends Disposable>(
+  event: TelemetryChangeEvent<TDisposable>,
+  restartLsp: () => Promise<void>,
+): TelemetryRestartSubscription {
+  let initialStartComplete = false
+  let restartPending = false
+  const subscription = event(() => {
+    if (!initialStartComplete) {
+      restartPending = true
+      return
+    }
+    void restartLsp()
+  })
+
+  return {
+    dispose: () => subscription.dispose(),
+    runDuringInitialStart: async task => {
+      try {
+        return await task()
+      } finally {
+        initialStartComplete = true
+        if (restartPending) {
+          restartPending = false
+          await restartLsp()
+        }
+      }
+    },
+  }
+}