feat: enhance Python installation prompts and responses via uv

karthiknadig · karthiknadig · commit 639c1cc89a9e · 2026-02-06T12:09:35.000-08:00
diff --git a/src/common/localize.ts b/src/common/localize.ts
@@ -221,6 +221,12 @@ export namespace UvInstallStrings {
     export const installingUv = l10n.t('Installing uv...');
     export const installingPython = l10n.t('Installing Python via uv...');
     export const installComplete = l10n.t('Python installed successfully');
+    export function installCompleteWithDetails(version: string, path: string): string {
+        return l10n.t('Python {0} installed successfully at {1}', version, path);
+    }
+    export function installCompleteWithPath(path: string): string {
+        return l10n.t('Python installed successfully at {0}', path);
+    }
     export const installFailed = l10n.t('Failed to install Python');
     export const uvInstallFailed = l10n.t('Failed to install uv');
     export const dontAskAgain = l10n.t("Don't ask again");
diff --git a/src/common/telemetry/constants.ts b/src/common/telemetry/constants.ts
@@ -128,7 +128,7 @@ export interface IEventNamePropertyMapping {
         }
     */
     [EventNames.UV_PYTHON_INSTALL_FAILED]: {
-        stage: 'uvInstall' | 'pythonInstall';
+        stage: 'uvInstall' | 'pythonInstall' | 'findPath';
     };
 
     /* __GDPR__
diff --git a/src/managers/builtin/sysPythonManager.ts b/src/managers/builtin/sysPythonManager.ts
@@ -75,7 +75,19 @@ export class SysPythonManager implements EnvironmentManager {
 
         // If no Python environments were found, offer to install via uv
         if (this.collection.length === 0) {
-            await promptInstallPythonViaUv('activation', this.api, this.log);
+            const pythonPath = await promptInstallPythonViaUv('activation', this.log);
+            if (pythonPath) {
+                const resolved = await resolveSystemPythonEnvironmentPath(
+                    pythonPath,
+                    this.nativeFinder,
+                    this.api,
+                    this,
+                );
+                if (resolved) {
+                    this.collection.push(resolved);
+                    this._onDidChangeEnvironments.fire([{ environment: resolved, kind: EnvironmentChangeKind.add }]);
+                }
+            }
         }
 
         this._initialized.resolve();
@@ -234,12 +246,17 @@ export class SysPythonManager implements EnvironmentManager {
         _scope: CreateEnvironmentScope,
         _options?: CreateEnvironmentOptions,
     ): Promise<PythonEnvironment | undefined> {
-        const success = await installPythonWithUv(this.api, this.log);
+        const pythonPath = await installPythonWithUv(this.log);
 
-        if (success) {
-            // Return the latest Python environment after installation
-            // The installPythonWithUv function already refreshes environments
-            return getLatest(this.collection);
+        if (pythonPath) {
+            // Resolve the installed Python using NativePythonFinder instead of full refresh
+            const resolved = await resolveSystemPythonEnvironmentPath(pythonPath, this.nativeFinder, this.api, this);
+            if (resolved) {
+                // Add to collection and fire change event
+                this.collection.push(resolved);
+                this._onDidChangeEnvironments.fire([{ environment: resolved, kind: EnvironmentChangeKind.add }]);
+                return resolved;
+            }
         }
 
         return undefined;
diff --git a/src/managers/builtin/uvPythonInstaller.ts b/src/managers/builtin/uvPythonInstaller.ts
@@ -7,7 +7,6 @@ import {
     TaskRevealKind,
     TaskScope,
 } from 'vscode';
-import { PythonEnvironmentApi } from '../../api';
 import { spawnProcess } from '../../common/childProcess.apis';
 import { UvInstallStrings } from '../../common/localize';
 import { traceError, traceInfo, traceLog } from '../../common/logging';
@@ -141,6 +140,27 @@ export async function installUv(_log?: LogOutputChannel): Promise<boolean> {
     return success;
 }
 
+/**
+ * Gets the path to the uv-managed Python installation.
+ * @returns Promise that resolves to the Python path, or undefined if not found
+ */
+export async function getUvPythonPath(): Promise<string | undefined> {
+    return new Promise((resolve) => {
+        const chunks: string[] = [];
+        const proc = spawnProcess('uv', ['python', 'find']);
+        proc.stdout?.on('data', (data) => chunks.push(data.toString()));
+        proc.on('error', () => resolve(undefined));
+        proc.on('exit', (code) => {
+            if (code === 0 && chunks.length > 0) {
+                const pythonPath = chunks.join('').trim();
+                resolve(pythonPath || undefined);
+            } else {
+                resolve(undefined);
+            }
+        });
+    });
+}
+
 /**
  * Installs Python using uv.
  * @param log Optional log output channel
@@ -171,21 +191,19 @@ export async function installPythonViaUv(_log?: LogOutputChannel, version?: stri
  * Respects the "Don't ask again" setting.
  *
  * @param trigger What triggered this prompt ('activation' or 'createEnvironment')
- * @param api The Python environment API (used to refresh environments after installation)
  * @param log Optional log output channel
- * @returns Promise that resolves to true if Python was successfully installed
+ * @returns Promise that resolves to the installed Python path, or undefined if not installed
  */
 export async function promptInstallPythonViaUv(
     trigger: 'activation' | 'createEnvironment',
-    api: PythonEnvironmentApi,
     log?: LogOutputChannel,
-): Promise<boolean> {
+): Promise<string | undefined> {
     const state = await getGlobalPersistentState();
     const dontAsk = await state.get<boolean>(UV_INSTALL_PYTHON_DONT_ASK_KEY);
 
     if (dontAsk) {
         traceLog('Skipping Python install prompt: user selected "Don\'t ask again"');
-        return false;
+        return undefined;
     }
 
     sendTelemetryEvent(EventNames.UV_PYTHON_INSTALL_PROMPTED, undefined, { trigger });
@@ -205,25 +223,24 @@ export async function promptInstallPythonViaUv(
     if (result === UvInstallStrings.dontAskAgain) {
         await state.set(UV_INSTALL_PYTHON_DONT_ASK_KEY, true);
         traceLog('User selected "Don\'t ask again" for Python install prompt');
-        return false;
+        return undefined;
     }
 
     if (result === UvInstallStrings.installPython) {
-        return await installPythonWithUv(api, log);
+        return await installPythonWithUv(log);
     }
 
-    return false;
+    return undefined;
 }
 
 /**
  * Installs Python using uv. If uv is not installed, installs it first.
  * This is the main entry point for programmatic Python installation.
  *
- * @param api The Python environment API (used to refresh environments after installation)
  * @param log Optional log output channel
- * @returns Promise that resolves to true if Python was successfully installed
+ * @returns Promise that resolves to the installed Python path, or undefined on failure
  */
-export async function installPythonWithUv(api: PythonEnvironmentApi, log?: LogOutputChannel): Promise<boolean> {
+export async function installPythonWithUv(log?: LogOutputChannel): Promise<string | undefined> {
     const uvInstalled = await isUvInstalled(log);
 
     sendTelemetryEvent(EventNames.UV_PYTHON_INSTALL_STARTED, undefined, { uvAlreadyInstalled: uvInstalled });
@@ -243,7 +260,7 @@ export async function installPythonWithUv(api: PythonEnvironmentApi, log?: LogOu
                 if (!uvSuccess) {
                     sendTelemetryEvent(EventNames.UV_PYTHON_INSTALL_FAILED, undefined, { stage: 'uvInstall' });
                     showErrorMessage(UvInstallStrings.uvInstallFailed);
-                    return false;
+                    return undefined;
                 }
             }
 
@@ -252,17 +269,23 @@ export async function installPythonWithUv(api: PythonEnvironmentApi, log?: LogOu
             if (!pythonSuccess) {
                 sendTelemetryEvent(EventNames.UV_PYTHON_INSTALL_FAILED, undefined, { stage: 'pythonInstall' });
                 showErrorMessage(UvInstallStrings.installFailed);
-                return false;
+                return undefined;
             }
 
-            // Step 3: Refresh environments to detect newly installed Python
-            traceInfo('Refreshing environments after Python installation...');
-            await api.refreshEnvironments(undefined);
+            // Step 3: Get the installed Python path using uv python find
+            const pythonPath = await getUvPythonPath();
+            if (!pythonPath) {
+                traceError('Python installed but could not find the path via uv python find');
+                sendTelemetryEvent(EventNames.UV_PYTHON_INSTALL_FAILED, undefined, { stage: 'findPath' });
+                showErrorMessage(UvInstallStrings.installFailed);
+                return undefined;
+            }
 
+            traceInfo(`Python installed successfully at: ${pythonPath}`);
             sendTelemetryEvent(EventNames.UV_PYTHON_INSTALL_COMPLETED);
-            showInformationMessage(UvInstallStrings.installComplete);
+            showInformationMessage(UvInstallStrings.installCompleteWithPath(pythonPath));
 
-            return true;
+            return pythonPath;
         },
     );
 }
diff --git a/src/managers/builtin/venvManager.ts b/src/managers/builtin/venvManager.ts
@@ -147,9 +147,11 @@ export class VenvManager implements EnvironmentManager {
 
             // If no Python environments found, offer to install Python via uv
             if (globals.length === 0) {
-                const installed = await promptInstallPythonViaUv('createEnvironment', this.api, this.log);
-                if (installed) {
-                    // Re-fetch environments after installation
+                const installedPath = await promptInstallPythonViaUv('createEnvironment', this.log);
+                if (installedPath) {
+                    // Refresh environments to detect the newly installed Python
+                    await this.api.refreshEnvironments(undefined);
+                    // Re-fetch environments after refresh
                     globals = await this.api.getEnvironments('global');
                     // Update globalEnv reference if we found any Python 3.x environments
                     const python3Envs = globals.filter((e) => e.version.startsWith('3.'));
diff --git a/src/test/managers/builtin/uvPythonInstaller.unit.test.ts b/src/test/managers/builtin/uvPythonInstaller.unit.test.ts
@@ -1,7 +1,6 @@
 import assert from 'assert';
 import * as sinon from 'sinon';
 import { LogOutputChannel } from 'vscode';
-import { PythonEnvironmentApi } from '../../../api';
 import { UvInstallStrings } from '../../../common/localize';
 import * as persistentState from '../../../common/persistentState';
 import { EventNames } from '../../../common/telemetry/constants';
@@ -17,15 +16,13 @@ import { createMockLogOutputChannel } from '../../mocks/helper';
 
 suite('uvPythonInstaller - promptInstallPythonViaUv', () => {
     let mockLog: LogOutputChannel;
-    let mockApi: Partial<PythonEnvironmentApi>;
     let isUvInstalledStub: sinon.SinonStub;
     let showInformationMessageStub: sinon.SinonStub;
     let sendTelemetryEventStub: sinon.SinonStub;
     let mockState: { get: sinon.SinonStub; set: sinon.SinonStub; clear: sinon.SinonStub };
 
     setup(() => {
         mockLog = createMockLogOutputChannel();
-        mockApi = { refreshEnvironments: sinon.stub().resolves() };
 
         mockState = {
             get: sinon.stub(),
@@ -42,12 +39,12 @@ suite('uvPythonInstaller - promptInstallPythonViaUv', () => {
         sinon.restore();
     });
 
-    test('should return false when "Don\'t ask again" is set', async () => {
+    test('should return undefined when "Don\'t ask again" is set', async () => {
         mockState.get.resolves(true);
 
-        const result = await promptInstallPythonViaUv('activation', mockApi as PythonEnvironmentApi, mockLog);
+        const result = await promptInstallPythonViaUv('activation', mockLog);
 
-        assert.strictEqual(result, false);
+        assert.strictEqual(result, undefined);
         assert(showInformationMessageStub.notCalled, 'Should not show message when dont ask again is set');
         assert(sendTelemetryEventStub.notCalled, 'Should not send telemetry when skipping prompt');
     });
@@ -57,7 +54,7 @@ suite('uvPythonInstaller - promptInstallPythonViaUv', () => {
         isUvInstalledStub.resolves(true);
         showInformationMessageStub.resolves(undefined); // User dismissed
 
-        await promptInstallPythonViaUv('activation', mockApi as PythonEnvironmentApi, mockLog);
+        await promptInstallPythonViaUv('activation', mockLog);
 
         assert(
             showInformationMessageStub.calledWith(
@@ -74,7 +71,7 @@ suite('uvPythonInstaller - promptInstallPythonViaUv', () => {
         isUvInstalledStub.resolves(false);
         showInformationMessageStub.resolves(undefined); // User dismissed
 
-        await promptInstallPythonViaUv('activation', mockApi as PythonEnvironmentApi, mockLog);
+        await promptInstallPythonViaUv('activation', mockLog);
 
         assert(
             showInformationMessageStub.calledWith(
@@ -91,28 +88,28 @@ suite('uvPythonInstaller - promptInstallPythonViaUv', () => {
         isUvInstalledStub.resolves(true);
         showInformationMessageStub.resolves(UvInstallStrings.dontAskAgain);
 
-        const result = await promptInstallPythonViaUv('activation', mockApi as PythonEnvironmentApi, mockLog);
+        const result = await promptInstallPythonViaUv('activation', mockLog);
 
-        assert.strictEqual(result, false);
+        assert.strictEqual(result, undefined);
         assert(mockState.set.calledWith('python-envs:uv:UV_INSTALL_PYTHON_DONT_ASK', true), 'Should set dont ask flag');
     });
 
-    test('should return false when user dismisses the dialog', async () => {
+    test('should return undefined when user dismisses the dialog', async () => {
         mockState.get.resolves(false);
         isUvInstalledStub.resolves(true);
         showInformationMessageStub.resolves(undefined); // User dismissed
 
-        const result = await promptInstallPythonViaUv('activation', mockApi as PythonEnvironmentApi, mockLog);
+        const result = await promptInstallPythonViaUv('activation', mockLog);
 
-        assert.strictEqual(result, false);
+        assert.strictEqual(result, undefined);
     });
 
     test('should send telemetry with correct trigger', async () => {
         mockState.get.resolves(false);
         isUvInstalledStub.resolves(true);
         showInformationMessageStub.resolves(undefined);
 
-        await promptInstallPythonViaUv('createEnvironment', mockApi as PythonEnvironmentApi, mockLog);
+        await promptInstallPythonViaUv('createEnvironment', mockLog);
 
         assert(
             sendTelemetryEventStub.calledWith(EventNames.UV_PYTHON_INSTALL_PROMPTED, undefined, {

Original file line number	Diff line number	Diff line change
`@@ -128,7 +128,7 @@ export interface IEventNamePropertyMapping {`
`128`	`128`	`}`
`129`	`129`	`*/`
`130`	`130`	`[EventNames.UV_PYTHON_INSTALL_FAILED]: {`
`131`		`- stage: 'uvInstall' \| 'pythonInstall';`
	`131`	`+ stage: 'uvInstall' \| 'pythonInstall' \| 'findPath';`
`132`	`132`	`};`
`133`	`133`
`134`	`134`	`/* __GDPR__`