Register test commands before first await to prevent activation race condition

src/client/extension.ts

@@ -45,6 +45,8 @@ import { buildProposedApi } from './proposedApi';
 import { GLOBAL_PERSISTENT_KEYS } from './common/persistentState';
 import { registerTools } from './chat';
 import { IRecommendedEnvironmentService } from './interpreter/configuration/types';
+import { registerTypes as unitTestsRegisterTypes } from './testing/serviceRegistry';
+import { registerTestCommands } from './testing/main';
 
 durations.codeLoadingTime = stopWatch.elapsedTime;
 
@@ -121,6 +123,11 @@ async function activateUnsafe(
     // Note standard utils especially experiment and platform code are fundamental to the extension
     // and should be available before we activate anything else.Hence register them first.
     initializeStandard(ext);
+
+    // Register test services and commands early to prevent race conditions.
+    unitTestsRegisterTypes(ext.legacyIOC.serviceManager);
+    registerTestCommands(activatedServiceContainer);
+
     // We need to activate experiments before initializing components as objects are created or not created based on experiments.
     const experimentService = activatedServiceContainer.get<IExperimentService>(IExperimentService);
     // This guarantees that all experiment information has loaded & all telemetry will contain experiment info.

src/client/extensionActivation.ts

@@ -33,7 +33,6 @@ import { setExtensionInstallTelemetryProperties } from './telemetry/extensionIns
 import { registerTypes as tensorBoardRegisterTypes } from './tensorBoard/serviceRegistry';
 import { registerTypes as commonRegisterTerminalTypes } from './terminals/serviceRegistry';
 import { ICodeExecutionHelper, ICodeExecutionManager, ITerminalAutoActivation } from './terminals/types';
-import { registerTypes as unitTestsRegisterTypes } from './testing/serviceRegistry';
 
 // components
 import * as pythonEnvironments from './pythonEnvironments';
@@ -144,7 +143,6 @@ async function activateLegacy(ext: ExtensionState, startupStopWatch: StopWatch):
     const { enableProposedApi } = applicationEnv.packageJson;
     serviceManager.addSingletonInstance<boolean>(UseProposedApi, enableProposedApi);
     // Feature specific registrations.
-    unitTestsRegisterTypes(serviceManager);
     installerRegisterTypes(serviceManager);
     commonRegisterTerminalTypes(serviceManager);
     debugConfigurationRegisterTypes(serviceManager);

src/client/testing/main.ts

@@ -18,7 +18,7 @@ import { IDisposableRegistry, Product } from '../common/types';
 import { IInterpreterService } from '../interpreter/contracts';
 import { IServiceContainer } from '../ioc/types';
 import { EventName } from '../telemetry/constants';
-import { captureTelemetry, sendTelemetryEvent } from '../telemetry/index';
+import { sendTelemetryEvent } from '../telemetry/index';
 import { selectTestWorkspace } from './common/testUtils';
 import { TestSettingsPropertyNames } from './configuration/types';
 import { ITestConfigurationService, ITestsHelper } from './common/types';
@@ -42,6 +42,91 @@ export class TestingService implements ITestingService {
     }
 }
 
