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.test.ts

@@ -0,0 +1,160 @@
+// SPDX-License-Identifier: Apache-2.0
+
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+import type { ExtensionContext } from 'vscode'
+import { err, ok } from '@bus/result'
+
+const mocks = vi.hoisted(() => ({
+  telemetryListener: undefined as ((enabled: boolean) => unknown) | undefined,
+  start: vi.fn(),
+  restart: vi.fn(),
+  registerCommand: vi.fn<(command: string) => { dispose(): void }>(() => ({
+    dispose: vi.fn(),
+  })),
+  registerContentProvider: vi.fn(() => ({ dispose: vi.fn() })),
+  registerViewProvider: vi.fn(() => ({ dispose: vi.fn() })),
+  setupTestController: vi.fn(() => ({ dispose: vi.fn() })),
+  testController: { dispose: vi.fn() },
+  handleError: vi.fn(),
+}))
+
+vi.mock('vscode', () => ({
+  env: {
+    onDidChangeTelemetryEnabled: (listener: (enabled: boolean) => unknown) => {
+      mocks.telemetryListener = listener
+      return { dispose: vi.fn() }
+    },
+  },
+  authentication: {
+    registerAuthenticationProvider: vi.fn(() => ({ dispose: vi.fn() })),
+  },
+  commands: { registerCommand: mocks.registerCommand },
+  workspace: {
+    registerTextDocumentContentProvider: mocks.registerContentProvider,
+    onDidSaveTextDocument: vi.fn(() => ({ dispose: vi.fn() })),
+  },
+  window: { registerWebviewViewProvider: mocks.registerViewProvider },
+}))
+
+vi.mock('./lsp/lsp', () => ({
+  LSPClient: class {
+    start = mocks.start
+    restart = mocks.restart
+    hasCompletionCapability = () => true
+  },
+}))
+vi.mock('./utilities/common/vscodeapi', () => ({
+  createOutputChannel: vi.fn(() => ({ dispose: vi.fn() })),
+  onDidChangeConfiguration: vi.fn(() => ({ dispose: vi.fn() })),
+  registerCommand: mocks.registerCommand,
+}))
+vi.mock('./utilities/common/log', () => ({
+  registerLogger: vi.fn(() => ({ dispose: vi.fn() })),
+  traceInfo: vi.fn(),
+  traceVerbose: vi.fn(),
+  traceError: vi.fn(),
+}))
+vi.mock('./utilities/common/python', () => ({
+  onDidChangePythonInterpreter: vi.fn(() => ({ dispose: vi.fn() })),
+}))
+vi.mock('./utilities/errors', () => ({ handleError: mocks.handleError }))
+vi.mock('./auth/auth', () => ({
+  AuthenticationProviderTobikoCloud: vi.fn(),
+}))
+vi.mock('./tests/tests', () => ({
+  controller: mocks.testController,
+  setupTestController: mocks.setupTestController,
+}))
+vi.mock('./providers/renderedModelProvider', () => ({
+  RenderedModelProvider: Object.assign(vi.fn(), {
+    getScheme: () => 'sqlmesh-rendered',
+  }),
+}))
+vi.mock('./webviews/lineagePanel', () => ({ LineagePanel: vi.fn() }))
+vi.mock('./completion/completion', () => ({
+  selector: [],
+  completionProvider: vi.fn(),
+}))
+vi.mock('./commands/format', () => ({ format: vi.fn() }))
+vi.mock('./commands/signout', () => ({ signOut: vi.fn() }))
+vi.mock('./commands/signin', () => ({ signIn: vi.fn() }))
+vi.mock('./commands/signinSpecifyFlow', () => ({ signInSpecifyFlow: vi.fn() }))
+vi.mock('./commands/renderModel', () => ({
+  renderModel: vi.fn(),
+  reRenderModelForSourceFile: vi.fn(),
+}))
+vi.mock('./commands/stop', () => ({ stop: vi.fn() }))
+vi.mock('./commands/printEnvironment', () => ({ printEnvironment: vi.fn() }))
+vi.mock('./commands/tableDiff', () => ({ showTableDiff: vi.fn() }))
+
+describe('activation after a telemetry change during startup', () => {
+  beforeEach(() => {
+    vi.resetModules()
+    vi.clearAllMocks()
+    mocks.telemetryListener = undefined
+    mocks.restart.mockResolvedValue(ok(undefined))
+    mocks.handleError.mockResolvedValue(undefined)
+  })
+
+  it.each([false, true])(
+    'registers the full extension after a successful retry, initial start failed: %s',
+    async initialStartFailed => {
+      mocks.start.mockImplementationOnce(() => {
+        mocks.telemetryListener?.(false)
+        return Promise.resolve(
+          initialStartFailed
+            ? err({ type: 'sqlmesh_lsp_not_found' })
+            : ok(undefined),
+        )
+      })
+      const context = {
+        subscriptions: [],
+        extensionUri: {},
+      } as unknown as ExtensionContext
+      const { activate } = await import('./extension')
+
+      await activate(context)
+
+      expect(mocks.restart).toHaveBeenCalledOnce()
+      expect(
+        mocks.registerCommand.mock.calls.map(([command]) => command),
+      ).toEqual([
+        'sqlmesh.signin',
+        'sqlmesh.signinSpecifyFlow',
+        'sqlmesh.signout',
+        'sqlmesh.renderModel',
+        'sqlmesh.showTableDiff',
+        'sqlmesh.format',
+        'sqlmesh.restart',
+        'sqlmesh.stop',
+        'sqlmesh.printEnvironment',
+      ])
+      expect(mocks.registerContentProvider).toHaveBeenCalledOnce()
+      expect(mocks.registerViewProvider).toHaveBeenCalledOnce()
+      expect(
+        context.subscriptions.filter(item => item === mocks.testController),
+      ).toHaveLength(1)
+    },
+  )
+
+  it('does not complete activation when both the initial start and retry fail', async () => {
+    const failure = err({ type: 'sqlmesh_lsp_not_found' })
+    mocks.start.mockImplementationOnce(() => {
+      mocks.telemetryListener?.(false)
+      return Promise.resolve(failure)
+    })
+    mocks.restart.mockResolvedValue(failure)
+    const context = {
+      subscriptions: [],
+      extensionUri: {},
+    } as unknown as ExtensionContext
+    const { activate } = await import('./extension')
+
+    await activate(context)
+
+    expect(mocks.restart).toHaveBeenCalledOnce()
+    expect(mocks.registerContentProvider).not.toHaveBeenCalled()
+    expect(mocks.registerViewProvider).not.toHaveBeenCalled()
+    expect(mocks.setupTestController).not.toHaveBeenCalled()
+  })
+})

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'
@@ -80,6 +83,8 @@ export async function activate(context: vscode.ExtensionContext) {
    * the run that is still waiting on the handler.
    */
   let restartError: ErrorType | undefined
+  // A retry during activation can succeed even when the first start failed.
+  let lspStartSucceeded = false
 
   const runRestart = async (): Promise<void> => {
     const invokedByUser = restartInvokedByUser
@@ -107,6 +112,7 @@ export async function activate(context: vscode.ExtensionContext) {
     }
     testControllerDisposable = setupTestController(lspClient)
     context.subscriptions.push(testControllerDisposable)
+    lspStartSucceeded = true
   }
 
   /**
@@ -128,6 +134,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 +156,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
+  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)
+    lspStartSucceeded = true
+  })
+  if (!lspStartSucceeded) {
+    return // abort activation – nothing else to do
+  }
+  context.subscriptions.push(testController)
 
   // Register the rendered model provider
   const renderedModelProvider = new RenderedModelProvider()