diff --git a/src/extension/interactive/interactive.ts b/src/extension/interactive/interactive.ts index b05ea6d..bdb1293 100644 --- a/src/extension/interactive/interactive.ts +++ b/src/extension/interactive/interactive.ts @@ -15,7 +15,7 @@ import { } from 'vscode'; import { commands } from 'vscode'; import { InteractiveWindowView, registerDisposable } from '../utils'; -import { KernelPerConnection, registerController } from '../kernel/provider'; +import { getSelectedController, KernelPerConnection, registerController } from '../kernel/provider'; import { FoldingRangesProvider } from '../languageServer'; import { ensureDocumentHasConnectionInfo } from '../kusto/connections/notebookConnection'; @@ -69,11 +69,16 @@ async function executeSelectedQuery(document: TextDocument, start: number, end: if (!info) { return; } - const [notebook, controller] = info; + const [notebook, cachedController] = info; // Ensure its visible. await commands.executeCommand('interactive.open', undefined, notebook.uri, undefined); const cell = await createCell(notebook, document, start, end); - await controller.executeInteractive([cell], document); + // Prefer the controller currently selected on the interactive notebook so + // that if the user picked a different kernel in the window we honor it. + // Fall back to the originally-resolved controller in case the selection + // signal hasn't landed yet (e.g. first run before any onDidChangeSelectedNotebooks). + const activeController = getSelectedController(notebook) ?? cachedController; + await activeController.executeInteractive([cell], document); } async function pickConnection(document: TextDocument) { diff --git a/src/extension/kernel/provider.ts b/src/extension/kernel/provider.ts index 302674c..fa716f0 100644 --- a/src/extension/kernel/provider.ts +++ b/src/extension/kernel/provider.ts @@ -67,6 +67,15 @@ function getControllerId(connection: IConnectionInfo, notebookType: string) { return `${notebookType}_${encodeConnectionInfo(connection)}`; } +// Currently-selected controller per notebook. Set by the +// onDidChangeSelectedNotebooks handler below; used by `getSelectedController` +// so callers can route execution to the kernel the user actually picked. +const selectedControllerByNotebook = new WeakMap(); + +export function getSelectedController(notebook: NotebookDocument): KernelPerConnection | undefined { + return selectedControllerByNotebook.get(notebook); +} + export class KernelPerConnection extends Disposable { public readonly notebookController: NotebookController; private readonly disposables: Disposable[] = []; @@ -86,6 +95,13 @@ export class KernelPerConnection extends Disposable { this.notebookController.description = displayInfo.description; this.disposables.push( this.notebookController.onDidChangeSelectedNotebooks(async ({ notebook, selected }) => { + if (selected) { + selectedControllerByNotebook.set(notebook, this); + } else if (selectedControllerByNotebook.get(notebook) === this) { + // Only clear if we still own the slot — handles A→B switches + // where B's `select` may fire before A's `deselect`. + selectedControllerByNotebook.delete(notebook); + } if (!selected) { return; } diff --git a/src/extension/kusto/client.ts b/src/extension/kusto/client.ts index 53f11b4..cdf0684 100644 --- a/src/extension/kusto/client.ts +++ b/src/extension/kusto/client.ts @@ -1,5 +1,6 @@ import type { KustoResponseDataSet } from 'azure-kusto-data'; import { workspace, NotebookDocument, TextDocument } from 'vscode'; +import { isEqual } from 'lodash'; // import { addDocumentConnectionHandler, ensureDocumentHasConnectionInfo } from './connections/notebookConnection'; import { IConnection, IConnectionInfo, IKustoClient } from './connections/types'; import { IDisposable } from '../types'; @@ -33,7 +34,17 @@ export class Client implements IDisposable { ): Promise { const client = clientMap.get(document); if (client) { - return client; + // If the caller specified a connection that doesn't match the cached + // client's connection (e.g. user switched kernel since last execution), + // drop the stale entry and recreate. Otherwise reuse. + if (!connectionInfo) { + return client; + } + const existing = await client; + if (existing && isEqual(existing.connectionInfo, connectionInfo)) { + return client; + } + clientMap.delete(document); } // eslint-disable-next-line no-async-promise-executor