From 2719448cc6d77270fe5908147a2e5223822fb585 Mon Sep 17 00:00:00 2001 From: Stella Huang Date: Thu, 23 Jul 2026 15:50:45 -0700 Subject: [PATCH] Add inline script environment persistence Persist and safely rehydrate per-script environment associations. Harden selection races, cache-lock handling, corrupt-state repair, scope validation, and central batch selection consistency without globally serializing environment operations. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 1a9f6ba1-9bd3-4664-bc25-a0d34d7a2e91 --- src/features/envManagers.ts | 142 ++-- .../builtin/inlineScript/envManager.ts | 565 +++++++++++++++- .../envManagers.lastKnown.unit.test.ts | 145 +++- .../inlineScript/envManager.unit.test.ts | 619 +++++++++++++++++- 4 files changed, 1419 insertions(+), 52 deletions(-) diff --git a/src/features/envManagers.ts b/src/features/envManagers.ts index 3e005a1e..4d800a98 100644 --- a/src/features/envManagers.ts +++ b/src/features/envManagers.ts @@ -10,7 +10,7 @@ import { PythonProject, SetEnvironmentScope, } from '../api'; -import { SYSTEM_MANAGER_ID } from '../common/constants'; +import { INLINE_SCRIPT_MANAGER_ID, SYSTEM_MANAGER_ID } from '../common/constants'; import { EnvironmentManagerAlreadyRegisteredError, PackageManagerAlreadyRegisteredError, @@ -20,6 +20,7 @@ import { StopWatch } from '../common/stopWatch'; import { EventNames } from '../common/telemetry/constants'; import { sendTelemetryEvent } from '../common/telemetry/sender'; import { getCallingExtension } from '../common/utils/frameUtils'; +import { normalizePath } from '../common/utils/pathUtils'; import { DidChangeEnvironmentManagerEventArgs, DidChangePackageManagerEventArgs, @@ -62,6 +63,7 @@ export class PythonEnvironmentManagers implements EnvironmentManagers { * Only mutated by setEnvironment() / setEnvironments() / refreshEnvironment(). */ private readonly _activeSelection = new Map(); + private readonly _selectionRevisions = new Map(); private _onDidChangeEnvironmentManager = new EventEmitter(); private _onDidChangePackageManager = new EventEmitter(); @@ -216,8 +218,11 @@ export class PythonEnvironmentManagers implements EnvironmentManagers { // Fall back to cached environment's manager if no user-configured settings const project = context ? this.pm.get(context) : undefined; - const key = project ? project.uri.toString() : 'global'; - const cachedEnv = this._activeSelection.get(key); + const cachedEnv = + (context instanceof Uri + ? this._activeSelection.get(this.getInlineScriptSelectionKey(context)) + : undefined) ?? + this._activeSelection.get(project ? project.uri.toString() : 'global'); if (cachedEnv) { const cachedManager = this._environmentManagers.get(cachedEnv.envId.managerId); if (cachedManager) { @@ -335,9 +340,11 @@ export class PythonEnvironmentManagers implements EnvironmentManagers { traceError(this.managers.map((m) => m.id).join(', ')); return; } + const project = scope ? this.pm.get(scope) : undefined; + const key = this.getActiveSelectionKey(scope, manager, project); await manager.set(scope, environment); + this.bumpSelectionRevision(key); - const project = scope ? this.pm.get(scope) : undefined; // Only persist to settings when explicitly requested if (shouldPersistSettings && scope) { const packageManager = this.getPackageManager(environment); @@ -359,7 +366,6 @@ export class PythonEnvironmentManagers implements EnvironmentManagers { ); } - const key = project ? project.uri.toString() : 'global'; const oldEnv = this._activeSelection.get(key); if (oldEnv?.envId.id !== environment?.envId.id) { this._activeSelection.set(key, environment); @@ -367,7 +373,7 @@ export class PythonEnvironmentManagers implements EnvironmentManagers { setImmediate(() => { try { this._onDidChangeActiveEnvironment.fire({ - uri: project?.uri, + uri: this.getActiveSelectionUri(scope, manager, project), new: environment, old: oldEnv, }); @@ -407,11 +413,14 @@ export class PythonEnvironmentManagers implements EnvironmentManagers { return; } - const promises: Promise[] = []; const settings: EditAllManagerSettings[] = []; const events: DidChangeEnvironmentEventArgs[] = []; if (Array.isArray(scope) && scope.every((s) => s instanceof Uri)) { - promises.push(manager.set(scope, environment)); + await manager.set(scope, environment); + scope.forEach((uri) => { + const project = this.pm.get(uri); + this.bumpSelectionRevision(this.getActiveSelectionKey(uri, manager, project)); + }); scope.forEach((uri) => { const m = this.getEnvironmentManager(uri); // Always add settings when persisting, OR when manager differs @@ -424,16 +433,21 @@ export class PythonEnvironmentManagers implements EnvironmentManagers { } const project = this.pm.get(uri); - const key = project ? project.uri.toString() : 'global'; + const key = this.getActiveSelectionKey(uri, manager, project); const oldEnv = this._activeSelection.get(key); if (oldEnv?.envId.id !== environment?.envId.id) { this._activeSelection.set(key, environment); - events.push({ uri: project?.uri, new: environment, old: oldEnv }); + events.push({ + uri: this.getActiveSelectionUri(uri, manager, project), + new: environment, + old: oldEnv, + }); } }); } else if (typeof scope === 'string' && scope === 'global') { const m = this.getEnvironmentManager(undefined); - promises.push(manager.set(undefined, environment)); + await manager.set(undefined, environment); + this.bumpSelectionRevision('global'); // Always add settings when persisting, OR when manager differs if (shouldPersistSettings || manager.id !== m?.id) { settings.push({ @@ -449,7 +463,6 @@ export class PythonEnvironmentManagers implements EnvironmentManagers { events.push({ uri: undefined, new: environment, old: oldEnv }); } } - await Promise.all(promises); // Only persist to settings when explicitly requested if (shouldPersistSettings) { await setAllManagerSettings(settings); @@ -467,49 +480,51 @@ export class PythonEnvironmentManagers implements EnvironmentManagers { }); } } else { - const promises: Promise[] = []; const events: DidChangeEnvironmentEventArgs[] = []; if (Array.isArray(scope) && scope.every((s) => s instanceof Uri)) { + const groupedScopes = new Map(); scope.forEach((uri) => { const manager = this.getEnvironmentManager(uri); if (manager) { - const setAndAddEvent = async () => { - await manager.set(uri); - + groupedScopes.set(manager, [...(groupedScopes.get(manager) ?? []), uri]); + } + }); + for (const [manager, uris] of groupedScopes) { + await manager.set(uris); + uris.forEach((uri) => { + const project = this.pm.get(uri); + this.bumpSelectionRevision(this.getActiveSelectionKey(uri, manager, project)); + }); + await Promise.all( + uris.map(async (uri) => { const project = this.pm.get(uri); - - // Always get the new first, then compare with the old. This has minor impact on the ordering of - // events. But it ensures that we always get the latest environment at the time of this call. const newEnv = await manager.get(uri); - const key = project ? project.uri.toString() : 'global'; + const key = this.getActiveSelectionKey(uri, manager, project); const oldEnv = this._activeSelection.get(key); if (oldEnv?.envId.id !== newEnv?.envId.id) { this._activeSelection.set(key, newEnv); - events.push({ uri: project?.uri, new: newEnv, old: oldEnv }); + events.push({ + uri: this.getActiveSelectionUri(uri, manager, project), + new: newEnv, + old: oldEnv, + }); } - }; - promises.push(setAndAddEvent()); - } - }); + }), + ); + } } else if (typeof scope === 'string' && scope === 'global') { const manager = this.getEnvironmentManager(undefined); if (manager) { - const setAndAddEvent = async () => { - await manager.set(undefined); - - // Always get the new first, then compare with the old. This has minor impact on the ordering of - // events. But it ensures that we always get the latest environment at the time of this call. - const newEnv = await manager.get(undefined); - const oldEnv = this._activeSelection.get('global'); - if (oldEnv?.envId.id !== newEnv?.envId.id) { - this._activeSelection.set('global', newEnv); - events.push({ uri: undefined, new: newEnv, old: oldEnv }); - } - }; - promises.push(setAndAddEvent()); + await manager.set(undefined); + this.bumpSelectionRevision('global'); + const newEnv = await manager.get(undefined); + const oldEnv = this._activeSelection.get('global'); + if (oldEnv?.envId.id !== newEnv?.envId.id) { + this._activeSelection.set('global', newEnv); + events.push({ uri: undefined, new: newEnv, old: oldEnv }); + } } } - await Promise.all(promises); if (events.length > 0) { await new Promise((resolve, reject) => { setImmediate(() => { @@ -593,24 +608,67 @@ export class PythonEnvironmentManagers implements EnvironmentManagers { } const project = scope ? this.pm.get(scope) : undefined; + const key = this.getActiveSelectionKey(scope, manager, project); + const revision = (this._selectionRevisions.get(key) ?? 0) + 1; const newEnv = await manager.get(scope); + if ( + this.getEnvironmentManager(scope) !== manager || + (this._selectionRevisions.get(key) ?? 0) >= revision + ) { + return; + } + this._selectionRevisions.set(key, revision); - const key = project ? project.uri.toString() : 'global'; const oldEnv = this._activeSelection.get(key); if (oldEnv?.envId.id !== newEnv?.envId.id) { this._activeSelection.set(key, newEnv); setImmediate(() => - this._onDidChangeActiveEnvironment.fire({ uri: project?.uri, new: newEnv, old: oldEnv }), + this._onDidChangeActiveEnvironment.fire({ + uri: this.getActiveSelectionUri(scope, manager, project), + new: newEnv, + old: oldEnv, + }), ); } } getLastKnownEnvironment(scope: GetEnvironmentScope): PythonEnvironment | undefined { const project = scope ? this.pm.get(scope) : undefined; - const key = project ? project.uri.toString() : 'global'; + const manager = this.getEnvironmentManager(scope); + const key = this.getActiveSelectionKey(scope, manager, project); return this._activeSelection.get(key); } + private getActiveSelectionKey( + scope: GetEnvironmentScope, + manager: InternalEnvironmentManager | undefined, + project: PythonProject | undefined, + ): string { + return scope instanceof Uri && manager?.id === INLINE_SCRIPT_MANAGER_ID + ? this.getInlineScriptSelectionKey(scope) + : project + ? project.uri.toString() + : 'global'; + } + + private getActiveSelectionUri( + scope: GetEnvironmentScope, + manager: InternalEnvironmentManager, + project: PythonProject | undefined, + ): Uri | undefined { + return scope instanceof Uri && manager.id === INLINE_SCRIPT_MANAGER_ID ? scope : project?.uri; + } + + private getInlineScriptSelectionKey(scope: Uri): string { + return `inline-script:${normalizePath(scope.fsPath)}`; + } + + private bumpSelectionRevision(key: string): number { + const revision = (this._selectionRevisions.get(key) ?? 0) + 1; + this._selectionRevisions.set(key, revision); + return revision; + } + getProjectEnvManagers(uris: Uri[]): InternalEnvironmentManager[] { const projectEnvManagers: InternalEnvironmentManager[] = []; uris.forEach((uri) => { diff --git a/src/managers/builtin/inlineScript/envManager.ts b/src/managers/builtin/inlineScript/envManager.ts index 345bb127..df19201a 100644 --- a/src/managers/builtin/inlineScript/envManager.ts +++ b/src/managers/builtin/inlineScript/envManager.ts @@ -22,6 +22,7 @@ import { import { getErrorMessage } from '../../../common/errors/utils'; import { computeCacheKey } from '../../../common/inlineScript/cacheKey'; import { + CacheEnvironmentInspection, META_SCHEMA_VERSION, getBaseInterpreterStatus, getScriptEnvCacheRoot, @@ -37,8 +38,15 @@ import { matchesPythonVersion, readInlineScriptMetadataFromFile, } from '../../../common/inlineScript/metadata'; -import { CONDA_MANAGER_ID, PYENV_MANAGER_ID, SYSTEM_MANAGER_ID } from '../../../common/constants'; +import { + CONDA_MANAGER_ID, + ENVS_EXTENSION_ID, + INLINE_SCRIPT_MANAGER_ID, + PYENV_MANAGER_ID, + SYSTEM_MANAGER_ID, +} from '../../../common/constants'; import { acquireFileLock, AcquiredFileLock } from '../../../common/lockfile.apis'; +import { getWorkspacePersistentState, PersistentState } from '../../../common/persistentState'; import { isFileNotFoundError } from '../../../common/utils/filesystem'; import { normalizePath } from '../../../common/utils/pathUtils'; import { compareReleaseSegments, parseReleaseSegments } from '../../../common/utils/pep440Release'; @@ -54,6 +62,17 @@ const BASE_INTERPRETER_MANAGER_IDS = new Set([ const CACHE_LOCK_TIMEOUT_MS = 5 * 60 * 1000; const CACHE_LOCK_RETRY_MS = 500; +const CACHED_ASSOCIATION_VALIDATION_INTERVAL_MS = 5_000; +/** Workspace-state key for PEP 723 script path to environment executable associations. */ +export const INLINE_SCRIPT_ENVS_KEY = `${ENVS_EXTENSION_ID}:inline-script:SCRIPT_ENVIRONMENTS`; + +type PersistedInlineScriptEnvironments = Record; + +interface PersistedAssociationChange { + readonly scriptPath: string; + readonly environmentPath?: string; + readonly expectedEnvironmentPath?: string; +} interface SelectedBaseInterpreter { readonly environment: PythonEnvironment; @@ -79,6 +98,13 @@ type CacheEntryInspection = /** Manages extension-owned PEP 723 script environments. */ export class InlineScriptEnvManager implements EnvironmentManager, Disposable { private readonly pendingCreations = new Map>(); + private readonly pendingRehydrations = new Map>(); + private readonly fsPathToEnv = new Map(); + private readonly fsPathToPersistedEnvPath = new Map(); + private readonly cachedAssociationValidatedAt = new Map(); + private readonly associationRevisions = new Map(); + private persistenceQueue: Promise = Promise.resolve(); + private selectionQueue: Promise = Promise.resolve(); private readonly _onDidChangeEnvironments = new EventEmitter(); public readonly onDidChangeEnvironments: Event = @@ -174,12 +200,12 @@ export class InlineScriptEnvManager implements EnvironmentManager, Disposable { return []; } - async set(_scope: SetEnvironmentScope, _environment?: PythonEnvironment): Promise { - return; + async set(scope: SetEnvironmentScope, environment?: PythonEnvironment): Promise { + return this.enqueueSelection(() => this.setInternal(scope, environment)); } - async get(_scope: GetEnvironmentScope): Promise { - return undefined; + async get(scope: GetEnvironmentScope): Promise { + return this.getInternal(scope); } async resolve(_context: ResolveEnvironmentContext): Promise { @@ -191,6 +217,523 @@ export class InlineScriptEnvManager implements EnvironmentManager, Disposable { return uri?.scheme === 'file' ? uri : undefined; } + private async setInternal(scope: SetEnvironmentScope, environment?: PythonEnvironment): Promise { + const scripts = this.getScriptUris(scope); + if (scripts.length === 0) { + return; + } + + let environmentPath: string | undefined; + if (environment) { + const ownership = await this.inspectAssociationOwnership(environment); + if (ownership !== 'expected') { + const message = `Inline-script environment is not an owned cache entry: ${environment.environmentPath.fsPath}.`; + this.log.warn(message); + throw new Error(message); + } + environmentPath = environment.environmentPath.fsPath; + } + + const updates: { + readonly uri: Uri; + readonly scriptPath: string; + readonly before: PythonEnvironment | undefined; + readonly needsPersistence: boolean; + readonly shouldNotify: boolean; + }[] = []; + for (const script of scripts) { + const before = await this.getAssociationForMutation(script.scriptPath); + const hadPersistedAssociation = this.fsPathToPersistedEnvPath.has(script.scriptPath); + const hasSamePersistedEnvironment = + environmentPath !== undefined && + normalizePath(this.fsPathToPersistedEnvPath.get(script.scriptPath) ?? '') === + normalizePath(environmentPath); + const needsPersistence = environment ? !hasSamePersistedEnvironment : hadPersistedAssociation; + const shouldNotify = + (!this.isSameEnvironment(before, environment) && !hasSamePersistedEnvironment) || + (!environment && hadPersistedAssociation); + const hasPendingRehydration = this.pendingRehydrations.has(script.scriptPath); + const cached = this.fsPathToEnv.get(script.scriptPath); + const needsMemoryUpdate = environment ? cached !== environment : cached !== undefined; + if (needsPersistence || shouldNotify || hasPendingRehydration || needsMemoryUpdate) { + updates.push({ + ...script, + before, + needsPersistence, + shouldNotify, + }); + } + } + if (updates.length === 0) { + return; + } + + try { + const persistenceUpdates = updates.filter((update) => update.needsPersistence); + if (persistenceUpdates.length > 0) { + await this.updatePersistedAssociations( + persistenceUpdates.map((update) => ({ + scriptPath: update.scriptPath, + environmentPath, + })), + ); + } + } catch (error) { + this.log.error(`Failed to persist inline-script environment association: ${getErrorMessage(error)}`); + throw error; + } + + for (const update of updates) { + this.bumpAssociationRevision(update.scriptPath); + this.pendingRehydrations.delete(update.scriptPath); + if (environment) { + this.fsPathToEnv.set(update.scriptPath, environment); + this.fsPathToPersistedEnvPath.set(update.scriptPath, environmentPath!); + this.cachedAssociationValidatedAt.set(update.scriptPath, Date.now()); + } else { + this.fsPathToEnv.delete(update.scriptPath); + this.fsPathToPersistedEnvPath.delete(update.scriptPath); + this.cachedAssociationValidatedAt.delete(update.scriptPath); + } + if (update.shouldNotify) { + this._onDidChangeEnvironment.fire({ + uri: update.uri, + old: update.before, + new: environment, + }); + } + } + } + + private async getInternal(scope: GetEnvironmentScope): Promise { + if (!(scope instanceof Uri) || scope.scheme !== 'file') { + return undefined; + } + + // An unreadable or invalid metadata block is indistinguishable from a transient + // read failure, so retain the association but do not return it. + const metadata = await readInlineScriptMetadataFromFile(scope); + if (!metadata) { + return undefined; + } + + const environment = await this.getAssociation(normalizePath(scope.fsPath), scope); + if (!environment) { + return undefined; + } + + const requiresPython = metadata.requiresPython?.trim(); + return requiresPython && !matchesPythonVersion(requiresPython, environment.version) ? undefined : environment; + } + + private getScriptUris(scope: SetEnvironmentScope): { readonly uri: Uri; readonly scriptPath: string }[] { + const candidates = scope instanceof Uri ? [scope] : Array.isArray(scope) ? scope : undefined; + if ( + !candidates || + candidates.length === 0 || + candidates.some((candidate) => !(candidate instanceof Uri) || candidate.scheme !== 'file') + ) { + throw new Error('Inline-script environment selection requires one or more local file URIs.'); + } + + const scripts: { readonly uri: Uri; readonly scriptPath: string }[] = []; + const seen = new Set(); + for (const candidate of candidates) { + const scriptPath = normalizePath(candidate.fsPath); + if (!seen.has(scriptPath)) { + seen.add(scriptPath); + scripts.push({ uri: candidate, scriptPath }); + } + } + return scripts; + } + + private async getAssociation(scriptPath: string, scriptUri: Uri): Promise { + const pending = this.pendingRehydrations.get(scriptPath); + if (pending) { + return pending; + } + + const cached = this.fsPathToEnv.get(scriptPath); + if (cached) { + const validatedAt = this.cachedAssociationValidatedAt.get(scriptPath); + if ( + validatedAt !== undefined && + Date.now() - validatedAt < CACHED_ASSOCIATION_VALIDATION_INTERVAL_MS + ) { + return cached; + } + const validation = this.validateCachedAssociation(scriptPath, scriptUri, cached); + this.pendingRehydrations.set(scriptPath, validation); + try { + return await validation; + } finally { + if (this.pendingRehydrations.get(scriptPath) === validation) { + this.pendingRehydrations.delete(scriptPath); + } + } + } + + const revision = this.associationRevisions.get(scriptPath) ?? 0; + const rehydration = this.rehydrateAssociation(scriptPath, scriptUri, revision); + this.pendingRehydrations.set(scriptPath, rehydration); + try { + return await rehydration; + } finally { + if (this.pendingRehydrations.get(scriptPath) === rehydration) { + this.pendingRehydrations.delete(scriptPath); + } + } + } + + private async getAssociationForMutation(scriptPath: string): Promise { + const cached = this.fsPathToEnv.get(scriptPath); + if (cached) { + return cached; + } + await this.getPersistedAssociation(scriptPath); + return this.fsPathToEnv.get(scriptPath); + } + + private async validateCachedAssociation( + scriptPath: string, + scriptUri: Uri, + cached: PythonEnvironment, + ): Promise { + const environmentPath = cached.environmentPath.fsPath; + const envDirPath = path.dirname(path.dirname(environmentPath)); + if (await this.isCacheEntryBusy(envDirPath)) { + return undefined; + } + try { + const stat = await fs.stat(environmentPath); + if (stat.isFile()) { + const revision = this.associationRevisions.get(scriptPath) ?? 0; + const resolved = await resolveVenvPythonEnvironmentPath( + environmentPath, + this.nativeFinder, + this.api, + this, + this.baseManager, + ); + if (!this.isCurrentAssociationRevision(scriptPath, revision)) { + return this.fsPathToEnv.get(scriptPath); + } + if (!resolved) { + return undefined; + } + const ownership = await this.inspectAssociationOwnership(resolved); + if (!this.isCurrentAssociationRevision(scriptPath, revision)) { + return this.fsPathToEnv.get(scriptPath); + } + if (ownership === 'stale') { + await this.removeStalePersistedAssociation( + scriptPath, + environmentPath, + revision, + scriptUri, + ); + return undefined; + } + if (ownership !== 'expected') { + return undefined; + } + this.cachedAssociationValidatedAt.set(scriptPath, Date.now()); + if (cached.version === resolved.version) { + return cached; + } + this.fsPathToEnv.set(scriptPath, resolved); + this._onDidChangeEnvironment.fire({ uri: scriptUri, old: cached, new: resolved }); + return resolved; + } + if (!(await this.isCacheEntryBusy(envDirPath))) { + await this.removeStalePersistedAssociation( + scriptPath, + environmentPath, + this.associationRevisions.get(scriptPath) ?? 0, + scriptUri, + ); + } + } catch (error) { + if (this.isDefinitivelyStalePathError(error)) { + if (!(await this.isCacheEntryBusy(envDirPath))) { + await this.removeStalePersistedAssociation( + scriptPath, + environmentPath, + this.associationRevisions.get(scriptPath) ?? 0, + scriptUri, + ); + } + } else { + this.log.warn( + `Unable to inspect cached inline-script environment ${environmentPath}: ${getErrorMessage(error)}`, + ); + } + } + return undefined; + } + + private async rehydrateAssociation( + scriptPath: string, + scriptUri: Uri, + revision: number, + ): Promise { + let environmentPath: string | undefined; + try { + environmentPath = await this.getPersistedAssociation(scriptPath); + } catch (error) { + this.log.warn(`Failed to read inline-script environment association: ${getErrorMessage(error)}`); + return undefined; + } + if (!environmentPath) { + return undefined; + } + if (!this.isCurrentAssociationRevision(scriptPath, revision)) { + return this.fsPathToEnv.get(scriptPath); + } + if (!path.isAbsolute(environmentPath)) { + await this.removeStalePersistedAssociation(scriptPath, environmentPath, revision, scriptUri); + return undefined; + } + const envDirPath = path.dirname(path.dirname(environmentPath)); + if (await this.isCacheEntryBusy(envDirPath)) { + return undefined; + } + + try { + const stat = await fs.stat(environmentPath); + if (!stat.isFile()) { + if (!(await this.isCacheEntryBusy(envDirPath))) { + await this.removeStalePersistedAssociation(scriptPath, environmentPath, revision, scriptUri); + } + return undefined; + } + } catch (error) { + if (this.isDefinitivelyStalePathError(error)) { + if (!(await this.isCacheEntryBusy(envDirPath))) { + await this.removeStalePersistedAssociation(scriptPath, environmentPath, revision, scriptUri); + } + } else { + this.log.warn( + `Unable to inspect persisted inline-script environment ${environmentPath}: ${getErrorMessage(error)}`, + ); + } + return undefined; + } + + const resolved = await resolveVenvPythonEnvironmentPath( + environmentPath, + this.nativeFinder, + this.api, + this, + this.baseManager, + ); + if (!resolved) { + // PET/API resolution can fail transiently. Keep the association for a later retry. + return undefined; + } + + if (!this.isCurrentAssociationRevision(scriptPath, revision)) { + return this.fsPathToEnv.get(scriptPath); + } + const ownership = await this.inspectAssociationOwnership(resolved); + if (ownership === 'stale') { + await this.removeStalePersistedAssociation(scriptPath, environmentPath, revision, scriptUri); + return undefined; + } + if (ownership !== 'expected') { + return undefined; + } + + if (!this.isCurrentAssociationRevision(scriptPath, revision) || this.fsPathToEnv.has(scriptPath)) { + return this.fsPathToEnv.get(scriptPath); + } + this.fsPathToEnv.set(scriptPath, resolved); + this.cachedAssociationValidatedAt.set(scriptPath, Date.now()); + this._onDidChangeEnvironment.fire({ uri: scriptUri, old: undefined, new: resolved }); + return resolved; + } + + private async inspectAssociationOwnership(environment: PythonEnvironment): Promise { + if (environment.envId.managerId !== INLINE_SCRIPT_MANAGER_ID || !path.isAbsolute(environment.sysPrefix)) { + return 'uncertain'; + } + const cacheRoot = getScriptEnvCacheRoot(this.globalStorageUri); + const envDir = Uri.file(environment.sysPrefix); + try { + if (!(await resolveCacheEntryPath(cacheRoot, envDir))) { + return 'stale'; + } + } catch { + return 'uncertain'; + } + return inspectOwnedCacheEntry( + environment, + cacheRoot, + envDir, + ); + } + + private async getPersistedAssociation(scriptPath: string): Promise { + await this.persistenceQueue; + const state = await getWorkspacePersistentState(); + const raw = await state.get(INLINE_SCRIPT_ENVS_KEY); + if (raw === undefined) { + this.fsPathToPersistedEnvPath.delete(scriptPath); + return undefined; + } + const associations = this.asPersistedAssociations(raw); + if (!associations) { + await this.removeInvalidPersistedAssociation(scriptPath); + this.fsPathToPersistedEnvPath.delete(scriptPath); + return undefined; + } + const rawValue = (raw as Record)[scriptPath]; + if (rawValue !== undefined && (typeof rawValue !== 'string' || rawValue.length === 0)) { + await this.removeInvalidPersistedAssociation(scriptPath); + this.fsPathToPersistedEnvPath.delete(scriptPath); + return undefined; + } + const environmentPath = associations[scriptPath]; + if (environmentPath) { + this.fsPathToPersistedEnvPath.set(scriptPath, environmentPath); + } else { + this.fsPathToPersistedEnvPath.delete(scriptPath); + } + return environmentPath; + } + + private async removeStalePersistedAssociation( + scriptPath: string, + expectedEnvironmentPath: string, + revision: number, + scriptUri?: Uri, + ): Promise { + await this.enqueueSelection(async () => { + if (!this.isCurrentAssociationRevision(scriptPath, revision)) { + return; + } + try { + await this.updatePersistedAssociations([{ scriptPath, expectedEnvironmentPath }]); + if ( + this.fsPathToPersistedEnvPath.get(scriptPath) === expectedEnvironmentPath && + this.isCurrentAssociationRevision(scriptPath, revision) + ) { + const old = this.fsPathToEnv.get(scriptPath); + this.bumpAssociationRevision(scriptPath); + this.fsPathToEnv.delete(scriptPath); + this.fsPathToPersistedEnvPath.delete(scriptPath); + this.cachedAssociationValidatedAt.delete(scriptPath); + if (old && scriptUri) { + this._onDidChangeEnvironment.fire({ uri: scriptUri, old, new: undefined }); + } + } + } catch (error) { + this.log.warn( + `Failed to remove stale inline-script environment association: ${getErrorMessage(error)}`, + ); + } + }); + } + + private removeInvalidPersistedAssociation(scriptPath: string): Promise { + return this.enqueuePersistence(async (state) => { + const raw = await state.get(INLINE_SCRIPT_ENVS_KEY); + if (raw === undefined) { + return; + } + const associations = this.asPersistedAssociations(raw); + if (!associations) { + await state.set(INLINE_SCRIPT_ENVS_KEY, {}); + return; + } + const rawValue = (raw as Record)[scriptPath]; + if (rawValue !== undefined && (typeof rawValue !== 'string' || rawValue.length === 0)) { + delete associations[scriptPath]; + await state.set(INLINE_SCRIPT_ENVS_KEY, associations); + } + }); + } + + private updatePersistedAssociations(changes: readonly PersistedAssociationChange[]): Promise { + return this.enqueuePersistence(async (state) => { + const raw = await state.get(INLINE_SCRIPT_ENVS_KEY); + const associations = { ...(this.asPersistedAssociations(raw) ?? {}) }; + for (const change of changes) { + const current = associations[change.scriptPath]; + if (change.environmentPath) { + associations[change.scriptPath] = change.environmentPath; + } else if ( + change.expectedEnvironmentPath === undefined || + (current !== undefined && + normalizePath(current) === normalizePath(change.expectedEnvironmentPath)) + ) { + delete associations[change.scriptPath]; + } + } + await state.set(INLINE_SCRIPT_ENVS_KEY, associations); + }); + } + + private asPersistedAssociations(value: unknown): PersistedInlineScriptEnvironments | undefined { + if (!value || typeof value !== 'object' || Array.isArray(value)) { + return undefined; + } + const associations: PersistedInlineScriptEnvironments = {}; + for (const [scriptPath, environmentPath] of Object.entries(value)) { + if (typeof environmentPath === 'string' && environmentPath.length > 0) { + associations[scriptPath] = environmentPath; + } + } + return associations; + } + + private enqueuePersistence(operation: (state: PersistentState) => Promise): Promise { + const run = this.persistenceQueue.then(async () => operation(await getWorkspacePersistentState())); + this.persistenceQueue = run.catch(() => undefined); + return run; + } + + private enqueueSelection(operation: () => Promise): Promise { + const run = this.selectionQueue.then(operation); + this.selectionQueue = run.then( + () => undefined, + () => undefined, + ); + return run; + } + + private async isCacheEntryBusy(envDirPath: string): Promise { + return ( + this.pendingCreations.has(path.basename(envDirPath)) || + (await fs.pathExists(`${path.resolve(envDirPath)}.lock`)) + ); + } + + private bumpAssociationRevision(scriptPath: string): void { + this.associationRevisions.set(scriptPath, (this.associationRevisions.get(scriptPath) ?? 0) + 1); + } + + private isCurrentAssociationRevision(scriptPath: string, revision: number): boolean { + return (this.associationRevisions.get(scriptPath) ?? 0) === revision; + } + + private isSameEnvironment( + first: PythonEnvironment | undefined, + second: PythonEnvironment | undefined, + ): boolean { + if (first === second) { + return true; + } + if (!first || !second) { + return false; + } + return ( + first.envId.managerId === second.envId.managerId && + normalizePath(first.environmentPath.fsPath) === normalizePath(second.environmentPath.fsPath) + ); + } + private async selectBaseInterpreter(metadata: InlineScriptMetadata): Promise { const globalEnvironments = await this.api.getEnvironments('global'); const reported = globalEnvironments.filter( @@ -448,6 +991,18 @@ export class InlineScriptEnvManager implements EnvironmentManager, Disposable { } } + private isDefinitivelyStalePathError(error: unknown): boolean { + if (isFileNotFoundError(error)) { + return true; + } + return ( + typeof error === 'object' && + error !== null && + 'code' in error && + ['ENOTDIR', 'EINVAL', 'ERR_INVALID_ARG_VALUE'].includes((error as NodeJS.ErrnoException).code ?? '') + ); + } + private areEqualPythonReleases(actual: string, expected: string): boolean { const actualRelease = parseReleaseSegments(actual); const expectedRelease = parseReleaseSegments(expected); diff --git a/src/test/features/envManagers.lastKnown.unit.test.ts b/src/test/features/envManagers.lastKnown.unit.test.ts index eb44755b..d1f9efce 100644 --- a/src/test/features/envManagers.lastKnown.unit.test.ts +++ b/src/test/features/envManagers.lastKnown.unit.test.ts @@ -69,18 +69,22 @@ suite('PythonEnvironmentManagers getLastKnownEnvironment', () => { envManagers.dispose(); }); - function registerManager(getImpl: (scope: GetEnvironmentScope) => Promise): string { + function registerManager( + getImpl: (scope: GetEnvironmentScope) => Promise, + setImpl: EnvironmentManager['set'] = async () => undefined, + name = 'test-env-mgr', + ): string { const onDidChangeEnvironment = new EventEmitter(); const onDidChangeEnvironments = new EventEmitter(); const manager = { - name: 'test-env-mgr', + name, displayName: 'Test Env Manager', preferredPackageManagerId: 'ms-python.python:pip', onDidChangeEnvironment: onDidChangeEnvironment.event, onDidChangeEnvironments: onDidChangeEnvironments.event, get: getImpl, getEnvironments: async () => [], - set: async () => undefined, + set: setImpl, resolve: async () => undefined, refresh: async () => undefined, } as unknown as EnvironmentManager; @@ -119,4 +123,139 @@ suite('PythonEnvironmentManagers getLastKnownEnvironment', () => { await envManagers.refreshEnvironment(undefined); assert.strictEqual(envManagers.getLastKnownEnvironment(undefined)?.envId.id, 'env2'); }); + + test('does not update selection, settings, or events when a registered manager rejects a selection', async () => { + const scope = Uri.file('/workspace/script.py'); + const project = { name: 'script.py', uri: scope }; + projectManager.setup((pm) => pm.get(scope)).returns(() => project); + const managerSet = sinon.stub().rejects(new Error('Inline-script environment is not an owned cache entry.')); + const managerId = registerManager(async () => undefined, managerSet); + const rejected = { + ...makeEnv('unowned'), + envId: { id: 'unowned', managerId }, + }; + const settings = sinon.stub(settingHelpers, 'setAllManagerSettings').resolves(); + const events: DidChangeEnvironmentEventArgs[] = []; + envManagers.onDidChangeActiveEnvironment((event) => events.push(event)); + + await assert.rejects(envManagers.setEnvironment(scope, rejected), /not an owned cache entry/); + await new Promise((resolve) => setImmediate(resolve)); + + assert.strictEqual(envManagers.getLastKnownEnvironment(scope), undefined); + assert.strictEqual(settings.callCount, 0); + assert.strictEqual(events.length, 0); + }); + + test('does not publish batch or global selections when the manager rejects', async () => { + const scope = Uri.file('/workspace/script.py'); + const managerSet = sinon.stub().rejects(new Error('selection rejected')); + const managerId = registerManager(async () => undefined, managerSet); + const rejected = { + ...makeEnv('rejected'), + envId: { id: 'rejected', managerId }, + }; + const settings = sinon.stub(settingHelpers, 'setAllManagerSettings').resolves(); + const events: DidChangeEnvironmentEventArgs[] = []; + envManagers.onDidChangeActiveEnvironment((event) => events.push(event)); + + await assert.rejects(envManagers.setEnvironments([scope], rejected, false), /selection rejected/); + await assert.rejects(envManagers.setEnvironments('global', rejected, false), /selection rejected/); + await new Promise((resolve) => setImmediate(resolve)); + + assert.strictEqual(envManagers.getLastKnownEnvironment(scope), undefined); + assert.strictEqual(envManagers.getLastKnownEnvironment(undefined), undefined); + assert.strictEqual(settings.callCount, 0); + assert.strictEqual(events.length, 0); + }); + + test('passes same-manager batch unsets to the manager atomically', async () => { + const first = Uri.file('/workspace/first.py'); + const second = Uri.file('/workspace/second.py'); + const managerSet = sinon.stub().resolves(); + registerManager(async () => undefined, managerSet); + + await envManagers.setEnvironments([first, second], undefined, false); + + sinon.assert.calledOnceWithExactly(managerSet, [first, second], undefined); + }); + + test('does not let an older same-manager refresh overwrite a newer selection', async () => { + const initial = makeEnv('initial'); + let resolveStaleRefresh: ((environment: PythonEnvironment) => void) | undefined; + const staleRefresh = new Promise((resolve) => { + resolveStaleRefresh = resolve; + }); + const managerGet = sinon.stub(); + managerGet.onFirstCall().resolves(initial); + managerGet.onSecondCall().returns(staleRefresh); + const managerId = registerManager(managerGet); + const selected = { + ...makeEnv('selected'), + envId: { id: 'selected', managerId }, + }; + + await envManagers.refreshEnvironment(undefined); + const refresh = envManagers.refreshEnvironment(undefined); + await envManagers.setEnvironment(undefined, selected, false); + resolveStaleRefresh!(initial); + await refresh; + + assert.strictEqual(envManagers.getLastKnownEnvironment(undefined), selected); + }); + + test('retains an in-flight refresh when a concurrent selection fails', async () => { + const refreshed = makeEnv('refreshed'); + let resolveRefresh: ((environment: PythonEnvironment) => void) | undefined; + const managerGet = sinon.stub().returns( + new Promise((resolve) => { + resolveRefresh = resolve; + }), + ); + const managerId = registerManager(managerGet, sinon.stub().rejects(new Error('selection rejected'))); + const rejected = { + ...makeEnv('rejected'), + envId: { id: 'rejected', managerId }, + }; + + const refresh = envManagers.refreshEnvironment(undefined); + await assert.rejects(envManagers.setEnvironment(undefined, rejected, false), /selection rejected/); + resolveRefresh!(refreshed); + await refresh; + + assert.strictEqual(envManagers.getLastKnownEnvironment(undefined), refreshed); + }); + + test('tracks inline-script selections independently for scripts in the same project', async () => { + const firstUri = Uri.file('/workspace/first.py'); + const secondUri = Uri.file('/workspace/second.py'); + const managerId = registerManager(async () => undefined, async () => undefined, 'inline-script'); + const first = { ...makeEnv('first'), envId: { id: 'first', managerId } }; + const second = { ...makeEnv('second'), envId: { id: 'second', managerId } }; + + await envManagers.setEnvironment(firstUri, first, false); + await envManagers.setEnvironment(secondUri, second, false); + + assert.strictEqual(envManagers.getLastKnownEnvironment(firstUri), first); + assert.strictEqual(envManagers.getLastKnownEnvironment(secondUri), second); + }); + + test('retains an earlier successful refresh when a later refresh fails', async () => { + const refreshed = makeEnv('refreshed'); + let resolveFirst: ((environment: PythonEnvironment) => void) | undefined; + const managerGet = sinon.stub(); + managerGet.onFirstCall().returns( + new Promise((resolve) => { + resolveFirst = resolve; + }), + ); + managerGet.onSecondCall().rejects(new Error('refresh rejected')); + registerManager(managerGet); + + const first = envManagers.refreshEnvironment(undefined); + await assert.rejects(envManagers.refreshEnvironment(undefined), /refresh rejected/); + resolveFirst!(refreshed); + await first; + + assert.strictEqual(envManagers.getLastKnownEnvironment(undefined), refreshed); + }); }); diff --git a/src/test/managers/builtin/inlineScript/envManager.unit.test.ts b/src/test/managers/builtin/inlineScript/envManager.unit.test.ts index 4baf669f..a82470ef 100644 --- a/src/test/managers/builtin/inlineScript/envManager.unit.test.ts +++ b/src/test/managers/builtin/inlineScript/envManager.unit.test.ts @@ -12,9 +12,14 @@ import * as cacheKey from '../../../../common/inlineScript/cacheKey'; import * as cacheLayout from '../../../../common/inlineScript/cacheLayout'; import * as metadataReader from '../../../../common/inlineScript/metadata'; import * as lockfileApis from '../../../../common/lockfile.apis'; +import * as persistentState from '../../../../common/persistentState'; import { isWindows } from '../../../../common/utils/platformUtils'; +import { normalizePath } from '../../../../common/utils/pathUtils'; import { getVenvPythonPath } from '../../../../common/utils/virtualEnvironment'; -import { InlineScriptEnvManager } from '../../../../managers/builtin/inlineScript/envManager'; +import { + InlineScriptEnvManager, + INLINE_SCRIPT_ENVS_KEY, +} from '../../../../managers/builtin/inlineScript/envManager'; import * as venvUtils from '../../../../managers/builtin/venvUtils'; import { NativePythonFinder } from '../../../../managers/common/nativePythonFinder'; @@ -71,6 +76,7 @@ suite('InlineScriptEnvManager', () => { let baseExecutable: string; let baseManager: EnvironmentManager; let computeCacheKeyStub: sinon.SinonStub; + let clock: sinon.SinonFakeTimers; let createWithProgressStub: sinon.SinonStub; let globalStorageUri: Uri; let lockStub: sinon.SinonStub; @@ -84,6 +90,12 @@ suite('InlineScriptEnvManager', () => { let tempRoot: string; let baseInterpreterStatusStub: sinon.SinonStub; let writeMetaStub: sinon.SinonStub; + let workspaceState: { + get: sinon.SinonStub; + set: sinon.SinonStub; + clear: sinon.SinonStub; + }; + let persistedAssociations: unknown; setup(async () => { tempRoot = await fs.realpath(await fs.mkdtemp(path.join(os.tmpdir(), 'inline-script-manager-'))); @@ -96,6 +108,19 @@ suite('InlineScriptEnvManager', () => { api = { getEnvironments: apiGetEnvironmentsStub } as unknown as PythonEnvironmentApi; nativeFinder = {} as NativePythonFinder; baseManager = {} as EnvironmentManager; + persistedAssociations = undefined; + workspaceState = { + get: sinon.stub().callsFake(async (key: string) => { + return key === INLINE_SCRIPT_ENVS_KEY ? persistedAssociations : undefined; + }), + set: sinon.stub().callsFake(async (key: string, value: unknown) => { + if (key === INLINE_SCRIPT_ENVS_KEY) { + persistedAssociations = value; + } + }), + clear: sinon.stub(), + }; + sinon.stub(persistentState, 'getWorkspacePersistentState').resolves(workspaceState); readMetadataStub = sinon.stub(metadataReader, 'readInlineScriptMetadataFromFile').resolves(VALID_METADATA); computeCacheKeyStub = sinon.stub(cacheKey, 'computeCacheKey').returns(CACHE_KEY); @@ -121,7 +146,7 @@ suite('InlineScriptEnvManager', () => { }; }); - sinon.useFakeTimers({ now: NOW, toFake: ['Date'] }); + clock = sinon.useFakeTimers({ now: NOW, toFake: ['Date'] }); manager = new InlineScriptEnvManager(nativeFinder, api, baseManager, globalStorageUri, makeFakeLog()); }); @@ -143,6 +168,33 @@ suite('InlineScriptEnvManager', () => { inspectMetaStub.resolves({ kind: 'valid', metadata }); } + async function createOwnedEnvironment( + cacheKey: string = CACHE_KEY, + envId: string = `inline-${cacheKey}`, + ): Promise { + const location = cacheLayout.getScriptEnvDir(globalStorageUri, cacheKey).fsPath; + const executable = getVenvPythonPath(location); + await fs.outputFile(executable, ''); + return { + ...makeEnvironment('ms-python.python:inline-script', '3.12.4', executable, location), + envId: { managerId: 'ms-python.python:inline-script', id: envId }, + }; + } + + async function waitForStubCall(stub: sinon.SinonStub): Promise { + for (let attempt = 0; attempt < 20; attempt += 1) { + if (stub.called) { + return; + } + await new Promise((resolve) => setTimeout(resolve, 5)); + } + assert.fail('Expected the stub to be called'); + } + + function nextTurn(): Promise { + return new Promise((resolve) => setImmediate(resolve)); + } + suite('static metadata and deferred methods', () => { test('exposes creation but leaves later-phase methods empty', async () => { const asInterface: EnvironmentManager = manager; @@ -830,4 +882,567 @@ suite('InlineScriptEnvManager', () => { assert.doesNotThrow(() => manager.dispose()); }); }); + + suite('script association persistence', () => { + test('sets, gets, unsets, persists, and reports only actual selection changes', async () => { + const uri = scriptUri(); + const environment = await createOwnedEnvironment(); + const listener = sinon.spy(); + manager.onDidChangeEnvironment(listener); + + await manager.set(uri, environment); + + assert.strictEqual(await manager.get(uri), environment); + assert.deepStrictEqual(persistedAssociations, { + [normalizePath(uri.fsPath)]: environment.environmentPath.fsPath, + }); + assert.strictEqual(workspaceState.set.firstCall.args[0], INLINE_SCRIPT_ENVS_KEY); + assert.strictEqual(listener.callCount, 1); + assert.deepStrictEqual(listener.firstCall.args[0], { uri, old: undefined, new: environment }); + + await manager.set(uri, environment); + assert.strictEqual(listener.callCount, 1); + + await manager.set(uri, undefined); + assert.deepStrictEqual(persistedAssociations, {}); + assert.strictEqual(listener.callCount, 2); + assert.deepStrictEqual(listener.secondCall.args[0], { uri, old: environment, new: undefined }); + }); + + test('persists a batch atomically and reports each distinct script URI exactly once', async () => { + const first = scriptUri('first.py'); + const second = scriptUri('second.py'); + const environment = await createOwnedEnvironment(); + const listener = sinon.spy(); + manager.onDidChangeEnvironment(listener); + + await manager.set([first, second, first], environment); + + assert.deepStrictEqual(persistedAssociations, { + [normalizePath(first.fsPath)]: environment.environmentPath.fsPath, + [normalizePath(second.fsPath)]: environment.environmentPath.fsPath, + }); + assert.strictEqual(workspaceState.set.callCount, 1); + assert.strictEqual(listener.callCount, 2); + assert.strictEqual(listener.firstCall.args[0].uri, first); + assert.strictEqual(listener.secondCall.args[0].uri, second); + assert.strictEqual(await manager.get(first), environment); + assert.strictEqual(await manager.get(second), environment); + }); + + test('serializes concurrent selections so neither persisted association is lost', async () => { + const firstUri = scriptUri('first.py'); + const secondUri = scriptUri('second.py'); + const firstEnvironment = await createOwnedEnvironment(); + const secondEnvironment = await createOwnedEnvironment('fedcba9876543210'); + + await Promise.all([ + manager.set(firstUri, firstEnvironment), + manager.set(secondUri, secondEnvironment), + ]); + + assert.deepStrictEqual(persistedAssociations, { + [normalizePath(firstUri.fsPath)]: firstEnvironment.environmentPath.fsPath, + [normalizePath(secondUri.fsPath)]: secondEnvironment.environmentPath.fsPath, + }); + assert.strictEqual(await manager.get(firstUri), firstEnvironment); + assert.strictEqual(await manager.get(secondUri), secondEnvironment); + }); + + test('rehydrates a persisted owned association on demand after restart', async () => { + const uri = scriptUri(); + const persistedEnvironment = await createOwnedEnvironment(); + persistedAssociations = { [normalizePath(uri.fsPath)]: persistedEnvironment.environmentPath.fsPath }; + const rehydrated = { ...persistedEnvironment, envId: { ...persistedEnvironment.envId, id: 'rehydrated' } }; + resolveVenvStub.resolves(rehydrated); + const restarted = new InlineScriptEnvManager(nativeFinder, api, baseManager, globalStorageUri, makeFakeLog()); + + assert.strictEqual(await restarted.get(uri), rehydrated); + assert.strictEqual(resolveVenvStub.callCount, 1); + assert.deepStrictEqual(persistedAssociations, { + [normalizePath(uri.fsPath)]: persistedEnvironment.environmentPath.fsPath, + }); + + const listener = sinon.spy(); + restarted.onDidChangeEnvironment(listener); + await restarted.set(uri, persistedEnvironment); + assert.strictEqual(listener.callCount, 0, 'different generated IDs for the same executable are not a change'); + + restarted.dispose(); + }); + + test('notifies when a slow persisted association finishes rehydrating', async () => { + const uri = scriptUri(); + const environment = await createOwnedEnvironment(); + persistedAssociations = { [normalizePath(uri.fsPath)]: environment.environmentPath.fsPath }; + let resolveRehydration: ((value: PythonEnvironment) => void) | undefined; + resolveVenvStub.callsFake( + () => + new Promise((resolve) => { + resolveRehydration = resolve; + }), + ); + const listener = sinon.spy(); + manager.onDidChangeEnvironment(listener); + + const pending = manager.get(uri); + await waitForStubCall(resolveVenvStub); + assert.strictEqual(listener.callCount, 0); + resolveRehydration!(environment); + + assert.strictEqual(await pending, environment); + sinon.assert.calledOnceWithExactly(listener, { uri, old: undefined, new: environment }); + }); + + test('does not rewrite or notify when a restart reselects the same persisted executable', async () => { + const uri = scriptUri(); + const environment = await createOwnedEnvironment(); + persistedAssociations = { [normalizePath(uri.fsPath)]: environment.environmentPath.fsPath }; + const restarted = new InlineScriptEnvManager(nativeFinder, api, baseManager, globalStorageUri, makeFakeLog()); + const listener = sinon.spy(); + restarted.onDidChangeEnvironment(listener); + + await restarted.set(uri, environment); + + assert.deepStrictEqual(persistedAssociations, { + [normalizePath(uri.fsPath)]: environment.environmentPath.fsPath, + }); + assert.strictEqual(workspaceState.set.callCount, 0); + assert.strictEqual(listener.callCount, 0); + assert.strictEqual(resolveVenvStub.callCount, 0); + + restarted.dispose(); + }); + + test('does not return a retained association when current metadata no longer accepts its Python version', async () => { + const uri = scriptUri(); + const environment = await createOwnedEnvironment(); + await manager.set(uri, environment); + readMetadataStub.resolves({ ...VALID_METADATA, requiresPython: '==3.11.*' }); + + assert.strictEqual(await manager.get(uri), undefined); + assert.deepStrictEqual(persistedAssociations, { + [normalizePath(uri.fsPath)]: environment.environmentPath.fsPath, + }); + + readMetadataStub.resolves(VALID_METADATA); + assert.strictEqual(await manager.get(uri), environment); + }); + + test('does not resolve or discard an association when metadata is absent or unreadable', async () => { + const uri = scriptUri(); + const environment = await createOwnedEnvironment(); + persistedAssociations = { [normalizePath(uri.fsPath)]: environment.environmentPath.fsPath }; + readMetadataStub.resolves(undefined); + + assert.strictEqual(await manager.get(uri), undefined); + assert.strictEqual(resolveVenvStub.callCount, 0); + assert.deepStrictEqual(persistedAssociations, { + [normalizePath(uri.fsPath)]: environment.environmentPath.fsPath, + }); + }); + + test('preserves a cold persisted association while its cache entry is locked', async () => { + const uri = scriptUri(); + const environment = await createOwnedEnvironment(); + persistedAssociations = { [normalizePath(uri.fsPath)]: environment.environmentPath.fsPath }; + const lockPath = `${path.resolve(environment.sysPrefix)}.lock`; + await fs.ensureDir(lockPath); + + assert.strictEqual(await manager.get(uri), undefined); + assert.deepStrictEqual(persistedAssociations, { + [normalizePath(uri.fsPath)]: environment.environmentPath.fsPath, + }); + assert.strictEqual(workspaceState.set.callCount, 0); + assert.strictEqual(resolveVenvStub.callCount, 0); + }); + + test('removes and notifies for a warm association whose executable was deleted', async () => { + const uri = scriptUri(); + const environment = await createOwnedEnvironment(); + await manager.set(uri, environment); + const listener = sinon.spy(); + manager.onDidChangeEnvironment(listener); + await fs.remove(environment.environmentPath.fsPath); + clock.tick(5_000); + + assert.strictEqual(await manager.get(uri), undefined); + assert.deepStrictEqual(persistedAssociations, {}); + sinon.assert.calledOnceWithExactly(listener, { uri, old: environment, new: undefined }); + }); + + test('preserves a warm association while its cache entry is locked', async () => { + const uri = scriptUri(); + const environment = await createOwnedEnvironment(); + await manager.set(uri, environment); + const listener = sinon.spy(); + manager.onDidChangeEnvironment(listener); + await fs.remove(environment.environmentPath.fsPath); + await fs.ensureDir(`${path.resolve(environment.sysPrefix)}.lock`); + clock.tick(5_000); + + assert.strictEqual(await manager.get(uri), undefined); + assert.deepStrictEqual(persistedAssociations, { + [normalizePath(uri.fsPath)]: environment.environmentPath.fsPath, + }); + assert.strictEqual(listener.callCount, 0); + }); + + test('refreshes a warm association rebuilt at the same cache path', async () => { + const uri = scriptUri(); + const environment = await createOwnedEnvironment(); + await manager.set(uri, environment); + const rebuilt = { + ...environment, + envId: { ...environment.envId, id: 'rebuilt' }, + version: '3.13.1', + }; + resolveVenvStub.resolves(rebuilt); + const listener = sinon.spy(); + manager.onDidChangeEnvironment(listener); + clock.tick(5_000); + + assert.strictEqual(await manager.get(uri), rebuilt); + sinon.assert.calledOnceWithExactly(listener, { uri, old: environment, new: rebuilt }); + }); + + test('retains warm environment identity when validation finds the same version', async () => { + const uri = scriptUri(); + const environment = await createOwnedEnvironment(); + await manager.set(uri, environment); + resolveVenvStub.resolves({ + ...environment, + envId: { ...environment.envId, id: 'new-generated-id' }, + }); + const listener = sinon.spy(); + manager.onDidChangeEnvironment(listener); + clock.tick(5_000); + + assert.strictEqual(await manager.get(uri), environment); + assert.strictEqual(listener.callCount, 0); + }); + + test('coalesces concurrent validation of an expired warm association', async () => { + const uri = scriptUri(); + const environment = await createOwnedEnvironment(); + await manager.set(uri, environment); + const rebuilt = { + ...environment, + envId: { ...environment.envId, id: 'rebuilt' }, + version: '3.13.1', + }; + let resolveValidation: ((value: PythonEnvironment) => void) | undefined; + resolveVenvStub.callsFake( + () => + new Promise((resolve) => { + resolveValidation = resolve; + }), + ); + const listener = sinon.spy(); + manager.onDidChangeEnvironment(listener); + clock.tick(5_000); + + const first = manager.get(uri); + const second = manager.get(uri); + await waitForStubCall(resolveVenvStub); + resolveValidation!(rebuilt); + + assert.deepStrictEqual(await Promise.all([first, second]), [rebuilt, rebuilt]); + assert.strictEqual(resolveVenvStub.callCount, 1); + sinon.assert.calledOnceWithExactly(listener, { uri, old: environment, new: rebuilt }); + }); + + test('unsets a persisted association after transient rehydration failure', async () => { + const uri = scriptUri(); + const environment = await createOwnedEnvironment(); + const listener = sinon.spy(); + manager.onDidChangeEnvironment(listener); + persistedAssociations = { [normalizePath(uri.fsPath)]: environment.environmentPath.fsPath }; + resolveVenvStub.resolves(undefined); + + assert.strictEqual(await manager.get(uri), undefined); + assert.strictEqual(resolveVenvStub.callCount, 1); + + await manager.set(uri, undefined); + assert.deepStrictEqual(persistedAssociations, {}); + assert.strictEqual(listener.callCount, 1); + + resolveVenvStub.resetHistory(); + assert.strictEqual(await manager.get(uri), undefined); + assert.strictEqual(resolveVenvStub.callCount, 0); + }); + + test('removes definitively stale or corrupt persisted paths but preserves transient resolution failures', async () => { + const staleUri = scriptUri('stale.py'); + persistedAssociations = { [normalizePath(staleUri.fsPath)]: path.join(tempRoot, 'missing-python') }; + + assert.strictEqual(await manager.get(staleUri), undefined); + assert.deepStrictEqual(persistedAssociations, {}); + + const corruptUri = scriptUri('corrupt.py'); + persistedAssociations = { [normalizePath(corruptUri.fsPath)]: 'not-an-absolute-path' }; + assert.strictEqual(await manager.get(corruptUri), undefined); + assert.deepStrictEqual(persistedAssociations, {}); + + const transientUri = scriptUri('transient.py'); + const environment = await createOwnedEnvironment(); + persistedAssociations = { [normalizePath(transientUri.fsPath)]: environment.environmentPath.fsPath }; + resolveVenvStub.resolves(undefined); + + assert.strictEqual(await manager.get(transientUri), undefined); + assert.deepStrictEqual(persistedAssociations, { + [normalizePath(transientUri.fsPath)]: environment.environmentPath.fsPath, + }); + + persistedAssociations = ['corrupt state']; + assert.strictEqual(await manager.get(scriptUri('corrupt-state.py')), undefined); + assert.deepStrictEqual(persistedAssociations, {}); + }); + + test('does not let stale corrupt-state repair delete a newer valid association', async () => { + const uri = scriptUri(); + const environment = await createOwnedEnvironment(); + const scriptPath = normalizePath(uri.fsPath); + persistedAssociations = { [scriptPath]: 42 }; + workspaceState.get.onSecondCall().callsFake(async () => { + persistedAssociations = { [scriptPath]: environment.environmentPath.fsPath }; + return persistedAssociations; + }); + + assert.strictEqual(await manager.get(uri), undefined); + assert.deepStrictEqual(persistedAssociations, { + [scriptPath]: environment.environmentPath.fsPath, + }); + assert.strictEqual(workspaceState.set.callCount, 0); + }); + + test('preserves an association when fallback resolution reports another manager', async () => { + const uri = scriptUri(); + const environment = await createOwnedEnvironment(); + persistedAssociations = { [normalizePath(uri.fsPath)]: environment.environmentPath.fsPath }; + resolveVenvStub.resolves({ + ...environment, + envId: { ...environment.envId, managerId: 'ms-python.python:system' }, + }); + + assert.strictEqual(await manager.get(uri), undefined); + assert.deepStrictEqual(persistedAssociations, { + [normalizePath(uri.fsPath)]: environment.environmentPath.fsPath, + }); + assert.strictEqual(workspaceState.set.callCount, 0); + }); + + test('rejects resolved and selected environments that are outside the owned cache', async () => { + const uri = scriptUri(); + const listener = sinon.spy(); + manager.onDidChangeEnvironment(listener); + const outsideDir = path.join(tempRoot, 'outside'); + const outsideExecutable = getVenvPythonPath(outsideDir); + await fs.outputFile(outsideExecutable, ''); + await fs.ensureDir(cacheLayout.getScriptEnvCacheRoot(globalStorageUri).fsPath); + const unowned = makeEnvironment( + 'ms-python.python:inline-script', + '3.12.4', + outsideExecutable, + outsideDir, + ); + persistedAssociations = { [normalizePath(uri.fsPath)]: outsideExecutable }; + resolveVenvStub.resolves(unowned); + + assert.strictEqual(await manager.get(uri), undefined); + assert.deepStrictEqual(persistedAssociations, {}); + workspaceState.set.resetHistory(); + + await assert.rejects(manager.set(uri, unowned), /not an owned cache entry/); + assert.deepStrictEqual(persistedAssociations, {}); + assert.strictEqual(workspaceState.set.callCount, 0); + assert.strictEqual(listener.callCount, 0); + }); + + test('normalizes script paths and treats same-ID environments at different paths as different selections', async function () { + if (!isWindows()) { + this.skip(); + } + const uri = scriptUri('CaseSensitive.py'); + const differentlyCased = Uri.file(uri.fsPath.toUpperCase()); + const first = await createOwnedEnvironment(CACHE_KEY, 'duplicate-id'); + const second = await createOwnedEnvironment('fedcba9876543210', 'duplicate-id'); + const listener = sinon.spy(); + manager.onDidChangeEnvironment(listener); + + await manager.set(uri, first); + assert.strictEqual(await manager.get(differentlyCased), first); + + await manager.set(differentlyCased, second); + assert.strictEqual(await manager.get(uri), second); + assert.strictEqual(listener.callCount, 2); + assert.strictEqual(listener.secondCall.args[0].uri, differentlyCased); + assert.strictEqual(listener.secondCall.args[0].old, first); + assert.strictEqual(listener.secondCall.args[0].new, second); + }); + + test('keeps the prior in-memory association and emits no event when persistence fails', async () => { + const uri = scriptUri(); + const first = await createOwnedEnvironment(); + const second = await createOwnedEnvironment('fedcba9876543210'); + const listener = sinon.spy(); + manager.onDidChangeEnvironment(listener); + + await manager.set(uri, first); + workspaceState.set.onSecondCall().rejects(new Error('Memento unavailable')); + await assert.rejects(manager.set(uri, second), /Memento unavailable/); + + assert.strictEqual(await manager.get(uri), first); + assert.deepStrictEqual(persistedAssociations, { + [normalizePath(uri.fsPath)]: first.environmentPath.fsPath, + }); + assert.strictEqual(listener.callCount, 1); + }); + + test('rejects a failed unset without changing its in-memory association or firing an event', async () => { + const uri = scriptUri(); + const environment = await createOwnedEnvironment(); + const listener = sinon.spy(); + manager.onDidChangeEnvironment(listener); + + await manager.set(uri, environment); + workspaceState.set.onSecondCall().rejects(new Error('Memento unavailable')); + + await assert.rejects(manager.set(uri, undefined), /Memento unavailable/); + assert.strictEqual(await manager.get(uri), environment); + assert.deepStrictEqual(persistedAssociations, { + [normalizePath(uri.fsPath)]: environment.environmentPath.fsPath, + }); + assert.strictEqual(listener.callCount, 1); + }); + + test('does not block a cached lookup behind another script rehydration', async () => { + const slowUri = scriptUri('slow.py'); + const cachedUri = scriptUri('cached.py'); + const slowEnvironment = await createOwnedEnvironment(); + const cachedEnvironment = await createOwnedEnvironment('fedcba9876543210'); + persistedAssociations = { [normalizePath(slowUri.fsPath)]: slowEnvironment.environmentPath.fsPath }; + await manager.set(cachedUri, cachedEnvironment); + + let resolveSlow: ((value: PythonEnvironment | undefined) => void) | undefined; + resolveVenvStub.callsFake( + () => + new Promise((resolve) => { + resolveSlow = resolve; + }), + ); + const slowGet = manager.get(slowUri); + await waitForStubCall(resolveVenvStub); + + const cachedResult = await Promise.race([ + manager.get(cachedUri).then((value) => ({ kind: 'cached' as const, value })), + nextTurn().then(() => ({ kind: 'blocked' as const, value: undefined })), + ]); + assert.strictEqual(cachedResult.kind, 'cached'); + assert.strictEqual(cachedResult.value, cachedEnvironment); + + resolveSlow!(slowEnvironment); + assert.strictEqual(await slowGet, slowEnvironment); + }); + + test('lets an unset win over a pending stale rehydration', async () => { + const uri = scriptUri(); + const environment = await createOwnedEnvironment(); + persistedAssociations = { [normalizePath(uri.fsPath)]: environment.environmentPath.fsPath }; + + let resolvePending: ((value: PythonEnvironment | undefined) => void) | undefined; + resolveVenvStub.callsFake( + () => + new Promise((resolve) => { + resolvePending = resolve; + }), + ); + const pendingGet = manager.get(uri); + await waitForStubCall(resolveVenvStub); + + await manager.set(uri, undefined); + assert.deepStrictEqual(persistedAssociations, {}); + + resolvePending!(environment); + assert.strictEqual(await pendingGet, undefined); + assert.strictEqual(await manager.get(uri), undefined); + }); + + test('lets a same-path selection supersede a pending stale rehydration', async () => { + const uri = scriptUri(); + const environment = await createOwnedEnvironment(); + persistedAssociations = { [normalizePath(uri.fsPath)]: environment.environmentPath.fsPath }; + + let resolvePending: ((value: PythonEnvironment | undefined) => void) | undefined; + resolveVenvStub.callsFake( + () => + new Promise((resolve) => { + resolvePending = resolve; + }), + ); + const pendingGet = manager.get(uri); + await waitForStubCall(resolveVenvStub); + + await manager.set(uri, environment); + const stale = { + ...environment, + envId: { ...environment.envId, managerId: 'ms-python.python:system' }, + }; + resolvePending!(stale); + + assert.strictEqual(await pendingGet, environment); + assert.strictEqual(await manager.get(uri), environment); + assert.deepStrictEqual(persistedAssociations, { + [normalizePath(uri.fsPath)]: environment.environmentPath.fsPath, + }); + assert.strictEqual(workspaceState.set.callCount, 0); + }); + + test('retains a pending rehydration when a competing persistence write fails', async () => { + const uri = scriptUri(); + const oldEnvironment = await createOwnedEnvironment(); + const newEnvironment = await createOwnedEnvironment('fedcba9876543210'); + persistedAssociations = { [normalizePath(uri.fsPath)]: oldEnvironment.environmentPath.fsPath }; + const listener = sinon.spy(); + manager.onDidChangeEnvironment(listener); + + let resolvePending: ((value: PythonEnvironment | undefined) => void) | undefined; + resolveVenvStub.callsFake( + () => + new Promise((resolve) => { + resolvePending = resolve; + }), + ); + const pendingGet = manager.get(uri); + await waitForStubCall(resolveVenvStub); + + workspaceState.set.onFirstCall().rejects(new Error('Memento unavailable')); + await assert.rejects(manager.set(uri, newEnvironment), /Memento unavailable/); + + resolvePending!(oldEnvironment); + assert.strictEqual(await pendingGet, oldEnvironment); + assert.strictEqual(await manager.get(uri), oldEnvironment); + assert.deepStrictEqual(persistedAssociations, { + [normalizePath(uri.fsPath)]: oldEnvironment.environmentPath.fsPath, + }); + sinon.assert.calledOnceWithExactly(listener, { uri, old: undefined, new: oldEnvironment }); + }); + + test('rejects invalid scopes atomically and never writes workspace state', async () => { + const environment = await createOwnedEnvironment(); + const valid = scriptUri(); + + await assert.rejects(manager.set(undefined, environment), /one or more local file URIs/); + await assert.rejects(manager.set(Uri.parse('untitled:script.py'), environment), /one or more local file URIs/); + await assert.rejects( + manager.set([valid, Uri.parse('untitled:script.py')], environment), + /one or more local file URIs/, + ); + + assert.strictEqual(workspaceState.set.callCount, 0); + assert.strictEqual(await manager.get(valid), undefined); + assert.strictEqual(await manager.get(undefined), undefined); + assert.strictEqual(await manager.get(Uri.parse('untitled:script.py')), undefined); + }); + }); });