diff --git a/src/extension.ts b/src/extension.ts index e088619c2..ccea52ae4 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -78,6 +78,7 @@ import { registerInlineScriptUx } from './features/inlineScript/setupEnvironment import { applyInitialEnvironmentSelection, registerInterpreterSettingsChangeListener, + registerLateManagerRegistrationListener, } from './features/interpreterSelection'; import type { PythonProjectManager } from './features/projectManager'; import { PythonProjectManagerImpl } from './features/projectManager'; @@ -776,6 +777,10 @@ export async function activate(context: ExtensionContext): Promise(); + +/** Serializes selection runs so a re-run triggered by a late registration cannot be overwritten by an older run. */ +let selectionQueue: Promise = Promise.resolve(); + +function enqueueSelection(run: () => Promise): Promise { + const result = selectionQueue.then(run); + selectionQueue = result.catch(() => undefined); + return result; +} + +/** Test hook: forget folders recorded as fallen back by earlier selection runs. */ +export function resetFallbackSelections(): void { + fallbackSelections.clear(); +} + /** * Result from the priority chain resolution. */ @@ -281,7 +304,33 @@ async function autoDiscoverEnvironment( * @param changeEvent - The settings change that triggered a re-run. The global scope is resolved again only when * a setting it depends on changed, so a user's global selection is kept otherwise. */ -export async function applyInitialEnvironmentSelection( +export function applyInitialEnvironmentSelection( + envManagers: EnvironmentManagers, + projectManager: PythonProjectManager, + nativeFinder: NativePythonFinder, + api: PythonEnvironmentApi, + activationToReadyDurationMs?: number, + globalScopeDeferredRef?: { value: 'deferred' | 'not_deferred' | 'unknown' }, + changeEvent?: ConfigurationChangeEvent, +): Promise { + return enqueueSelection(() => + applyEnvironmentSelectionNow( + envManagers, + projectManager, + nativeFinder, + api, + activationToReadyDurationMs, + globalScopeDeferredRef, + changeEvent, + ), + ); +} + +/** + * @param onlyFolders - When set, only these workspace folders (by URI string) are selected and the global scope is + * left alone. Used to re-select folders that fell back because their configured manager registered late. + */ +async function applyEnvironmentSelectionNow( envManagers: EnvironmentManagers, projectManager: PythonProjectManager, nativeFinder: NativePythonFinder, @@ -289,8 +338,9 @@ export async function applyInitialEnvironmentSelection( activationToReadyDurationMs?: number, globalScopeDeferredRef?: { value: 'deferred' | 'not_deferred' | 'unknown' }, changeEvent?: ConfigurationChangeEvent, + onlyFolders?: ReadonlySet, ): Promise { - const folders = getWorkspaceFolders() ?? []; + const folders = (getWorkspaceFolders() ?? []).filter((f) => !onlyFolders || onlyFolders.has(f.uri.toString())); traceInfo( `[interpreterSelection] Applying initial environment selection for ${folders.length} workspace folder(s)`, ); @@ -335,6 +385,15 @@ export async function applyInitialEnvironmentSelection( explicit: isExplicitSelection(result, env), }); + const missingManagerIds = errors + .filter((e) => e.kind === 'managerNotRegistered') + .map((e) => e.configuredValue); + if (missingManagerIds.length > 0) { + fallbackSelections.set(folder.uri.toString(), { missingManagerIds, fellBackAt: Date.now() }); + } else { + fallbackSelections.delete(folder.uri.toString()); + } + if (env) { workspaceFolderResolved = true; resolvedFolderCount++; @@ -388,9 +447,10 @@ export async function applyInitialEnvironmentSelection( // pythonProjects never applies to the global scope. Resolving it again when only that changed would replace // a user's global selection with the unchanged result of the other settings. const globalSettingsChanged = - !changeEvent || - changeEvent.affectsConfiguration('python-envs.defaultEnvManager') || - changeEvent.affectsConfiguration('python.defaultInterpreterPath'); + !onlyFolders && + (!changeEvent || + changeEvent.affectsConfiguration('python-envs.defaultEnvManager') || + changeEvent.affectsConfiguration('python.defaultInterpreterPath')); if (!globalSettingsChanged) { traceVerbose('[interpreterSelection] Settings for the global scope did not change, keeping its environment'); @@ -720,6 +780,71 @@ async function tryResolveInterpreterPath( return undefined; } +/** + * Re-apply environment selection when a manager named by a setting registers after the priority chain skipped it. + * + * Startup selection does not wait for managers contributed by other extensions, because there is no way to know + * which extensions will register or when. Instead of blocking on them, selection falls back immediately and is + * corrected here once the configured manager actually registers: + * - Only the workspace folders that fell back for that manager are selected again; other folders and the global + * scope are left alone. + * - The correction only applies within MANAGER_READY_TIMEOUT_MS of the fallback (the same window the API uses when + * waiting for managers). A manager that shows up later, e.g. an extension installed mid-session, does not switch + * the interpreter under a running session. + * - The re-selection is queued behind any selection that is still running, so it cannot be overwritten by it. + * + * Register this before startup selection runs so that registrations which happen while it is in progress are not + * missed. + */ +export function registerLateManagerRegistrationListener( + envManagers: EnvironmentManagers, + projectManager: PythonProjectManager, + nativeFinder: NativePythonFinder, + api: PythonEnvironmentApi, +): Disposable { + return envManagers.onDidChangeEnvironmentManager((e) => { + if (e.kind !== 'registered') { + return; + } + const managerId = e.manager.id; + // Queued, so a registration that happens during startup selection is handled once that run has recorded + // which folders fell back. + enqueueSelection(async () => { + const now = Date.now(); + const folders = new Set(); + for (const [folder, fallback] of fallbackSelections) { + if (!fallback.missingManagerIds.includes(managerId)) { + continue; + } + fallbackSelections.delete(folder); + if (now - fallback.fellBackAt <= MANAGER_READY_TIMEOUT_MS) { + folders.add(folder); + } else { + traceInfo( + `[interpreterSelection] Manager '${managerId}' registered too long after ${folder} fell back, keeping its environment`, + ); + } + } + if (folders.size === 0) { + return; + } + traceInfo( + `[interpreterSelection] Configured manager '${managerId}' registered late, re-applying selection for ${folders.size} workspace folder(s)`, + ); + await applyEnvironmentSelectionNow( + envManagers, + projectManager, + nativeFinder, + api, + undefined, + undefined, + undefined, + folders, + ); + }).catch((err) => traceError(`[interpreterSelection] Re-applying selection for '${managerId}' failed: ${err}`)); + }); +} + /** * Register a configuration change listener for interpreter-related settings. When relevant settings change (defaultInterpreterPath, defaultEnvManager, pythonProjects), * re-run the priority chain to apply the new settings immediately. diff --git a/src/test/features/interpreterSelection.unit.test.ts b/src/test/features/interpreterSelection.unit.test.ts index 7411ecfd6..d9b7195d4 100644 --- a/src/test/features/interpreterSelection.unit.test.ts +++ b/src/test/features/interpreterSelection.unit.test.ts @@ -21,6 +21,8 @@ import * as frameUtils from '../../common/utils/frameUtils'; import { applyInitialEnvironmentSelection, registerInterpreterSettingsChangeListener, + registerLateManagerRegistrationListener, + resetFallbackSelections, resetSettingWarnings, resolveEnvironmentByPriority, resolveGlobalEnvironmentByPriority, @@ -2603,3 +2605,212 @@ suite('Interpreter Selection - Multi-Root Workspace', () => { assert.strictEqual(mockEnvManagers.setEnvironment.callCount, 2, 'Both folders should be processed'); }); }); + +suite('Interpreter Selection - late manager registration', () => { + let sandbox: sinon.SinonSandbox; + let mockEnvManagers: sinon.SinonStubbedInstance; + let mockProjectManager: sinon.SinonStubbedInstance; + let mockNativeFinder: sinon.SinonStubbedInstance; + let mockApi: sinon.SinonStubbedInstance; + let mockVenvManager: sinon.SinonStubbedInstance; + let mockPixiManager: sinon.SinonStubbedInstance; + let pixiRegistered: boolean; + let fireManagerEvent: (e: { kind: 'registered' | 'unregistered'; manager: InternalEnvironmentManager }) => void; + + const testUri = Uri.file('/test/workspace'); + const pixiManagerId = 'renan-r-santos.pixi-code:pixi'; + const makeEnv = (id: string, managerId: string): PythonEnvironment => ({ + envId: { id, managerId }, + name: id, + displayName: id, + version: '3.12.0', + displayPath: `/test/${id}`, + environmentPath: Uri.file(`/test/${id}`), + sysPrefix: `/test/${id}`, + execInfo: { run: { executable: `/test/${id}/bin/python` } }, + }); + const venvEnv = makeEnv('venv-env', VENV_MANAGER_ID); + const pixiEnv = makeEnv('pixi-env', pixiManagerId); + + setup(() => { + sandbox = sinon.createSandbox(); + resetSettingWarnings(); + resetFallbackSelections(); + pixiRegistered = false; + + mockVenvManager = { + id: VENV_MANAGER_ID, + name: 'venv', + displayName: 'Venv', + get: sandbox.stub().resolves(venvEnv), + set: sandbox.stub().resolves(), + } as unknown as sinon.SinonStubbedInstance; + mockPixiManager = { + id: pixiManagerId, + name: 'pixi', + displayName: 'Pixi', + get: sandbox.stub().resolves(pixiEnv), + set: sandbox.stub().resolves(), + } as unknown as sinon.SinonStubbedInstance; + + mockEnvManagers = { + getEnvironmentManager: sandbox.stub().callsFake((scope: unknown) => { + if (scope === VENV_MANAGER_ID) { + return mockVenvManager; + } + if (scope === pixiManagerId && pixiRegistered) { + return mockPixiManager; + } + return undefined; + }), + setEnvironment: sandbox.stub().resolves(), + setEnvironments: sandbox.stub().resolves(), + onDidChangeEnvironmentManager: sandbox.stub().callsFake((handler) => { + fireManagerEvent = handler; + return { dispose: () => undefined }; + }), + managers: [mockVenvManager], + } as unknown as sinon.SinonStubbedInstance; + mockProjectManager = { get: sandbox.stub() } as unknown as sinon.SinonStubbedInstance; + mockNativeFinder = { resolve: sandbox.stub() } as unknown as sinon.SinonStubbedInstance; + mockApi = { resolveEnvironment: sandbox.stub() } as unknown as sinon.SinonStubbedInstance; + + sandbox.stub(workspaceApis, 'getWorkspaceFolders').returns([{ uri: testUri, name: 'test', index: 0 }]); + sandbox.stub(workspaceApis, 'getConfiguration').returns(createMockConfig([]) as WorkspaceConfiguration); + sandbox.stub(helpers, 'getUserConfiguredSetting').callsFake((section: string, key: string) => { + if (section === 'python-envs' && key === 'defaultEnvManager') { + return pixiManagerId; + } + return undefined; + }); + sandbox.stub(windowApis, 'showWarningMessage').resolves(undefined); + }); + + teardown(() => { + sandbox.restore(); + }); + + const run = () => + applyInitialEnvironmentSelection( + mockEnvManagers as unknown as EnvironmentManagers, + mockProjectManager as unknown as PythonProjectManager, + mockNativeFinder as unknown as NativePythonFinder, + mockApi as unknown as PythonEnvironmentApi, + ); + + const listen = () => + registerLateManagerRegistrationListener( + mockEnvManagers as unknown as EnvironmentManagers, + mockProjectManager as unknown as PythonProjectManager, + mockNativeFinder as unknown as NativePythonFinder, + mockApi as unknown as PythonEnvironmentApi, + ); + + // Lets the selection re-run queued by the listener finish (it is not awaited by the event handler). + const settle = async () => { + for (let i = 0; i < 20; i++) { + await new Promise((resolve) => setImmediate(resolve)); + } + }; + + test('re-applies selection when the configured manager registers after startup selection', async () => { + listen(); + await run(); + + // Startup selection fell back because the configured manager was not registered yet. + assert.strictEqual(mockEnvManagers.setEnvironment.firstCall.args[1], venvEnv); + + pixiRegistered = true; + fireManagerEvent({ kind: 'registered', manager: mockPixiManager }); + await settle(); + + const envs = mockEnvManagers.setEnvironment.getCalls().map((c) => c.args[1]); + assert.ok(envs.includes(pixiEnv), 'selection should be re-applied with the configured manager'); + assert.strictEqual(envs[envs.length - 1], pixiEnv, 'the configured manager should win'); + }); + + test('picks up a manager that registers while startup selection is still running', async () => { + listen(); + // The manager registers during the fallback discovery await, i.e. after it was found missing. + mockVenvManager.get.callsFake(async () => { + if (!pixiRegistered) { + pixiRegistered = true; + fireManagerEvent({ kind: 'registered', manager: mockPixiManager }); + } + return venvEnv; + }); + + await run(); + await settle(); + + const envs = mockEnvManagers.setEnvironment.getCalls().map((c) => c.args[1]); + assert.strictEqual(envs[0], venvEnv, 'the in-flight startup selection still completes first'); + assert.strictEqual(envs[envs.length - 1], pixiEnv, 'the queued re-run applies the configured manager last'); + }); + + test('ignores registrations of managers no setting is waiting for', async () => { + listen(); + await run(); + const callsBefore = mockEnvManagers.setEnvironment.callCount; + + const otherManager = { ...mockPixiManager, id: 'other.extension:other' }; + fireManagerEvent({ + kind: 'registered', + manager: otherManager as unknown as InternalEnvironmentManager, + }); + await settle(); + + assert.strictEqual(mockEnvManagers.setEnvironment.callCount, callsBefore, 'no re-run expected'); + }); + + test('does not switch the environment when the manager registers after the timeout', async () => { + let now = 1_000_000; + sandbox.stub(Date, 'now').callsFake(() => now); + listen(); + await run(); + const callsBefore = mockEnvManagers.setEnvironment.callCount; + + now += 30_001; // just past MANAGER_READY_TIMEOUT_MS + pixiRegistered = true; + fireManagerEvent({ kind: 'registered', manager: mockPixiManager }); + await settle(); + + assert.strictEqual( + mockEnvManagers.setEnvironment.callCount, + callsBefore, + 'a manager that registers long after startup must not switch the environment', + ); + }); + + test('only re-selects the workspace folders that fell back for the registered manager', async () => { + const otherUri = Uri.file('/test/other'); + (workspaceApis.getWorkspaceFolders as sinon.SinonStub).returns([ + { uri: testUri, name: 'test', index: 0 }, + { uri: otherUri, name: 'other', index: 1 }, + ]); + // Only the first folder is configured for Pixi; the other one uses venv, which is registered. + (helpers.getUserConfiguredSetting as sinon.SinonStub).callsFake( + (section: string, key: string, scope?: Uri) => { + if (section === 'python-envs' && key === 'defaultEnvManager') { + return scope?.toString() === otherUri.toString() ? VENV_MANAGER_ID : pixiManagerId; + } + return undefined; + }, + ); + listen(); + await run(); + const callsBefore = mockEnvManagers.setEnvironment.callCount; + + pixiRegistered = true; + fireManagerEvent({ kind: 'registered', manager: mockPixiManager }); + await settle(); + + const reselected = mockEnvManagers.setEnvironment + .getCalls() + .slice(callsBefore) + .map((c) => (c.args[0] as Uri).toString()); + assert.deepStrictEqual(reselected, [testUri.toString()], 'only the fallen-back folder is re-selected'); + assert.strictEqual(mockEnvManagers.setEnvironment.lastCall.args[1], pixiEnv); + assert.strictEqual(mockEnvManagers.setEnvironments.callCount, 1, 'global scope is not re-selected'); + }); +});