+/**
+ * Registers command handlers but defers service resolution until the commands are actually invoked,
+ * allowing registration to happen before all services are fully initialized.
+ */
+export function registerTestCommands(serviceContainer: IServiceContainer): void {
+    // Resolve only the essential services needed for command registration itself
+    const disposableRegistry = serviceContainer.get<Disposable[]>(IDisposableRegistry);
+    const commandManager = serviceContainer.get<ICommandManager>(ICommandManager);
+
+    // Helper function to configure tests - services are resolved when invoked, not at registration time
+    const configureTestsHandler = async (resource?: Uri) => {
+        sendTelemetryEvent(EventName.UNITTEST_CONFIGURE);
+
+        // Resolve services lazily when the command is invoked
+        const workspaceService = serviceContainer.get<IWorkspaceService>(IWorkspaceService);
+
+        let wkspace: Uri | undefined;
+        if (resource) {
+            const wkspaceFolder = workspaceService.getWorkspaceFolder(resource);
+            wkspace = wkspaceFolder ? wkspaceFolder.uri : undefined;
+        } else {
+            const appShell = serviceContainer.get<IApplicationShell>(IApplicationShell);
+            wkspace = await selectTestWorkspace(appShell);
+        }
+        if (!wkspace) {
+            return;
+        }
+        const interpreterService = serviceContainer.get<IInterpreterService>(IInterpreterService);
+        const cmdManager = serviceContainer.get<ICommandManager>(ICommandManager);
+        if (!(await interpreterService.getActiveInterpreter(wkspace))) {
+            cmdManager.executeCommand(constants.Commands.TriggerEnvironmentSelection, wkspace);
+            return;
+        }
+        const configurationService = serviceContainer.get<ITestConfigurationService>(ITestConfigurationService);
+        await configurationService.promptToEnableAndConfigureTestFramework(wkspace);
+    };
+
+    disposableRegistry.push(
+        // Command: python.configureTests - prompts user to configure test framework
+        commandManager.registerCommand(
+            constants.Commands.Tests_Configure,
+            (_, _cmdSource: constants.CommandSource = constants.CommandSource.commandPalette, resource?: Uri) => {
+                // Invoke configuration handler (errors are ignored as this can be called from multiple places)
+                configureTestsHandler(resource).ignoreErrors();
+                traceVerbose('Testing: Trigger refresh after config change');
+                // Refresh test data if test controller is available (resolved lazily)
+                if (tests && !!tests.createTestController) {
+                    const testController = serviceContainer.get<ITestController>(ITestController);
+                    testController?.refreshTestData(resource, { forceRefresh: true });
+                }
+            },
+        ),
+        // Command: python.tests.copilotSetup - Copilot integration for test setup
+        commandManager.registerCommand(constants.Commands.Tests_CopilotSetup, (resource?: Uri):
+            | { message: string; command: Command }
+            | undefined => {
+            // Resolve services lazily when the command is invoked
+            const workspaceService = serviceContainer.get<IWorkspaceService>(IWorkspaceService);
+            const wkspaceFolder =
+                workspaceService.getWorkspaceFolder(resource) || workspaceService.workspaceFolders?.at(0);
+            if (!wkspaceFolder) {
+                return undefined;
+            }
+
+            const configurationService = serviceContainer.get<ITestConfigurationService>(ITestConfigurationService);
+            if (configurationService.hasConfiguredTests(wkspaceFolder.uri)) {
+                return undefined;
+            }
+
+            return {
+                message: Testing.copilotSetupMessage,
+                command: {
+                    title: Testing.configureTests,
+                    command: constants.Commands.Tests_Configure,
+                    arguments: [undefined, constants.CommandSource.ui, resource],
+                },
+            };
+        }),
+        // Command: python.copyTestId - copies test ID to clipboard
+        commandManager.registerCommand(constants.Commands.CopyTestId, async (testItem: TestItem) => {
+            writeTestIdToClipboard(testItem);
+        }),
+    );
+}
+
 @injectable()
 export class UnitTestManagementService implements IExtensionActivationService {
     private activatedOnce: boolean = false;
@@ -80,7 +165,6 @@ export class UnitTestManagementService implements IExtensionActivationService {
         this.activatedOnce = true;
 
         this.registerHandlers();
-        this.registerCommands();
 
         if (!!tests.testResults) {
             await this.updateTestUIButtons();
@@ -130,73 +214,6 @@ export class UnitTestManagementService implements IExtensionActivationService {
         await Promise.all(changedWorkspaces.map((u) => this.testController?.refreshTestData(u)));
     }
 
-    @captureTelemetry(EventName.UNITTEST_CONFIGURE, undefined, false)
-    private async configureTests(resource?: Uri) {
-        let wkspace: Uri | undefined;
-        if (resource) {
-            const wkspaceFolder = this.workspaceService.getWorkspaceFolder(resource);
-            wkspace = wkspaceFolder ? wkspaceFolder.uri : undefined;
-        } else {
-            const appShell = this.serviceContainer.get<IApplicationShell>(IApplicationShell);
-            wkspace = await selectTestWorkspace(appShell);
-        }
-        if (!wkspace) {
-            return;
-        }
-        const interpreterService = this.serviceContainer.get<IInterpreterService>(IInterpreterService);
-        const commandManager = this.serviceContainer.get<ICommandManager>(ICommandManager);
-        if (!(await interpreterService.getActiveInterpreter(wkspace))) {
-            commandManager.executeCommand(constants.Commands.TriggerEnvironmentSelection, wkspace);
-            return;
-        }
-        const configurationService = this.serviceContainer.get<ITestConfigurationService>(ITestConfigurationService);
-        await configurationService.promptToEnableAndConfigureTestFramework(wkspace!);
-    }
-
-    private registerCommands(): void {
-        const commandManager = this.serviceContainer.get<ICommandManager>(ICommandManager);
-        this.disposableRegistry.push(
-            commandManager.registerCommand(
-                constants.Commands.Tests_Configure,
-                (_, _cmdSource: constants.CommandSource = constants.CommandSource.commandPalette, resource?: Uri) => {
-                    // Ignore the exceptions returned.
-                    // This command will be invoked from other places of the extension.
-                    this.configureTests(resource).ignoreErrors();
-                    traceVerbose('Testing: Trigger refresh after config change');
-                    this.testController?.refreshTestData(resource, { forceRefresh: true });
-                },
-            ),
-            commandManager.registerCommand(constants.Commands.Tests_CopilotSetup, (resource?: Uri):
-                | { message: string; command: Command }
-                | undefined => {
-                const wkspaceFolder =
-                    this.workspaceService.getWorkspaceFolder(resource) || this.workspaceService.workspaceFolders?.at(0);
-                if (!wkspaceFolder) {
-                    return undefined;
-                }
-
-                const configurationService = this.serviceContainer.get<ITestConfigurationService>(
-                    ITestConfigurationService,
-                );
-                if (configurationService.hasConfiguredTests(wkspaceFolder.uri)) {
-                    return undefined;
-                }
-
-                return {
-                    message: Testing.copilotSetupMessage,
-                    command: {
-                        title: Testing.configureTests,
-                        command: constants.Commands.Tests_Configure,
-                        arguments: [undefined, constants.CommandSource.ui, resource],
-                    },
-                };
-            }),
-            commandManager.registerCommand(constants.Commands.CopyTestId, async (testItem: TestItem) => {
-                writeTestIdToClipboard(testItem);
-            }),
-        );
-    }
-
     private registerHandlers() {
         const interpreterService = this.serviceContainer.get<IInterpreterService>(IInterpreterService);
         this.disposableRegistry.push(