From 153e8e72f9e296963559e8651f799e4ee699dbb6 Mon Sep 17 00:00:00 2001
From: Don Jayamanne <don.jayamanne@outlook.com>
Date: Wed, 21 May 2025 06:45:17 +1000
Subject: [PATCH] Handle cases when failing to resolve env

---
 src/client/chat/configurePythonEnvTool.ts | 99 +++++++++++++----------
 src/client/chat/getExecutableTool.ts      | 10 ++-
 src/client/chat/getPythonEnvTool.ts       |  9 ++-
 src/client/chat/installPackagesTool.ts    | 12 ++-
 src/client/chat/utils.ts                  | 30 ++++++-
 5 files changed, 110 insertions(+), 50 deletions(-)

diff --git a/src/client/chat/configurePythonEnvTool.ts b/src/client/chat/configurePythonEnvTool.ts
index 1a22fed83140..aef2272f357c 100644
--- a/src/client/chat/configurePythonEnvTool.ts
+++ b/src/client/chat/configurePythonEnvTool.ts
@@ -19,7 +19,7 @@ import { PythonExtension, ResolvedEnvironment } from '../api/types';
 import { IServiceContainer } from '../ioc/types';
 import { ICodeExecutionService } from '../terminals/types';
 import { TerminalCodeExecutionProvider } from '../terminals/codeExecution/terminalCodeExecution';
-import { getEnvironmentDetails, raceCancellationError } from './utils';
+import { getEnvironmentDetails, NoEnvironmentError, raceCancellationError } from './utils';
 import { resolveFilePath } from './utils';
 import { IRecommendedEnvironmentService } from '../interpreter/configuration/types';
 import { ITerminalHelper } from '../common/terminal/types';
