Attempt to specify a venv for jupyter

source/vscode/src/learning/commands.ts

@@ -271,14 +271,47 @@ async function openCourseNotebook(
     log.warn("No notebook associated with the current position.");
     return;
   }
-  const cellId = service.getCurrentExerciseCellId();
 
-  await vscode.commands.executeCommand(
-    "vscode.openWith",
-    notebookUri,
-    "jupyter-notebook",
-    { viewColumn: vscode.ViewColumn.Active, preview: false },
-  );
+  let opened = false;
+
+  // Try to open via the Jupyter extension's unstable API so the course's
+  // Python environment is automatically set as the active kernel.
+  try {
+    // Grab before awaiting, in case it changes while we're yielding
+    // and gets out of sync with notebookUri
+    const courseId = service.getActiveCourseId();
+    const jupyter = vscode.extensions.getExtension("ms-toolsai.jupyter");
+    const api = await jupyter?.activate();
+    if (api && typeof api.openNotebook === "function") {
+      const envPath = await service.getJupyterEnvironmentPath(courseId);
+      if (envPath) {
+        await api.openNotebook(notebookUri, envPath);
+        opened = true;
+      } else {
+        log.info("Didn't find a course virtual environment to use in notebook");
+      }
+    }
+    if (!opened) {
+      log.warn(
+        "Jupyter openNotebook API is not available; falling back to generic open.",
+      );
+    }
+  } catch (e) {
+    log.warn(
+      `Jupyter openNotebook API call failed: ${e}; falling back to generic open.`,
+    );
+  }
+
+  if (!opened) {
+    await vscode.commands.executeCommand(
+      "vscode.openWith",
+      notebookUri,
+      "jupyter-notebook",
+      { viewColumn: vscode.ViewColumn.Active, preview: false },
+    );
+  }
+
+  const cellId = service.getCurrentExerciseCellId();
 
   if (options?.reveal === "top") {
     revealNotebookTop(notebookUri);

source/vscode/src/learning/service.ts

@@ -5,6 +5,7 @@ import { log } from "qsharp-lang";
 import { getExerciseSources } from "qsharp-lang/katas-md";
 import * as vscode from "vscode";
 import { FullProgramConfig, getProgramForDocument } from "../programConfig.js";
+import { getVenvInFolder } from "../pythonEnvs.js";
 import { ProgramRunStatus, runProgram } from "../run.js";
 import { EventType, sendTelemetryEvent } from "../telemetry.js";
 import { createCourseProvider, toDescriptor } from "./courseProvider.js";
@@ -127,7 +128,9 @@ export class LearningService {
   readonly onDidChangeProgress = this._onDidChangeProgress.event;
 
   private _progressFileWatcher: vscode.FileSystemWatcher | undefined;
+  private _jupyterEnvSubscription: vscode.Disposable | undefined;
   private _writingProgress = false;
+  private _writeQueue: Promise<void> = Promise.resolve();
   private _initPromise: Promise<boolean> | undefined;
   /** Whether {@link _initPromise} was started with `createIfMissing`. */
   private _initCreates = false;
@@ -163,6 +166,22 @@ export class LearningService {
     return this.requireWorkspace().workspaceRoot;
   }
 
+  async getJupyterEnvironmentPath(
+    courseId: string,
+  ): Promise<{ id: string; path: string } | undefined> {
+    const ws = this.requireWorkspace();
+    const saved = ws.progressData.pythonEnvironments[courseId];
+    if (saved) {
+      if (await uriExists(vscode.Uri.file(saved.path))) {
+        return saved;
+      }
+      log.info(
+        `Persisted venv at ${saved.path} no longer exists; falling back to discovery.`,
+      );
+    }
+    return await getVenvInFolder(ws.workspaceRoot);
+  }
+
   /**
    * Try to initialize the service. Returns `true` when ready, `false`
    * when no learning workspace could be found (or created).
@@ -255,6 +274,7 @@ export class LearningService {
     this._onDidChangeState.dispose();
     this._onDidChangeProgress.dispose();
     this._progressFileWatcher?.dispose();
+    this._jupyterEnvSubscription?.dispose();
     for (const d of this._disposables) {
       d.dispose();
     }
@@ -1617,6 +1637,7 @@ export class LearningService {
     }
 
     ws.progressData = parsed as ProgressFileData;
+    ws.progressData.pythonEnvironments ??= {};
 
     // Validate saved position references a known course, unit, and activity.
     // Accept the passed-in default if they're not valid.
@@ -1672,31 +1693,44 @@ export class LearningService {
       },
       completions: {},
       startedAt: new Date().toISOString(),
+      pythonEnvironments: {},
     };
   }
 
   private async saveProgress(): Promise<void> {
-    const ws = this.requireWorkspace();
-    const json = JSON.stringify(ws.progressData, null, 2);
-    this._writingProgress = true;
-    try {
-      await vscode.workspace.fs.writeFile(
-        ws.learningFile,
-        new TextEncoder().encode(json),
-      );
-    } catch (err) {
-      throw new Error(
-        `Failed to save learning progress: ${
-          err instanceof Error ? err.message : String(err)
-        }`,
-        { cause: err },
-      );
-    } finally {
-      this._writingProgress = false;
-    }
+    await this.writeProgressFile();
     this.emitProgress();
   }
 
+  private async writeProgressFile(): Promise<void> {
+    // Chain writes so only one is in flight at a time. Each call
+    // re-serializes progressData when its turn comes, so the last
+    // writer always captures every preceding in-memory mutation.
+    const prev = this._writeQueue;
+    const task = (async () => {
+      await prev;
+      const ws = this.requireWorkspace();
+      const json = JSON.stringify(ws.progressData, null, 2);
+      this._writingProgress = true;
+      try {
+        await vscode.workspace.fs.writeFile(
+          ws.learningFile,
+          new TextEncoder().encode(json),
+        );
+      } catch (err) {
+        log.error(
+          `Failed to save learning progress: ${
+            err instanceof Error ? err.message : String(err)
+          }`,
+        );
+      } finally {
+        this._writingProgress = false;
+      }
+    })();
+    this._writeQueue = task;
+    await task;
+  }
+
   async reloadProgress(): Promise<void> {
     const ws = this.requireWorkspace();
     await this.loadProgress(ws);
@@ -1774,6 +1808,60 @@ export class LearningService {
     this._progressFileWatcher.onDidDelete(onDelete);
 
     this.emitProgress();
+    void this.startJupyterEnvWatcher();
+  }
+
+  private async startJupyterEnvWatcher(): Promise<void> {
+    if (this._jupyterEnvSubscription) {
+      return;
+    }
+    try {
+      const jupyter = vscode.extensions.getExtension("ms-toolsai.jupyter");
+      const api = await jupyter?.activate();
+      if (!api?.onDidChangePythonEnvironment || !api?.getPythonEnvironment) {
+        return;
+      }
+      this._jupyterEnvSubscription = api.onDidChangePythonEnvironment(
+        (uri: vscode.Uri) => {
+          void this.onJupyterEnvChanged(api, uri);
+        },
+      );
+    } catch {
+      log.info("Jupyter extension not available; skipping kernel watch.");
+    }
+  }
+
+  private async onJupyterEnvChanged(
+    jupyterApi: {
+      getPythonEnvironment(
+        uri: vscode.Uri,
+      ): { id: string; path: string } | undefined;
+    },
+    uri: vscode.Uri,
+  ): Promise<void> {
+    if (!this.workspace) {
+      return;
+    }
+    const resolved = this.resolveWorkbookLocation(uri);
+    if (!resolved) {
+      return;
+    }
+    const env = jupyterApi.getPythonEnvironment(uri);
+    if (!env) {
+      return;
+    }
+    const { id, path } = env;
+    // The Jupyter API is undocumented; guard against shape changes.
+    if (typeof id !== "string" || typeof path !== "string" || !id || !path) {
+      return;
+    }
+    const { course } = resolved;
+    const existing = this.workspace.progressData.pythonEnvironments[course.id];
+    if (existing?.id === id && existing?.path === path) {
+      return;
+    }
+    this.workspace.progressData.pythonEnvironments[course.id] = { id, path };
+    await this.writeProgressFile();
   }
 
   private emitProgress(): void {

source/vscode/src/learning/types.d.ts

@@ -135,6 +135,7 @@ export interface ProgressFileData {
   position: ActivityLocation;
   completions: Record<string, { completedAt: string }>;
   startedAt: string;
+  pythonEnvironments: Record<string, { id: string; path: string }>;
 }
 
 // ─── Bundled state ───

source/vscode/src/pythonEnvs.ts

@@ -94,45 +94,45 @@ async function getPythonEnvsApi(): Promise<PythonEnvironmentApi | undefined> {
   }
 }
 
-function isInWorkspaceRoot(env: PythonEnvironment, root: vscode.Uri): boolean {
+function isEnvInFolder(env: PythonEnvironment, folder: vscode.Uri): boolean {
   const envStr = vscode.Uri.file(env.sysPrefix).toString();
-  const rootStr = root.toString();
-  return envStr.startsWith(rootStr + "/");
+  const folderStr = folder.toString();
+  return envStr.startsWith(folderStr + "/");
 }
 
-// Look for an environment in the workspace root.
-// Prefer the active environment if it's in the root.
-async function getWorkspaceRootEnv(
+// Find a venv whose sysPrefix is in the given folder.
+// Prefer the active environment if it qualifies.
+async function getEnvInFolder(
   api: PythonEnvironmentApi,
-  root: vscode.Uri,
+  folder: vscode.Uri,
 ): Promise<PythonEnvironment | undefined> {
-  log.trace(`Searching for existing venvs in ${root.fsPath}`);
+  log.trace(`Searching for existing venvs in ${folder.fsPath}`);
 
-  await api.refreshEnvironments(root);
+  await api.refreshEnvironments(folder);
 
   // This can return the global environment, for example, so we have to check
-  // whether it's actually
-  const activeEnv = await api.getEnvironment(root);
-  if (activeEnv && isInWorkspaceRoot(activeEnv, root)) {
+  // whether it's actually in the folder.
+  const activeEnv = await api.getEnvironment(folder);
+  if (activeEnv && isEnvInFolder(activeEnv, folder)) {
     log.trace(`Preferring active venv in ${activeEnv.environmentPath.fsPath}`);
     return activeEnv;
   }
 
-  const allEnvs = await api.getEnvironments(root);
-  const matchingEnvs = allEnvs.filter((env) => isInWorkspaceRoot(env, root));
+  const allEnvs = await api.getEnvironments(folder);
+  const matchingEnvs = allEnvs.filter((env) => isEnvInFolder(env, folder));
   if (matchingEnvs.length == 0) {
-    log.trace(`Found no venvs in ${root.fsPath}`);
+    log.trace(`Found no venvs in ${folder.fsPath}`);
     return undefined;
   }
 
   if (matchingEnvs.length > 1) {
     log.warn(
-      `Found multiple venvs in ${root.fsPath} - using ${matchingEnvs[0].environmentPath}`,
+      `Found multiple venvs in ${folder.fsPath} - using ${matchingEnvs[0].environmentPath}`,
     );
   }
 
   log.trace(
-    `Found existing venv in ${root.fsPath} - using ${matchingEnvs[0].environmentPath}`,
+    `Found existing venv in ${folder.fsPath} - using ${matchingEnvs[0].environmentPath}`,
   );
   return matchingEnvs[0];
 }
@@ -157,19 +157,29 @@ function getActiveWorkspaceRoot(): vscode.Uri | undefined {
   return result;
 }
 
-export async function getExistingQuantumVenv(): Promise<
-  vscode.Uri | undefined
-> {
+export async function getVenvInFolder(
+  folder: vscode.Uri,
+): Promise<{ id: string; path: string } | undefined> {
   const api = await getPythonEnvsApi();
   if (!api) {
     return undefined;
   }
+  const env = await getEnvInFolder(api, folder);
+  if (!env || !isEnvInFolder(env, folder)) {
+    return undefined;
+  }
+  return { id: env.envId.id, path: env.environmentPath.fsPath };
+}
+
+export async function getExistingQuantumVenv(): Promise<
+  vscode.Uri | undefined
+> {
   const root = getActiveWorkspaceRoot();
   if (!root) {
     return undefined;
   }
-  const env = await getWorkspaceRootEnv(api, root);
-  return env?.environmentPath;
+  const info = await getVenvInFolder(root);
+  return info ? vscode.Uri.file(info.path) : undefined;
 }
 
 export async function createQuantumVenv(): Promise<{ action: string }> {
@@ -192,7 +202,7 @@ export async function createQuantumVenv(): Promise<{ action: string }> {
     selectedPackages.map((item) => item.label),
   );
 
-  const existingEnv = await getWorkspaceRootEnv(api, root);
+  const existingEnv = await getEnvInFolder(api, root);
   if (existingEnv) {
     // Copilot should already have confirmed that the user is willing to update
     // the existing workspace
@@ -254,7 +264,7 @@ export async function createQuantumVenvForCommand(): Promise<void> {
     root = picked.uri;
   }
 
-  const existingEnv = await getWorkspaceRootEnv(api, root);
+  const existingEnv = await getEnvInFolder(api, root);
   if (existingEnv) {
     const choice = await vscode.window.showQuickPick(
       ["Update existing environment", "Cancel"],