Skip to content

Commit 6ea1f11

Browse files
StellaHuang95StellaHuang95
committed
Address uv fallback review
Coalesce identical inline-script setup requests before interpreter installation and cover uv bootstrap success and failure paths. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 1a9f6ba1-9bd3-4664-bc25-a0d34d7a2e91
1 parent c6abcf8 commit 6ea1f11

3 files changed

Lines changed: 193 additions & 29 deletions

File tree

‎src/managers/builtin/inlineScript/envManager.ts‎

Lines changed: 64 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ import {
2121
SetEnvironmentScope,
2222
} from '../../../api';
2323
import { getErrorMessage } from '../../../common/errors/utils';
24-
import { computeCacheKey } from '../../../common/inlineScript/cacheKey';
24+
import { computeCacheKey, normalizeDependency } from '../../../common/inlineScript/cacheKey';
2525
import {
2626
META_SCHEMA_VERSION,
2727
getBaseInterpreterStatus,
@@ -77,6 +77,7 @@ type CacheEntryInspection =
7777

7878
/** Manages extension-owned PEP 723 script environments. */
7979
export class InlineScriptEnvManager implements EnvironmentManager, Disposable {
80+
private readonly pendingSetups = new Map<string, Promise<PythonEnvironment | undefined>>();
8081
private readonly pendingCreations = new Map<string, Promise<PythonEnvironment | undefined>>();
8182
private readonly directlyResolvedBaseInterpreters = new Map<string, PythonEnvironment>();
8283
private baseInterpreterInstallationQueue: Promise<void> = Promise.resolve();
@@ -132,36 +133,19 @@ export class InlineScriptEnvManager implements EnvironmentManager, Disposable {
132133
return undefined;
133134
}
134135

135-
let selectedBase = await this.selectBaseInterpreter(metadata);
136-
if (!selectedBase && options?.quickCreate !== true) {
137-
selectedBase = await this.installAndSelectBaseInterpreter(metadata);
138-
}
139-
if (!selectedBase) {
140-
this.log.warn(`No compatible Python is available for inline-script environment creation: ${scriptUri.fsPath}.`);
141-
return undefined;
142-
}
143-
144-
const cacheKey = computeCacheKey({
145-
dependencies: packages,
146-
interpreterPath: selectedBase.canonicalPath,
147-
});
148-
const pending = this.pendingCreations.get(cacheKey);
136+
const setupKey = this.getPendingSetupKey(scriptUri, metadata, packages, options);
137+
const pending = this.pendingSetups.get(setupKey);
149138
if (pending) {
150139
return await pending;
151140
}
152141

153-
const creation = this.createOrReuseEnvironment({
154-
cacheKey,
155-
packages,
156-
metadata,
157-
selectedBase,
158-
});
159-
this.pendingCreations.set(cacheKey, creation);
142+
const setup = this.createForScript(scriptUri, metadata, packages, options);
143+
this.pendingSetups.set(setupKey, setup);
160144
try {
161-
return await creation;
145+
return await setup;
162146
} finally {
163-
if (this.pendingCreations.get(cacheKey) === creation) {
164-
this.pendingCreations.delete(cacheKey);
147+
if (this.pendingSetups.get(setupKey) === setup) {
148+
this.pendingSetups.delete(setupKey);
165149
}
166150
}
167151
} catch (error) {
@@ -170,6 +154,61 @@ export class InlineScriptEnvManager implements EnvironmentManager, Disposable {
170154
}
171155
}
172156

157+
private async createForScript(
158+
scriptUri: Uri,
159+
metadata: InlineScriptMetadata,
160+
packages: readonly string[],
161+
options?: CreateEnvironmentOptions,
162+
): Promise<PythonEnvironment | undefined> {
163+
let selectedBase = await this.selectBaseInterpreter(metadata);
164+
if (!selectedBase && options?.quickCreate !== true) {
165+
selectedBase = await this.installAndSelectBaseInterpreter(metadata);
166+
}
167+
if (!selectedBase) {
168+
this.log.warn(`No compatible Python is available for inline-script environment creation: ${scriptUri.fsPath}.`);
169+
return undefined;
170+
}
171+
172+
const cacheKey = computeCacheKey({
173+
dependencies: packages,
174+
interpreterPath: selectedBase.canonicalPath,
175+
});
176+
const pending = this.pendingCreations.get(cacheKey);
177+
if (pending) {
178+
return await pending;
179+
}
180+
181+
const creation = this.createOrReuseEnvironment({
182+
cacheKey,
183+
packages,
184+
metadata,
185+
selectedBase,
186+
});
187+
this.pendingCreations.set(cacheKey, creation);
188+
try {
189+
return await creation;
190+
} finally {
191+
if (this.pendingCreations.get(cacheKey) === creation) {
192+
this.pendingCreations.delete(cacheKey);
193+
}
194+
}
195+
}
196+
197+
private getPendingSetupKey(
198+
scriptUri: Uri,
199+
metadata: InlineScriptMetadata,
200+
packages: readonly string[],
201+
options: CreateEnvironmentOptions | undefined,
202+
): string {
203+
const normalizedPackages = Array.from(new Set(packages.map(normalizeDependency))).sort();
204+
return JSON.stringify([
205+
normalizePath(scriptUri.fsPath),
206+
metadata.requiresPython?.trim() ?? '',
207+
options?.quickCreate === true ? 'quick' : 'interactive',
208+
normalizedPackages,
209+
]);
210+
}
211+
173212
async refresh(_scope: RefreshEnvironmentsScope): Promise<void> {
174213
return;
175214
}

‎src/test/managers/builtin/inlineScript/envManager.unit.test.ts‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -658,6 +658,70 @@ suite('InlineScriptEnvManager', () => {
658658
assert.strictEqual(lockStub.callCount, 0);
659659
});
660660

661+
test('coalesces the full concurrent setup for the same script', async () => {
662+
const uri = scriptUri();
663+
const uvExecutable = path.join(tempRoot, 'uv-python', isWindows() ? 'python.exe' : 'python');
664+
await fs.outputFile(uvExecutable, '');
665+
const uvBase = makeEnvironment('ms-python.python:system', '3.13.1', uvExecutable);
666+
readMetadataStub.resolves({ ...VALID_METADATA, requiresPython: '>=3.13' });
667+
let installed = false;
668+
apiGetEnvironmentsStub.callsFake(async () => (installed ? [uvBase] : [baseEnvironment]));
669+
let releaseInstall: (() => void) | undefined;
670+
let signalPrompt: (() => void) | undefined;
671+
const promptShown = new Promise<void>((resolve) => {
672+
signalPrompt = resolve;
673+
});
674+
const installGate = new Promise<void>((resolve) => {
675+
releaseInstall = resolve;
676+
});
677+
promptInstallPythonViaUvStub.callsFake(async () => {
678+
signalPrompt!();
679+
await installGate;
680+
installed = true;
681+
return uvExecutable;
682+
});
683+
684+
const first = manager.create(uri);
685+
await promptShown;
686+
const second = manager.create(uri);
687+
releaseInstall!();
688+
const [firstResult, secondResult] = await Promise.all([first, second]);
689+
690+
assert.ok(firstResult);
691+
assert.strictEqual(firstResult, secondResult);
692+
assert.strictEqual(promptInstallPythonViaUvStub.callCount, 1);
693+
assert.strictEqual(apiRefreshEnvironmentsStub.callCount, 1);
694+
assert.strictEqual(createWithProgressStub.callCount, 1);
695+
});
696+
697+
test('coalesces concurrent setup requests for the same script when installation is declined', async () => {
698+
const uri = scriptUri();
699+
readMetadataStub.resolves({ ...VALID_METADATA, requiresPython: '>=3.13' });
700+
apiGetEnvironmentsStub.resolves([baseEnvironment]);
701+
let finishPrompt: ((value: undefined) => void) | undefined;
702+
let signalPrompt: (() => void) | undefined;
703+
const promptShown = new Promise<void>((resolve) => {
704+
signalPrompt = resolve;
705+
});
706+
promptInstallPythonViaUvStub.callsFake(
707+
() =>
708+
new Promise<undefined>((resolve) => {
709+
signalPrompt!();
710+
finishPrompt = resolve;
711+
}),
712+
);
713+
714+
const first = manager.create(uri);
715+
await promptShown;
716+
const second = manager.create(uri);
717+
finishPrompt!(undefined);
718+
assert.deepStrictEqual(await Promise.all([first, second]), [undefined, undefined]);
719+
assert.deepStrictEqual(await Promise.all([first, second]), [undefined, undefined]);
720+
assert.strictEqual(promptInstallPythonViaUvStub.callCount, 1);
721+
assert.strictEqual(apiRefreshEnvironmentsStub.callCount, 0);
722+
assert.strictEqual(createWithProgressStub.callCount, 0);
723+
});
724+
661725
test('coalesces simultaneous fallback requests for the same Python version', async () => {
662726
const uvExecutable = path.join(tempRoot, 'uv-python', isWindows() ? 'python.exe' : 'python');
663727
await fs.outputFile(uvExecutable, '');

‎src/test/managers/builtin/uvPythonInstaller.unit.test.ts‎

Lines changed: 65 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import { MockChildProcess } from '../../mocks/mockChildProcess';
2525
suite('uvPythonInstaller - promptInstallPythonViaUv', () => {
2626
let mockLog: LogOutputChannel;
2727
let isUvInstalledStub: sinon.SinonStub;
28+
let showErrorMessageStub: sinon.SinonStub;
2829
let showInformationMessageStub: sinon.SinonStub;
2930
let sendTelemetryEventStub: sinon.SinonStub;
3031
let mockState: { get: sinon.SinonStub; set: sinon.SinonStub; clear: sinon.SinonStub };
@@ -39,6 +40,7 @@ suite('uvPythonInstaller - promptInstallPythonViaUv', () => {
3940
};
4041
sinon.stub(persistentState, 'getGlobalPersistentState').resolves(mockState);
4142
isUvInstalledStub = sinon.stub(helpers, 'isUvInstalled');
43+
showErrorMessageStub = sinon.stub(windowApis, 'showErrorMessage');
4244
showInformationMessageStub = sinon.stub(windowApis, 'showInformationMessage');
4345
sendTelemetryEventStub = sinon.stub(telemetrySender, 'sendTelemetryEvent');
4446
});
@@ -47,6 +49,25 @@ suite('uvPythonInstaller - promptInstallPythonViaUv', () => {
4749
sinon.restore();
4850
});
4951

52+
function stubUvInstallTask(exitCode: number | undefined): sinon.SinonStub {
53+
let taskEndListener: ((event: TaskProcessEndEvent) => unknown) | undefined;
54+
sinon.stub(taskApis, 'onDidEndTaskProcess').callsFake((listener) => {
55+
taskEndListener = listener;
56+
return { dispose: () => undefined };
57+
});
58+
const executeTaskStub = sinon.stub(taskApis, 'executeTask').callsFake(async (task) => {
59+
const execution = { task, terminate: () => undefined } as TaskExecution;
60+
setImmediate(() => taskEndListener?.({ execution, exitCode } as TaskProcessEndEvent));
61+
return execution;
62+
});
63+
64+
const commandCheck = new MockChildProcess('curl', ['--version']);
65+
const spawnStub: sinon.SinonStub = sinon.stub(childProcessApis, 'spawnProcess');
66+
spawnStub.returns(commandCheck);
67+
setImmediate(() => commandCheck.emit('exit', 0, null));
68+
return executeTaskStub;
69+
}
70+
5071
test('should return undefined when "Don\'t ask again" is set', async () => {
5172
mockState.get.resolves(true);
5273

@@ -219,6 +240,50 @@ suite('uvPythonInstaller - promptInstallPythonViaUv', () => {
219240
);
220241
});
221242

243+
test('should install uv for version lookup after consent', async () => {
244+
isUvInstalledStub.onFirstCall().resolves(false);
245+
isUvInstalledStub.onSecondCall().resolves(true);
246+
showInformationMessageStub.resolves(UvInstallStrings.installUv);
247+
const executeTaskStub = stubUvInstallTask(0);
248+
249+
assert.strictEqual(await ensureUvForInlineScriptVersionLookup('>=3.13,<3.14', mockLog), true);
250+
assert.strictEqual(isUvInstalledStub.callCount, 2);
251+
assert.strictEqual(executeTaskStub.callCount, 1);
252+
assert.strictEqual(showErrorMessageStub.callCount, 0);
253+
});
254+
255+
test('should stop version lookup when uv installation fails', async () => {
256+
isUvInstalledStub.resolves(false);
257+
showInformationMessageStub.resolves(UvInstallStrings.installUv);
258+
const executeTaskStub = stubUvInstallTask(1);
259+
260+
assert.strictEqual(await ensureUvForInlineScriptVersionLookup('>=3.13,<3.14', mockLog), false);
261+
assert.strictEqual(isUvInstalledStub.callCount, 1);
262+
assert.strictEqual(executeTaskStub.callCount, 1);
263+
assert.strictEqual(showErrorMessageStub.callCount, 0);
264+
});
265+
266+
test('should show restart guidance when installed uv remains unavailable', async () => {
267+
isUvInstalledStub.resolves(false);
268+
showInformationMessageStub.resolves(UvInstallStrings.installUv);
269+
stubUvInstallTask(0);
270+
271+
assert.strictEqual(await ensureUvForInlineScriptVersionLookup('>=3.13,<3.14', mockLog), false);
272+
assert.strictEqual(isUvInstalledStub.callCount, 2);
273+
sinon.assert.calledOnceWithExactly(showErrorMessageStub, UvInstallStrings.uvInstallRestartRequired);
274+
});
275+
276+
test('should stop version lookup when uv installation is cancelled', async () => {
277+
isUvInstalledStub.resolves(false);
278+
showInformationMessageStub.resolves(UvInstallStrings.installUv);
279+
const executeTaskStub = stubUvInstallTask(undefined);
280+
281+
assert.strictEqual(await ensureUvForInlineScriptVersionLookup('>=3.13,<3.14', mockLog), false);
282+
assert.strictEqual(isUvInstalledStub.callCount, 1);
283+
assert.strictEqual(executeTaskStub.callCount, 1);
284+
assert.strictEqual(showErrorMessageStub.callCount, 0);
285+
});
286+
222287
test('should trim inline-script context before displaying it', async () => {
223288
mockState.get.resolves(false);
224289
isUvInstalledStub.resolves(true);
@@ -401,10 +466,6 @@ suite('uvPythonInstaller - isDontAskAgainSet and clearDontAskAgain', () => {
401466
});
402467
});
403468

404-
// NOTE: Installation functions (installUv, installPythonViaUv, installPythonWithUv) require
405-
// VS Code's Task API which cannot be fully mocked in unit tests.
406-
// These should be tested via integration tests in a real VS Code environment.
407-
408469
/**
409470
* Helper to build a UvPythonVersion object for testing.
410471
*/

0 commit comments

Comments
 (0)