From 2cd5f924a74b5c7bde203bf5fc4a6fa0001cf78a Mon Sep 17 00:00:00 2001 From: PawelJanowski Date: Tue, 19 May 2026 08:57:54 -0700 Subject: [PATCH] Fix: dispatch interactive .kql queries to currently selected controller When a query is launched from a .kql file via kusto.executeSelectedQuery, the controller used to execute it is currently the one cached at the time the interactive window was first opened. Changing the kernel via the interactive window's picker updates the UI and globalState but does not affect subsequent runs from the source .kql -- they keep routing through the cached (now stale) controller, AND the per-document Kusto client cache in Client.create also returns a stale client built for the original connection, even if the caller now passes a different connection. Two-part fix: * Track the currently-selected controller per notebook via the existing onDidChangeSelectedNotebooks signal, expose getSelectedController(), and prefer it over the cached controller in executeSelectedQuery. * In Client.create, when the caller passes a connectionInfo that doesn't match the cached client's connectionInfo, drop the stale entry and recreate. Preserves caching when nothing changed. --- src/extension/interactive/interactive.ts | 11 ++++++++--- src/extension/kernel/provider.ts | 16 ++++++++++++++++ src/extension/kusto/client.ts | 13 ++++++++++++- 3 files changed, 36 insertions(+), 4 deletions(-) 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