@@ -67,33 +67,51 @@ export class ConfigurePythonEnvTool implements LanguageModelTool<IResourceRefere
         options: LanguageModelToolInvocationOptions<IResourceReference>,
         token: CancellationToken,
     ): Promise<LanguageModelToolResult> {
-        const resource = resolveFilePath(options.input.resourcePath);
-        const recommededEnv = await this.recommendedEnvService.getRecommededEnvironment(resource);
-        // Already selected workspace env, hence nothing to do.
-        if (recommededEnv?.reason === 'workspaceUserSelected' && workspace.workspaceFolders?.length) {
-            return await getEnvDetailsForResponse(
-                recommededEnv.environment,
-                this.api,
-                this.terminalExecutionService,
-                this.terminalHelper,
-                resource,
-                token,
-            );
-        }
-        // No workspace folders, and the user selected a global environment.
-        if (recommededEnv?.reason === 'globalUserSelected' && !workspace.workspaceFolders?.length) {
-            return await getEnvDetailsForResponse(
-                recommededEnv.environment,
-                this.api,
-                this.terminalExecutionService,
-                this.terminalHelper,
-                resource,
-                token,
-            );
-        }
+        try {
+            const resource = resolveFilePath(options.input.resourcePath);
+            const recommededEnv = await this.recommendedEnvService.getRecommededEnvironment(resource);
+            // Already selected workspace env, hence nothing to do.
+            if (recommededEnv?.reason === 'workspaceUserSelected' && workspace.workspaceFolders?.length) {
+                return await getEnvDetailsForResponse(
+                    recommededEnv.environment,
+                    this.api,
+                    this.terminalExecutionService,
+                    this.terminalHelper,
+                    resource,
+                    token,
+                );
+            }
+            // No workspace folders, and the user selected a global environment.
+            if (recommededEnv?.reason === 'globalUserSelected' && !workspace.workspaceFolders?.length) {
+                return await getEnvDetailsForResponse(
+                    recommededEnv.environment,
+                    this.api,
+                    this.terminalExecutionService,
+                    this.terminalHelper,
+                    resource,
+                    token,
+                );
+            }
 
-        if (!workspace.workspaceFolders?.length) {
-            const selected = await Promise.resolve(commands.executeCommand(Commands.Set_Interpreter));
+            if (!workspace.workspaceFolders?.length) {
+                const selected = await Promise.resolve(commands.executeCommand(Commands.Set_Interpreter));
+                const env = await this.api.resolveEnvironment(this.api.getActiveEnvironmentPath(resource));
+                if (selected && env) {
+                    return await getEnvDetailsForResponse(
+                        env,
+                        this.api,
+                        this.terminalExecutionService,
+                        this.terminalHelper,
+                        resource,
+                        token,
+                    );
+                }
+                return new LanguageModelToolResult([
+                    new LanguageModelTextPart('User did not select a Python environment.'),
+                ]);
+            }
+
+            const selected = await showCreateAndSelectEnvironmentQuickPick(resource, this.serviceContainer);
             const env = await this.api.resolveEnvironment(this.api.getActiveEnvironmentPath(resource));
             if (selected && env) {
                 return await getEnvDetailsForResponse(
@@ -106,25 +124,18 @@ export class ConfigurePythonEnvTool implements LanguageModelTool<IResourceRefere
                 );
             }
             return new LanguageModelToolResult([
-                new LanguageModelTextPart('User did not select a Python environment.'),
+                new LanguageModelTextPart('User did not create nor select a Python environment.'),
             ]);
+        } catch (ex) {
+            if (ex instanceof NoEnvironmentError) {
+                return new LanguageModelToolResult([
+                    new LanguageModelTextPart(
+                        'Failed to configure a Python Environment, as the environment could not be found.',
+                    ),
+                ]);
+            }
+            throw ex;
         }
-
-        const selected = await showCreateAndSelectEnvironmentQuickPick(resource, this.serviceContainer);
-        const env = await this.api.resolveEnvironment(this.api.getActiveEnvironmentPath(resource));
-        if (selected && env) {
-            return await getEnvDetailsForResponse(
-                env,
-                this.api,
-                this.terminalExecutionService,
-                this.terminalHelper,
-                resource,
-                token,
-            );
-        }
-        return new LanguageModelToolResult([
-            new LanguageModelTextPart('User did not create nor select a Python environment.'),
-        ]);
     }
 
     async prepareInvocation?(
diff --git a/src/client/chat/getExecutableTool.ts b/src/client/chat/getExecutableTool.ts
index 350598e5ca36..c91177991557 100644
--- a/src/client/chat/getExecutableTool.ts
+++ b/src/client/chat/getExecutableTool.ts
@@ -16,7 +16,7 @@ import { PythonExtension } from '../api/types';
 import { IServiceContainer } from '../ioc/types';
 import { ICodeExecutionService } from '../terminals/types';
 import { TerminalCodeExecutionProvider } from '../terminals/codeExecution/terminalCodeExecution';
-import { getEnvDisplayName, getEnvironmentDetails, raceCancellationError } from './utils';
+import { getEnvDisplayName, getEnvironmentDetails, NoEnvironmentError, raceCancellationError } from './utils';
 import { resolveFilePath } from './utils';
 import { traceError } from '../logging';
 import { ITerminalHelper } from '../common/terminal/types';
@@ -58,6 +58,14 @@ export class GetExecutableTool implements LanguageModelTool<IResourceReference>
             );
             return new LanguageModelToolResult([new LanguageModelTextPart(message)]);
         } catch (error) {
+            if (error instanceof NoEnvironmentError) {
+                return new LanguageModelToolResult([
+                    new LanguageModelTextPart(
+                        'Failed to configure a Python Environment, as the environment could not be found.',
+                    ),
+                ]);
+            }
+
             if (error instanceof CancellationError) {
                 throw error;
             }
diff --git a/src/client/chat/getPythonEnvTool.ts b/src/client/chat/getPythonEnvTool.ts
index 11af29f18be7..8f0a5349a117 100644
--- a/src/client/chat/getPythonEnvTool.ts
+++ b/src/client/chat/getPythonEnvTool.ts
@@ -17,7 +17,7 @@ import { IServiceContainer } from '../ioc/types';
 import { ICodeExecutionService } from '../terminals/types';
 import { TerminalCodeExecutionProvider } from '../terminals/codeExecution/terminalCodeExecution';
 import { IProcessServiceFactory, IPythonExecutionFactory } from '../common/process/types';
-import { getEnvironmentDetails, raceCancellationError } from './utils';
+import { getEnvironmentDetails, NoEnvironmentError, raceCancellationError } from './utils';
 import { resolveFilePath } from './utils';
 import { getPythonPackagesResponse } from './listPackagesTool';
 import { ITerminalHelper } from '../common/terminal/types';
@@ -82,6 +82,13 @@ export class GetEnvironmentInfoTool implements LanguageModelTool<IResourceRefere
 
             return new LanguageModelToolResult([new LanguageModelTextPart(message)]);
         } catch (error) {
+            if (error instanceof NoEnvironmentError) {
+                return new LanguageModelToolResult([
+                    new LanguageModelTextPart(
+                        'Failed to configure a Python Environment, as the environment could not be found.',
+                    ),
+                ]);
+            }
             if (error instanceof CancellationError) {
                 throw error;
             }
diff --git a/src/client/chat/installPackagesTool.ts b/src/client/chat/installPackagesTool.ts
index d0531cb015b2..e2f8500edc3e 100644
--- a/src/client/chat/installPackagesTool.ts
+++ b/src/client/chat/installPackagesTool.ts
@@ -14,7 +14,7 @@ import {
 } from 'vscode';
 import { PythonExtension } from '../api/types';
 import { IServiceContainer } from '../ioc/types';
-import { getEnvDisplayName, raceCancellationError } from './utils';
+import { getEnvDisplayName, NoEnvironmentError, raceCancellationError } from './utils';
 import { resolveFilePath } from './utils';
 import { IModuleInstaller } from '../common/installer/types';
 import { ModuleInstallerType } from '../pythonEnvironments/info';
@@ -51,7 +51,7 @@ export class InstallPackagesTool implements LanguageModelTool<IInstallPackageArg
             const envPath = this.api.getActiveEnvironmentPath(resourcePath);
             const environment = await raceCancellationError(this.api.resolveEnvironment(envPath), token);
             if (!environment || !environment.version) {
-                throw new Error('No environment found for the provided resource path: ' + resourcePath?.fsPath);
+                throw new NoEnvironmentError('No environment found');
             }
             const isConda = (environment.environment?.type || '').toLowerCase() === 'conda';
             const installers = this.serviceContainer.getAll<IModuleInstaller>(IModuleInstaller);
@@ -70,6 +70,14 @@ export class InstallPackagesTool implements LanguageModelTool<IInstallPackageArg
             const resultMessage = `Successfully installed ${packagePlurality}: ${options.input.packageList.join(', ')}`;
             return new LanguageModelToolResult([new LanguageModelTextPart(resultMessage)]);
         } catch (error) {
+            if (error instanceof NoEnvironmentError) {
+                return new LanguageModelToolResult([
+                    new LanguageModelTextPart(
+                        'Failed to configure a Python Environment, as the environment could not be found.',
+                    ),
+                ]);
+            }
+
             if (error instanceof CancellationError) {
                 throw error;
             }
diff --git a/src/client/chat/utils.ts b/src/client/chat/utils.ts
index 2973e748aee4..cccade7edfd5 100644
--- a/src/client/chat/utils.ts
+++ b/src/client/chat/utils.ts
@@ -1,7 +1,14 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import { CancellationError, CancellationToken, Uri, workspace } from 'vscode';
+import {
+    CancellationError,
+    CancellationToken,
+    LanguageModelTextPart,
+    LanguageModelToolResult,
+    Uri,
+    workspace,
+} from 'vscode';
 import { IDiscoveryAPI } from '../pythonEnvironments/base/locator';
 import { PythonExtension, ResolvedEnvironment } from '../api/types';
 import { ITerminalHelper, TerminalShellType } from '../common/terminal/types';
@@ -52,6 +59,12 @@ export function isCondaEnv(env: ResolvedEnvironment) {
     return (env.environment?.type || '').toLowerCase() === 'conda';
 }
 
+export class NoEnvironmentError extends Error {
+    constructor(message: string) {
+        super(message);
+    }
+}
+
 export async function getEnvironmentDetails(
     resourcePath: Uri | undefined,
     api: PythonExtension['environments'],
@@ -64,7 +77,7 @@ export async function getEnvironmentDetails(
     const envPath = api.getActiveEnvironmentPath(resourcePath);
     const environment = await raceCancellationError(api.resolveEnvironment(envPath), token);
     if (!environment || !environment.version) {
-        throw new Error('No environment found for the provided resource path: ' + resourcePath?.fsPath);
+        throw new NoEnvironmentError('No environment found for the provided resource path: ' + resourcePath?.fsPath);
     }
     const runCommand = await raceCancellationError(
         getTerminalCommand(environment, resourcePath, terminalExecutionService, terminalHelper),
@@ -115,3 +128,16 @@ async function getCondaRunCommand(environment: ResolvedEnvironment) {
     }
     return { command: cmd[0], args: cmd.slice(1) };
 }
+
+function throwErrorIfUriIsNotebook(uri: Uri | undefined) {
+    if (!uri) {
+        return;
+    }
+    if (
+        uri.scheme === 'vscode-notebook-cell' ||
+        uri.fsPath.endsWith('.ipynb') ||
+        workspace.notebookDocuments.some((doc) => doc.uri.fsPath === uri.fsPath)
+    ) {
+        throw new Error('This tool does not support Notebooks, use a relevant tool instead.');
+    }
+}