diff --git a/extensions/mssql/src/controllers/connectionManager.ts b/extensions/mssql/src/controllers/connectionManager.ts index 7a67398f1f..023ff94c1c 100644 --- a/extensions/mssql/src/controllers/connectionManager.ts +++ b/extensions/mssql/src/controllers/connectionManager.ts @@ -480,6 +480,45 @@ export default class ConnectionManager { return await this.client.sendRequest(requestType, params); } + /** + * Registers interest in the `connection/complete` notification for a URI and returns a + * promise that resolves with the completion params. The `connection/connect` request only + * acknowledges that a connection attempt started; the actual outcome arrives later via + * this notification. Callers that send ConnectionRequest directly (e.g. notebook cell + * IntelliSense registration) use this to observe the real result — call BEFORE sending + * the request, and pair with {@link cancelConnectionCompleteExpectation} on timeout so + * the pending entry doesn't leak. + */ + public expectConnectionComplete( + ownerUri: string, + ): Promise { + const deferred = new Deferred(); + this._uriToConnectionCompleteParamsMap.set(ownerUri, deferred); + return deferred.promise; + } + + /** + * Removes a pending `connection/complete` expectation registered via + * {@link expectConnectionComplete} (no-op if it already resolved). + * @param expectation When provided, the entry is only removed if it still + * belongs to this expectation — a later expectConnectionComplete call for + * the same URI supersedes the old one, and a stale caller's cleanup must + * not cancel the newer expectation. + */ + public cancelConnectionCompleteExpectation( + ownerUri: string, + expectation?: Promise, + ): void { + const pending = this._uriToConnectionCompleteParamsMap.get(ownerUri); + if (!pending) { + return; + } + if (expectation && pending.promise !== expectation) { + return; + } + this._uriToConnectionCompleteParamsMap.delete(ownerUri); + } + /** * Exposed for testing purposes. */ diff --git a/extensions/mssql/src/notebooks/notebookConnectionManager.ts b/extensions/mssql/src/notebooks/notebookConnectionManager.ts index 189ba8c644..bb407011c2 100644 --- a/extensions/mssql/src/notebooks/notebookConnectionManager.ts +++ b/extensions/mssql/src/notebooks/notebookConnectionManager.ts @@ -9,7 +9,13 @@ import ConnectionManager from "../controllers/connectionManager"; import { ConnectionSharingService } from "../connectionSharing/connectionSharingService"; import SqlToolsServiceClient from "../languageservice/serviceclient"; import { QueryNotificationHandler } from "../controllers/queryNotificationHandler"; -import { ConnectionRequest, ConnectParams } from "../models/contracts/connection"; +import { + ConnectionRequest, + ConnectParams, + ConnectionCompleteParams, + DisconnectRequest, + DisconnectParams, +} from "../models/contracts/connection"; import { generateQueryUri } from "../models/utils"; import * as LocalizedConstants from "../constants/locConstants"; import { sendActionEvent, startActivity } from "../telemetry/telemetry"; @@ -37,6 +43,13 @@ import { * actually uses for query execution. Our connection (under an adhoc URI) is the * authoritative one for SQL execution. */ +/** + * How long to wait for STS to report `connection/complete` for a cell's + * IntelliSense registration before treating the attempt as failed so a + * later trigger can retry it. + */ +const CELL_CONNECT_COMPLETE_TIMEOUT_MS = 30000; + export class NotebookConnectionManager implements vscode.Disposable { private connectionUri: string | undefined; private connectionInfo: IConnectionInfo | undefined; @@ -343,7 +356,9 @@ export class NotebookConnectionManager implements vscode.Disposable { this.connectionUri = undefined; this.connectionInfo = undefined; this.connectionLabel = ""; - this.invalidateCellRegistrations(); + // Release (not just forget) cell registrations: each registered cell + // holds its own STS-side connection that nothing else ever closes. + this.releaseCellRegistrations(); } /** @@ -407,9 +422,18 @@ export class NotebookConnectionManager implements vscode.Disposable { const generation = this.connectionGeneration; + // Prefer the live credentials that STS validated for the notebook's execution + // connection — they carry any refreshed auth token and the actual database — + // over the stored profile, so the cell binds to the same database the + // notebook executes against. + const sourceInfo = + (this.connectionUri + ? this.connectionMgr.getConnectionInfoFromUri(this.connectionUri) + : undefined) ?? this.connectionInfo; + let connectionDetails: ConnectionDetails; try { - connectionDetails = this.connectionMgr.createConnectionDetails(this.connectionInfo); + connectionDetails = this.connectionMgr.createConnectionDetails(sourceInfo); } catch (err: any) { this.log.warn( `[connectCellForIntellisense] createConnectionDetails failed: ${err.message} cell=${cellDocumentUri}`, @@ -425,11 +449,16 @@ export class NotebookConnectionManager implements vscode.Disposable { } const authType = connectionDetails.options?.authenticationType ?? "unknown"; this.log.debug( - `[connectCellForIntellisense] Sending connect request scheme=${cellUriScheme} server=${this.connectionInfo.server} database=${this.connectionInfo.database || "(default)"} auth=${authType} cell=${cellDocumentUri}`, + `[connectCellForIntellisense] Sending connect request scheme=${cellUriScheme} server=${sourceInfo.server} database=${sourceInfo.database || "(default)"} auth=${authType} cell=${cellDocumentUri}`, ); this.registeredCellUris.add(cellDocumentUri); + // The connect request only acknowledges that a connection attempt STARTED; + // the actual outcome arrives via the connection/complete notification. + // Register for it before sending the request so the result isn't missed. + const completePromise = this.connectionMgr.expectConnectionComplete(cellDocumentUri); + try { const params: ConnectParams = { ownerUri: cellDocumentUri, @@ -438,25 +467,67 @@ export class NotebookConnectionManager implements vscode.Disposable { const result = await this.connectionMgr.sendRequest(ConnectionRequest.type, params); + if (result !== true) { + this.connectionMgr.cancelConnectionCompleteExpectation( + cellDocumentUri, + completePromise, + ); + if (generation === this.connectionGeneration) { + this.registeredCellUris.delete(cellDocumentUri); + } + this.log.warn( + `[connectCellForIntellisense] STS did not accept connect request (result=${String(result)}) cell=${cellDocumentUri}`, + ); + return; + } + + const completeParams = await this.waitForCellConnectionComplete( + cellDocumentUri, + completePromise, + ); + if (generation !== this.connectionGeneration) { this.log.debug( `[connectCellForIntellisense] connection changed mid-request (gen ${generation} → ${this.connectionGeneration}); dropping stale registration cell=${cellDocumentUri}`, ); + if (completeParams?.connectionId) { + // The stale connect actually landed in STS after this cell's + // registration was released — disconnect it so it doesn't + // linger untracked (a released cell URI is never retried + // under this generation's bookkeeping). + this.disconnectCellUri(cellDocumentUri); + } return; } - if (result !== true) { + if (!completeParams?.connectionId) { + // Connection failed or timed out — deregister so the next + // IntelliSense trigger (cell execution, connection change, + // cell document open) retries instead of silently leaving the + // cell with keyword-only completions. this.registeredCellUris.delete(cellDocumentUri); + if (!completeParams) { + // Timed out: the STS-side connect may still land later, + // and this URI is no longer tracked for release on + // disconnect — best-effort disconnect so a late-completing + // connection doesn't linger untracked. + this.disconnectCellUri(cellDocumentUri); + } this.log.warn( - `[connectCellForIntellisense] STS connect did not succeed (result=${String(result)}) cell=${cellDocumentUri}`, + `[connectCellForIntellisense] STS connection did not complete ` + + `(${completeParams ? `error=${completeParams.errorMessage ?? completeParams.messages ?? "unknown"}` : "timed out"}) cell=${cellDocumentUri}`, ); return; } this.log.debug( - `[connectCellForIntellisense] STS connect returned true cell=${cellDocumentUri}`, + `[connectCellForIntellisense] STS connection complete connectionId=${completeParams.connectionId} cell=${cellDocumentUri}`, ); } catch (err: any) { + this.connectionMgr.cancelConnectionCompleteExpectation( + cellDocumentUri, + completePromise, + ); if (generation === this.connectionGeneration) { this.registeredCellUris.delete(cellDocumentUri); } @@ -466,6 +537,71 @@ export class NotebookConnectionManager implements vscode.Disposable { } } + /** + * Awaits the connection/complete notification for a cell registration, + * resolving undefined if it doesn't arrive within the timeout. + */ + private async waitForCellConnectionComplete( + cellDocumentUri: string, + completePromise: Promise, + ): Promise { + let timer: ReturnType | undefined; + const timeoutPromise = new Promise((resolve) => { + timer = setTimeout(() => { + this.connectionMgr.cancelConnectionCompleteExpectation( + cellDocumentUri, + completePromise, + ); + resolve(undefined); + }, CELL_CONNECT_COMPLETE_TIMEOUT_MS); + }); + try { + return await Promise.race([completePromise, timeoutPromise]); + } finally { + if (timer !== undefined) { + clearTimeout(timer); + } + } + } + + /** + * Releases all current cell IntelliSense registrations: clears the + * memoization (so future registrations re-connect) and best-effort + * disconnects the previously registered cell URIs from STS. + * + * Called when the notebook's cell URIs change wholesale — i.e. when an + * untitled notebook is saved to disk and every cell document is re-created + * under a file-based URI. This mirrors the query editor's save flow + * (ConnectionManager.transferConnectionToFile), which connects the new URI + * and disconnects the old one. + */ + releaseCellRegistrations(): void { + const staleUris = [...this.registeredCellUris]; + this.invalidateCellRegistrations(); + + for (const uri of staleUris) { + this.disconnectCellUri(uri); + } + } + + /** + * Best-effort disconnect of a single cell's STS-side IntelliSense + * connection (fire-and-forget with logging). + */ + private disconnectCellUri(uri: string): void { + const params: DisconnectParams = { ownerUri: uri }; + void this.connectionMgr.sendRequest(DisconnectRequest.type, params).then( + (disconnected) => + this.log.debug( + `[disconnectCellUri] disconnect result=${String(disconnected)} cell=${uri}`, + ), + (err: any) => + this.log.warn( + `[disconnectCellUri] disconnect failed: ${err?.message ?? err} cell=${uri}`, + ), + ); + } + dispose(): void { this.disconnect(); } diff --git a/extensions/mssql/src/notebooks/sqlNotebookController.ts b/extensions/mssql/src/notebooks/sqlNotebookController.ts index 84dd48efc0..cc350560bf 100644 --- a/extensions/mssql/src/notebooks/sqlNotebookController.ts +++ b/extensions/mssql/src/notebooks/sqlNotebookController.ts @@ -30,6 +30,13 @@ import { TelemetryViews, TelemetryActions, ActivityStatus } from "../sharedInter const NOTEBOOK_RESULT_RENDERER_ID = "ms-mssql.sql-result-renderer"; +/** + * How long a connection manager parked by the close of an untitled notebook + * waits to be adopted by the file-based notebook that replaces it on save, + * before being disposed (which closes its connection). + */ +const SAVE_ADOPTION_TTL_MS = 10000; + const MIME_TEXT_PLAIN = "text/plain"; const MIME_NOTEBOOK_QUERY_RESULT = "application/vnd.mssql.query-result"; type NotebookTextualResultBlock = Exclude; @@ -78,6 +85,30 @@ export class SqlNotebookController implements vscode.Disposable { // SQL kernelspec/language_info metadata so the .ipynb file identifies as SQL // (not Python) when reopened. private readonly selectedNotebooks = new WeakSet(); + /** + * Connection managers parked when a connected untitled notebook closes. + * Saving an untitled notebook REPLACES its NotebookDocument with a new + * file-based document (close + open) rather than updating the URI in + * place, so disposing on close would drop the connection mid-save. The + * file-based notebook that opens with matching cell content adopts the + * parked manager; unclaimed managers are disposed after a short TTL. + * Keyed by the closed notebook's URI string. + */ + private readonly pendingSaveAdoptions = new Map< + string, + { + mgr: NotebookConnectionManager; + signature: string; + timer: ReturnType; + } + >(); + /** + * Notebook URIs recently seen opening, with their open timestamps. + * Used by the park-time adoption scan so a parked connection is only + * offered to notebooks that appeared around the save transition — not to + * long-open unrelated notebooks that happen to have matching content. + */ + private readonly recentNotebookOpens = new Map(); constructor( private connectionMgr: ConnectionManager, @@ -128,12 +159,12 @@ export class SqlNotebookController implements vscode.Disposable { // Auto-detect SQL notebooks and set affinity so VS Code // auto-selects our kernel instead of showing "Detecting Kernels". + // Also adopt any connection manager parked by the close of the + // untitled notebook this document replaced on save. this.disposables.push( - vscode.workspace.onDidOpenNotebookDocument((notebook) => { - if (notebook.notebookType === "jupyter-notebook") { - this.setAffinityIfSql(notebook); - } - }), + vscode.workspace.onDidOpenNotebookDocument((notebook) => + this.handleNotebookOpened(notebook), + ), ); // Set affinity for notebooks already open when the extension activates @@ -174,23 +205,33 @@ export class SqlNotebookController implements vscode.Disposable { this.disposables.push( vscode.workspace.onDidSaveNotebookDocument((notebook) => { // Persist connection metadata under the final file URI (handles untitled → saved file URI change) - this.rekeyConnectionOnSave(notebook); + const uriChanged = this.rekeyConnectionOnSave(notebook); + if (uriChanged) { + // The notebook URI changed (untitled → file), which re-created every + // cell document under a new URI. Transfer the IntelliSense + // registrations: disconnect the stale cell URIs from STS and + // register the new ones — the notebook equivalent of the query + // editor's connection transfer on save. + const mgr = this.connections.get(notebook.uri.toString()); + mgr?.releaseCellRegistrations(); + this.connectCellsForIntellisense(notebook, "didSaveNotebook"); + this.updateStatusBar(notebook); + this.codeLensProvider.refresh(); + } else if (!this.connections.has(notebook.uri.toString())) { + // The save produced a NEW document (untitled → file replacement) + // that this controller has no manager for — claim the connection + // from the untitled notebook it replaced. + this.tryAdoptConnection(notebook); + } this.saveConnectionMetadataIfConnected(notebook); }), ); - // Clean up connection managers when notebooks are closed + // Clean up connection managers when notebooks are closed. this.disposables.push( - vscode.workspace.onDidCloseNotebookDocument((notebook) => { - const key = notebook.uri.toString(); - const mgr = this.connections.get(key); - if (mgr) { - this.log.info(`[onDidCloseNotebookDocument] Disposing manager for ${key}`); - mgr.dispose(); - this.connections.delete(key); - } - // Note: WeakMap entry will be garbage collected automatically - }), + vscode.workspace.onDidCloseNotebookDocument((notebook) => + this.handleNotebookClosed(notebook), + ), ); const messaging = vscode.notebooks.createRendererMessaging(NOTEBOOK_RESULT_RENDERER_ID); @@ -204,6 +245,30 @@ export class SqlNotebookController implements vscode.Disposable { }), ); + // When a cell document opens after the notebook is already connected — + // e.g. cells re-created on save, or re-opened when their language flips + // to SQL after kernel selection — register it with STS for IntelliSense. + // Connect-time registration (connectCellsForIntellisense) can't cover + // these because the cell document didn't exist under that URI yet. + this.disposables.push( + vscode.workspace.onDidOpenTextDocument((doc) => { + if (doc.uri.scheme !== "vscode-notebook-cell" || doc.languageId !== "sql") { + return; + } + const notebook = this.findNotebookForCellDocument(doc); + if (!notebook) { + return; + } + const mgr = this.connections.get(notebook.uri.toString()); + if (!mgr?.isConnected()) { + return; + } + const cellUri = doc.uri.toString(); + this.log.debug(`[onDidOpenTextDocument] Registering opened cell ${cellUri}`); + void mgr.connectCellForIntellisense(cellUri); + }), + ); + // When new cells are added to a notebook, connect them to STS // for IntelliSense if the notebook already has an active connection. this.disposables.push( @@ -469,8 +534,10 @@ export class SqlNotebookController implements vscode.Disposable { * When a notebook is saved (untitled → file), its URI changes but the * connections map still has the entry under the old key. Re-key it to * the new URI by tracking the notebook document object. + * @returns true when the notebook URI changed and a tracked connection + * manager was re-keyed to it. */ - private rekeyConnectionOnSave(notebook: vscode.NotebookDocument): void { + private rekeyConnectionOnSave(notebook: vscode.NotebookDocument): boolean { const newKey = notebook.uri.toString(); const oldKey = this.notebookToUri.get(notebook); @@ -483,7 +550,288 @@ export class SqlNotebookController implements vscode.Disposable { this.log.info(`[rekeyConnectionOnSave] Re-keying connection: ${oldKey} → ${newKey}`); this.connections.delete(oldKey); this.connections.set(newKey, mgr); + return true; + } + return false; + } + + /** + * Handles a notebook document opening: sets kernel affinity for SQL + * notebooks and adopts any connection manager parked by the close of the + * untitled notebook this document replaced on save. + * Public for testing purposes. + */ + public handleNotebookOpened(notebook: vscode.NotebookDocument): void { + if (notebook.notebookType !== "jupyter-notebook") { + return; } + this.recordNotebookOpen(notebook); + this.setAffinityIfSql(notebook); + this.tryAdoptConnection(notebook); + } + + /** + * Finds the open notebook that owns a cell text document. VS Code derives + * cell URIs from the notebook URI (scheme swapped, fragment added), so + * path-matching notebooks are checked first and membership is always + * verified against the notebook's actual cells; a full scan remains as + * fallback in case the derivation ever changes. + */ + private findNotebookForCellDocument( + doc: vscode.TextDocument, + ): vscode.NotebookDocument | undefined { + const cellUri = doc.uri.toString(); + const ownsCell = (nb: vscode.NotebookDocument) => + nb.getCells().some((cell) => cell.document.uri.toString() === cellUri); + + for (const nb of vscode.workspace.notebookDocuments) { + if (nb.uri.path === doc.uri.path && ownsCell(nb)) { + return nb; + } + } + return vscode.workspace.notebookDocuments.find( + (nb) => nb.uri.path !== doc.uri.path && ownsCell(nb), + ); + } + + /** + * Record when a notebook opened (pruning stale entries) so the park-time + * adoption scan can restrict itself to notebooks opened around the save + * transition. + */ + private recordNotebookOpen(notebook: vscode.NotebookDocument): void { + const now = Date.now(); + for (const [key, openedAt] of this.recentNotebookOpens) { + if (now - openedAt > SAVE_ADOPTION_TTL_MS) { + this.recentNotebookOpens.delete(key); + } + } + this.recentNotebookOpens.set(notebook.uri.toString(), now); + } + + /** + * Whether the notebook opened recently enough to plausibly be the + * file-based replacement created by an untitled notebook's save. + */ + private wasRecentlyOpened(uriString: string): boolean { + const openedAt = this.recentNotebookOpens.get(uriString); + return openedAt !== undefined && Date.now() - openedAt <= SAVE_ADOPTION_TTL_MS; + } + + /** + * Whether any editor tab currently shows the given notebook URI. A + * notebook document replaced by a save keeps no tab, while a notebook the + * user still has open (even as a background tab) does. + * Public for testing purposes. + */ + public isNotebookOpenInTab(uriString: string): boolean { + for (const group of vscode.window.tabGroups.all) { + for (const tab of group.tabs) { + if ( + tab.input instanceof vscode.TabInputNotebook && + tab.input.uri.toString() === uriString + ) { + return true; + } + } + } + return false; + } + + /** + * Handles a notebook document closing. A connected UNTITLED notebook + * closing is usually the save-to-disk transition (VS Code replaces the + * document rather than renaming it), so its manager is parked for + * adoption instead of disposed. + * Public for testing purposes. + */ + public handleNotebookClosed(notebook: vscode.NotebookDocument): void { + const key = notebook.uri.toString(); + const mgr = this.connections.get(key); + if (!mgr) { + return; + } + this.connections.delete(key); + // Note: WeakMap entry will be garbage collected automatically + + if (notebook.uri.scheme === "untitled" && mgr.isConnected()) { + this.parkManagerForAdoption(key, notebook, mgr); + return; + } + + this.log.info(`[handleNotebookClosed] Disposing manager for ${key}`); + mgr.dispose(); + } + + /** + * Content signature used to match a closing untitled notebook with the + * file-based notebook that replaces it on save. Cell content is identical + * across the transition; cap the length to bound comparison cost. + */ + private notebookContentSignature(notebook: vscode.NotebookDocument): string { + // Accumulate the capped prefix and total length incrementally so a + // large notebook never materializes its full concatenated content + // just to produce a 10k-char signature. + const maxChars = 10000; + const cells = notebook.getCells(); + const separator = "\u0000"; + let totalLength = 0; + let prefix = ""; + for (let i = 0; i < cells.length; i++) { + const text = + i === 0 ? cells[i].document.getText() : separator + cells[i].document.getText(); + totalLength += text.length; + if (prefix.length < maxChars) { + prefix += text.slice(0, maxChars - prefix.length); + } + } + // Include the full content length so notebooks that differ only past + // the comparison cap still get distinct signatures. + return `${cells.length}:${totalLength}:${prefix}`; + } + + /** + * Park the connection manager of a closing untitled notebook so the + * file-based notebook created by the save can adopt it. If no notebook + * adopts it within the TTL (e.g. the untitled notebook was genuinely + * discarded), the manager is disposed, closing its connection. + */ + private parkManagerForAdoption( + oldKey: string, + notebook: vscode.NotebookDocument, + mgr: NotebookConnectionManager, + ): void { + // VS Code reuses untitled names, so a second park can arrive under the + // same key while the first is still pending. The superseded manager + // can never be adopted once replaced — dispose it and cancel its + // timer so it cannot fire against the new entry. + const superseded = this.pendingSaveAdoptions.get(oldKey); + if (superseded) { + clearTimeout(superseded.timer); + this.pendingSaveAdoptions.delete(oldKey); + this.log.info( + `[parkManagerForAdoption] Superseding parked manager for reused key ${oldKey}`, + ); + superseded.mgr.dispose(); + } + + const signature = this.notebookContentSignature(notebook); + const timer = setTimeout(() => { + // Only dispose the entry this timer was created for — a newer park + // under the same (reused) untitled URI must not be torn down by a + // stale timer. + const parked = this.pendingSaveAdoptions.get(oldKey); + if (parked?.mgr === mgr) { + this.pendingSaveAdoptions.delete(oldKey); + this.log.info( + `[parkManagerForAdoption] No adoption within ${SAVE_ADOPTION_TTL_MS}ms for ${oldKey}; disposing manager`, + ); + mgr.dispose(); + } + }, SAVE_ADOPTION_TTL_MS); + this.pendingSaveAdoptions.set(oldKey, { mgr, signature, timer }); + this.log.info(`[parkManagerForAdoption] Parked connected manager for ${oldKey}`); + + // If the file-based notebook opened BEFORE the untitled one closed, + // it is already in the workspace — adopt immediately. Only notebooks + // that opened around the save transition are considered, so a + // long-open unrelated notebook with matching content is never picked. + for (const candidate of vscode.workspace.notebookDocuments) { + if ( + candidate.notebookType === "jupyter-notebook" && + this.wasRecentlyOpened(candidate.uri.toString()) && + this.tryAdoptConnection(candidate) + ) { + break; + } + } + } + + /** + * Transfer a connection stranded on the untitled notebook this file-based + * notebook replaced on save. Two sources, matched by content signature: + * 1. Managers parked by the untitled notebook's close event. + * 2. Live managers still keyed by an untitled URI whose document lingers + * open — VS Code may keep the replaced untitled NotebookDocument open + * (or close it late), in which case no close event has fired yet. + * @returns true when an adoption occurred. + */ + private tryAdoptConnection(notebook: vscode.NotebookDocument): boolean { + if (notebook.uri.scheme === "untitled") { + return false; + } + const newKey = notebook.uri.toString(); + if (this.connections.has(newKey)) { + return false; + } + // Skip the (cell-content) signature computation when there is nothing + // that could possibly be adopted. + const hasLiveUntitledCandidate = [...this.connections.keys()].some((key) => + key.startsWith("untitled:"), + ); + if (this.pendingSaveAdoptions.size === 0 && !hasLiveUntitledCandidate) { + return false; + } + const signature = this.notebookContentSignature(notebook); + + // Source 1: manager parked by the untitled notebook's close event. + for (const [oldKey, parked] of this.pendingSaveAdoptions) { + if (parked.signature !== signature) { + continue; + } + clearTimeout(parked.timer); + this.pendingSaveAdoptions.delete(oldKey); + this.adoptManager(notebook, newKey, oldKey, parked.mgr); + return true; + } + + // Source 2: live manager whose untitled notebook document lingers open. + for (const [oldKey, mgr] of this.connections) { + if (!oldKey.startsWith("untitled:") || !mgr.isConnected()) { + continue; + } + // An untitled notebook that still has an editor tab is one the + // user is actively using — NOT a document orphaned by a save + // replacement. Never steal its connection just because content + // matches. + if (this.isNotebookOpenInTab(oldKey)) { + continue; + } + const oldNotebook = vscode.workspace.notebookDocuments.find( + (nb) => nb.uri.toString() === oldKey, + ); + if (!oldNotebook || this.notebookContentSignature(oldNotebook) !== signature) { + continue; + } + this.connections.delete(oldKey); + this.adoptManager(notebook, newKey, oldKey, mgr); + return true; + } + return false; + } + + /** + * Complete an adoption: re-key the manager to the new notebook, re-register + * the (new) cell document URIs with STS for IntelliSense, persist metadata, + * and refresh UI state. + */ + private adoptManager( + notebook: vscode.NotebookDocument, + newKey: string, + oldKey: string, + mgr: NotebookConnectionManager, + ): void { + this.connections.set(newKey, mgr); + this.notebookToUri.set(notebook, newKey); + this.log.info(`[adoptManager] Transferred connection ${oldKey} → ${newKey}`); + + // The old cell URIs belong to the replaced untitled notebook; release + // their STS registrations and register the new cell URIs. + mgr.releaseCellRegistrations(); + this.connectCellsForIntellisense(notebook, "adoptAfterSave"); + this.saveConnectionMetadataIfConnected(notebook); + this.updateStatusBar(notebook); + this.codeLensProvider.refresh(); } /** @@ -1195,6 +1543,11 @@ export class SqlNotebookController implements vscode.Disposable { mgr.dispose(); } this.connections.clear(); + for (const parked of this.pendingSaveAdoptions.values()) { + clearTimeout(parked.timer); + parked.mgr.dispose(); + } + this.pendingSaveAdoptions.clear(); this.disposables.forEach((d) => d.dispose()); this.statusBarItem.dispose(); this.controller.dispose(); diff --git a/extensions/mssql/test/unit/notebooks/notebookConnectionManager.test.ts b/extensions/mssql/test/unit/notebooks/notebookConnectionManager.test.ts index 4a124b922b..45d7556507 100644 --- a/extensions/mssql/test/unit/notebooks/notebookConnectionManager.test.ts +++ b/extensions/mssql/test/unit/notebooks/notebookConnectionManager.test.ts @@ -14,6 +14,11 @@ chai.use(sinonChai); import { NotebookConnectionManager } from "../../../src/notebooks/notebookConnectionManager"; import { ILogger } from "../../../src/sharedInterfaces/logger"; import ConnectionManager from "../../../src/controllers/connectionManager"; +import { + ConnectionCompleteParams, + ConnectionRequest, + DisconnectRequest, +} from "../../../src/models/contracts/connection"; import { ConnectionSharingService } from "../../../src/connectionSharing/connectionSharingService"; import { ConnectionStore } from "../../../src/models/connectionStore"; import { ConnectionUI } from "../../../src/views/connectionUI"; @@ -111,6 +116,11 @@ suite("NotebookConnectionManager", () => { connectionMgr.createConnectionDetails.returns(stubDetails); connectionMgr.sendRequest.resolves(true); connectionMgr.getConnectionInfoFromUri.returns(makeConnectionInfo({ database: "TestDB" })); + // Cell registration awaits the async connection/complete notification; + // default to a successful completion so registrations succeed. + connectionMgr.expectConnectionComplete.resolves({ + connectionId: "test-connection-id", + } as ConnectionCompleteParams); // --- ConnectionStore (getter stub) --- stubStore = sandbox.createStubInstance(ConnectionStore); @@ -796,6 +806,183 @@ suite("NotebookConnectionManager", () => { await mgr.connectCellForIntellisense("vscode-notebook-cell://cell1"); expect(connectionMgr.sendRequest).to.not.have.been.called; }); + + test("builds cell connection details from live adhoc connection credentials", async () => { + // The adhoc (execution) connection's credentials are the ones STS + // validated — including refreshed tokens and the actual database — + // so cell registration must prefer them over the stored profile. + await mgr.connectWith(makeConnectionInfo({ database: "RequestedDB" })); + connectionMgr.getConnectionInfoFromUri.returns( + makeConnectionInfo({ database: "LiveDB" }), + ); + connectionMgr.createConnectionDetails.resetHistory(); + + await mgr.connectCellForIntellisense("vscode-notebook-cell://cell1"); + + expect(connectionMgr.createConnectionDetails).to.have.been.calledWithMatch( + sinon.match({ database: "LiveDB" }), + ); + }); + + test("re-sends after connection/complete reports failure (no false memoization)", async () => { + // The connect REQUEST resolving true only means the attempt started. + // A failed completion (no connectionId) must deregister the cell so + // the next IntelliSense trigger retries instead of leaving the cell + // with keyword-only completions. + await mgr.connectWith(makeConnectionInfo()); + connectionMgr.expectConnectionComplete.resolves({ + errorMessage: "Login failed", + } as ConnectionCompleteParams); + + await mgr.connectCellForIntellisense("vscode-notebook-cell://cell1"); + expect(log.warn).to.have.been.called; + + // Recover and retry — must fire a new connect request. + connectionMgr.sendRequest.resetHistory(); + connectionMgr.expectConnectionComplete.resolves({ + connectionId: "recovered-connection-id", + } as ConnectionCompleteParams); + await mgr.connectCellForIntellisense("vscode-notebook-cell://cell1"); + expect(connectionMgr.sendRequest).to.have.been.calledWith( + ConnectionRequest.type, + sinon.match({ ownerUri: "vscode-notebook-cell://cell1" }), + ); + }); + + test("re-sends after connection/complete times out (no false memoization)", async () => { + await mgr.connectWith(makeConnectionInfo()); + + const clock = sandbox.useFakeTimers(); + // Completion never arrives. + connectionMgr.expectConnectionComplete.returns( + new Promise(() => {}), + ); + + const pendingRegistration = mgr.connectCellForIntellisense( + "vscode-notebook-cell://cell1", + ); + await clock.tickAsync(31000); + await pendingRegistration; + + // The pending expectation must be cleaned up on timeout. + expect(connectionMgr.cancelConnectionCompleteExpectation).to.have.been.calledWith( + "vscode-notebook-cell://cell1", + ); + + // A best-effort disconnect must be sent — the connect may still + // land in STS later and this URI is no longer tracked for release. + expect(connectionMgr.sendRequest).to.have.been.calledWith( + DisconnectRequest.type, + sinon.match({ ownerUri: "vscode-notebook-cell://cell1" }), + ); + + clock.restore(); + + // Retry — must fire a new connect request. + connectionMgr.sendRequest.resetHistory(); + connectionMgr.expectConnectionComplete.resolves({ + connectionId: "recovered-connection-id", + } as ConnectionCompleteParams); + await mgr.connectCellForIntellisense("vscode-notebook-cell://cell1"); + expect(connectionMgr.sendRequest).to.have.been.calledWith( + ConnectionRequest.type, + sinon.match({ ownerUri: "vscode-notebook-cell://cell1" }), + ); + }); + }); + + // ---------------------------------------------------------------- + // releaseCellRegistrations + // ---------------------------------------------------------------- + suite("releaseCellRegistrations", () => { + test("disconnects registered cell URIs and clears memoization", async () => { + await mgr.connectWith(makeConnectionInfo()); + await mgr.connectCellForIntellisense("vscode-notebook-cell://cell1"); + await mgr.connectCellForIntellisense("vscode-notebook-cell://cell2"); + connectionMgr.sendRequest.resetHistory(); + + mgr.releaseCellRegistrations(); + + expect(connectionMgr.sendRequest).to.have.been.calledWith( + DisconnectRequest.type, + sinon.match({ ownerUri: "vscode-notebook-cell://cell1" }), + ); + expect(connectionMgr.sendRequest).to.have.been.calledWith( + DisconnectRequest.type, + sinon.match({ ownerUri: "vscode-notebook-cell://cell2" }), + ); + + // Memoization cleared — cells can be re-registered under new URIs + // (or the same URI) with a fresh connect request. + connectionMgr.sendRequest.resetHistory(); + await mgr.connectCellForIntellisense("vscode-notebook-cell://cell1"); + expect(connectionMgr.sendRequest).to.have.been.calledWith( + ConnectionRequest.type, + sinon.match({ ownerUri: "vscode-notebook-cell://cell1" }), + ); + }); + + test("is a no-op when no cells are registered", () => { + mgr.releaseCellRegistrations(); + expect(connectionMgr.sendRequest).to.not.have.been.called; + }); + + test("survives disconnect request failures", async () => { + await mgr.connectWith(makeConnectionInfo()); + await mgr.connectCellForIntellisense("vscode-notebook-cell://cell1"); + connectionMgr.sendRequest.rejects(new Error("STS unavailable")); + + expect(() => mgr.releaseCellRegistrations()).to.not.throw(); + + // Let the rejected fire-and-forget promise settle so the warn logs. + await new Promise((resolve) => setImmediate(resolve)); + expect(log.warn).to.have.been.called; + }); + + test("disconnect releases registered cell connections from STS", async () => { + // Each registered cell holds its own STS-side connection; notebook + // disconnect (including notebook close → dispose) must close them, + // not just forget them. + await mgr.connectWith(makeConnectionInfo()); + await mgr.connectCellForIntellisense("vscode-notebook-cell://cell1"); + connectionMgr.sendRequest.resetHistory(); + + mgr.disconnect(); + + expect(connectionMgr.sendRequest).to.have.been.calledWith( + DisconnectRequest.type, + sinon.match({ ownerUri: "vscode-notebook-cell://cell1" }), + ); + }); + + test("stale registration whose connect completed after a connection change is disconnected", async () => { + await mgr.connectWith(makeConnectionInfo()); + + // The cell's connect request is accepted, but before the + // connection/complete arrives the notebook connection changes. + let resolveComplete: (params: ConnectionCompleteParams) => void = () => {}; + connectionMgr.expectConnectionComplete.returns( + new Promise((resolve) => { + resolveComplete = resolve; + }), + ); + const staleRegistration = mgr.connectCellForIntellisense( + "vscode-notebook-cell://cell1", + ); + + await mgr.changeDatabase("OtherDB"); + connectionMgr.sendRequest.resetHistory(); + + // The stale connect then completes successfully — its now-unwanted + // STS connection must be explicitly disconnected. + resolveComplete({ connectionId: "stale-conn-id" } as ConnectionCompleteParams); + await staleRegistration; + + expect(connectionMgr.sendRequest).to.have.been.calledWith( + DisconnectRequest.type, + sinon.match({ ownerUri: "vscode-notebook-cell://cell1" }), + ); + }); }); // ---------------------------------------------------------------- diff --git a/extensions/mssql/test/unit/notebooks/sqlNotebookController.test.ts b/extensions/mssql/test/unit/notebooks/sqlNotebookController.test.ts index 1e2fcda0a6..b2763d4166 100644 --- a/extensions/mssql/test/unit/notebooks/sqlNotebookController.test.ts +++ b/extensions/mssql/test/unit/notebooks/sqlNotebookController.test.ts @@ -121,6 +121,7 @@ suite("SqlNotebookController", () => { function makeNotebook( cells?: Array<{ text: string; languageId?: string; kind?: vscode.NotebookCellKind }>, metadata?: Record, + uri: vscode.Uri = notebookUri, ): vscode.NotebookDocument { const cellObjs = (cells ?? []).map((c, i) => ({ index: i, @@ -128,11 +129,13 @@ suite("SqlNotebookController", () => { document: { getText: () => c.text, languageId: c.languageId ?? "sql", - uri: vscode.Uri.parse(`vscode-notebook-cell://test-notebook#cell${i}`), + // Mirror VS Code's cell URI shape: notebook URI with the + // vscode-notebook-cell scheme and a per-cell fragment. + uri: uri.with({ scheme: "vscode-notebook-cell", fragment: `cell${i}` }), }, })); return { - uri: notebookUri, + uri, notebookType: "jupyter-notebook", metadata: metadata ?? {}, getCells: () => cellObjs, @@ -145,6 +148,15 @@ suite("SqlNotebookController", () => { keys: sinon.SinonStub; }; + // Handlers captured from the controller's event subscriptions so tests can + // simulate workspace events, plus a mutable open-notebooks array. + // Note: onDidOpen/CloseNotebookDocument cannot be sinon-stubbed on + // vscode.workspace in the test host; tests invoke the controller's + // handleNotebookOpened/handleNotebookClosed methods directly instead. + let didSaveNotebookHandler: (notebook: vscode.NotebookDocument) => void; + let didOpenTextDocumentHandler: (doc: vscode.TextDocument) => void; + let openNotebookDocuments: vscode.NotebookDocument[]; + function setupVscodeMocks(sb: sinon.SinonSandbox): void { sb.stub(vscode.notebooks, "createNotebookController").returns( mockController as unknown as vscode.NotebookController, @@ -175,8 +187,13 @@ suite("SqlNotebookController", () => { sb.stub(vscode.workspace, "onDidChangeNotebookDocument").returns({ dispose: sb.stub(), }); - sb.stub(vscode.workspace, "onDidSaveNotebookDocument").returns({ - dispose: sb.stub(), + sb.stub(vscode.workspace, "onDidSaveNotebookDocument").callsFake((listener) => { + didSaveNotebookHandler = listener as (notebook: vscode.NotebookDocument) => void; + return { dispose: sb.stub() }; + }); + sb.stub(vscode.workspace, "onDidOpenTextDocument").callsFake((listener) => { + didOpenTextDocumentHandler = listener as (doc: vscode.TextDocument) => void; + return { dispose: sb.stub() }; }); sb.stub(vscode.workspace, "onDidCloseNotebookDocument").returns({ dispose: sb.stub(), @@ -184,7 +201,8 @@ suite("SqlNotebookController", () => { sb.stub(vscode.languages, "registerCodeLensProvider").returns({ dispose: sb.stub(), }); - sb.stub(vscode.workspace, "notebookDocuments").value([]); + openNotebookDocuments = []; + sb.stub(vscode.workspace, "notebookDocuments").value(openNotebookDocuments); sb.stub(vscode.workspace, "applyEdit").resolves(true); } @@ -1020,6 +1038,309 @@ suite("SqlNotebookController", () => { }); }); + suite("save URI transfer", () => { + test("re-keys manager and re-registers cells when notebook URI changes on save", async () => { + const untitledUri = vscode.Uri.parse("untitled:Untitled-1.ipynb"); + const notebook = makeNotebook([{ text: "SELECT 1" }], undefined, untitledUri); + + // Execute once so the manager is tracked under the untitled URI. + await mockController.executeHandler(notebook.getCells(), notebook, mockController); + expect(controller.connections.get(untitledUri.toString())).to.equal( + mockNotebookConnMgr, + ); + + // Simulate the save assigning the final file URI. + const fileUri = vscode.Uri.parse("file:///c:/temp/notebook1.ipynb"); + (notebook as { uri: vscode.Uri }).uri = fileUri; + mockNotebookConnMgr.connectCellForIntellisense.resetHistory(); + + didSaveNotebookHandler(notebook); + + // Manager re-keyed to the new URI. + expect(controller.connections.has(untitledUri.toString())).to.be.false; + expect(controller.connections.get(fileUri.toString())).to.equal(mockNotebookConnMgr); + + // Stale cell registrations released and current cells re-registered + // so IntelliSense keeps working under the new cell URIs. + expect(mockNotebookConnMgr.releaseCellRegistrations).to.have.been.calledOnce; + expect(mockNotebookConnMgr.connectCellForIntellisense).to.have.been.calledWith( + notebook.getCells()[0].document.uri.toString(), + ); + }); + + test("does not release registrations when URI is unchanged on save", async () => { + const notebook = makeNotebook([{ text: "SELECT 1" }]); + await mockController.executeHandler(notebook.getCells(), notebook, mockController); + + didSaveNotebookHandler(notebook); + + expect(mockNotebookConnMgr.releaseCellRegistrations).to.not.have.been.called; + }); + }); + + suite("untitled save adoption (close + reopen transition)", () => { + const untitledUri = vscode.Uri.parse("untitled:Untitled-1.ipynb"); + const savedUri = vscode.Uri.parse("file:///c:/src/note1.ipynb"); + + /** Connect an untitled notebook so its manager is tracked by the controller. */ + async function connectUntitledNotebook( + text = "SELECT 1", + ): Promise { + const notebook = makeNotebook([{ text }], undefined, untitledUri); + await mockController.executeHandler(notebook.getCells(), notebook, mockController); + expect(controller.connections.get(untitledUri.toString())).to.equal( + mockNotebookConnMgr, + ); + return notebook; + } + + test("close then open: parked manager is adopted by matching file notebook", async () => { + const untitled = await connectUntitledNotebook(); + + // Save replaces the document: the untitled notebook closes... + controller.handleNotebookClosed(untitled); + expect(controller.connections.has(untitledUri.toString())).to.be.false; + // ...but the connected manager must survive the transition. + expect(mockNotebookConnMgr.dispose).to.not.have.been.called; + + // ...and the file-based notebook opens with the same content. + const saved = makeNotebook([{ text: "SELECT 1" }], undefined, savedUri); + mockNotebookConnMgr.connectCellForIntellisense.resetHistory(); + controller.handleNotebookOpened(saved); + + expect(controller.connections.get(savedUri.toString())).to.equal(mockNotebookConnMgr); + expect(mockNotebookConnMgr.dispose).to.not.have.been.called; + // Old cell registrations released, new cell URIs registered. + expect(mockNotebookConnMgr.releaseCellRegistrations).to.have.been.calledOnce; + expect(mockNotebookConnMgr.connectCellForIntellisense).to.have.been.calledWith( + saved.getCells()[0].document.uri.toString(), + ); + // Connection metadata persisted under the new file URI. + expect(mockWorkspaceState.update).to.have.been.calledWith( + `notebook.connection.${savedUri.toString()}`, + sinon.match({ server: "test-server", database: "TestDB" }), + ); + }); + + test("open then close: adoption still occurs when the file notebook opens first", async () => { + const untitled = await connectUntitledNotebook(); + + // The file-based notebook opens (and is seen opening) before the + // untitled one closes. + const saved = makeNotebook([{ text: "SELECT 1" }], undefined, savedUri); + openNotebookDocuments.push(saved); + controller.handleNotebookOpened(saved); + + controller.handleNotebookClosed(untitled); + + expect(controller.connections.get(savedUri.toString())).to.equal(mockNotebookConnMgr); + expect(mockNotebookConnMgr.dispose).to.not.have.been.called; + expect(mockNotebookConnMgr.releaseCellRegistrations).to.have.been.calledOnce; + }); + + test("park-time scan ignores notebooks that were not recently opened", async () => { + const untitled = await connectUntitledNotebook(); + + // A matching file notebook is open but was NEVER seen opening in + // the adoption window (e.g. a long-open unrelated notebook). + const oldNotebook = makeNotebook([{ text: "SELECT 1" }], undefined, savedUri); + openNotebookDocuments.push(oldNotebook); + + controller.handleNotebookClosed(untitled); + + // Parked, not adopted into the stale notebook. + expect(controller.connections.has(savedUri.toString())).to.be.false; + expect(mockNotebookConnMgr.dispose).to.not.have.been.called; + }); + + test("does not steal a connected untitled notebook that still has an open tab", async () => { + const untitled = await connectUntitledNotebook(); + openNotebookDocuments.push(untitled); + // The untitled notebook still has an editor tab — the user is + // actively using it; it is not a save-replacement leftover. + sandbox.stub(controller, "isNotebookOpenInTab").returns(true); + + const saved = makeNotebook([{ text: "SELECT 1" }], undefined, savedUri); + controller.handleNotebookOpened(saved); + + expect(controller.connections.has(savedUri.toString())).to.be.false; + expect(controller.connections.get(untitledUri.toString())).to.equal( + mockNotebookConnMgr, + ); + }); + + test("adopts a live manager when the untitled document lingers open (no close event)", async () => { + const untitled = await connectUntitledNotebook(); + // The replaced untitled document stays open — no close event fires. + openNotebookDocuments.push(untitled); + + const saved = makeNotebook([{ text: "SELECT 1" }], undefined, savedUri); + mockNotebookConnMgr.connectCellForIntellisense.resetHistory(); + controller.handleNotebookOpened(saved); + + expect(controller.connections.get(savedUri.toString())).to.equal(mockNotebookConnMgr); + expect(controller.connections.has(untitledUri.toString())).to.be.false; + expect(mockNotebookConnMgr.dispose).to.not.have.been.called; + expect(mockNotebookConnMgr.releaseCellRegistrations).to.have.been.calledOnce; + expect(mockNotebookConnMgr.connectCellForIntellisense).to.have.been.calledWith( + saved.getCells()[0].document.uri.toString(), + ); + }); + + test("save event adopts a lingering manager when the open event found no match", async () => { + const untitled = await connectUntitledNotebook(); + openNotebookDocuments.push(untitled); + + const saved = makeNotebook([{ text: "SELECT 1" }], undefined, savedUri); + didSaveNotebookHandler(saved); + + expect(controller.connections.get(savedUri.toString())).to.equal(mockNotebookConnMgr); + expect(mockNotebookConnMgr.releaseCellRegistrations).to.have.been.calledOnce; + }); + + test("does not steal a live manager whose untitled document content differs", async () => { + const untitled = await connectUntitledNotebook("SELECT 1"); + openNotebookDocuments.push(untitled); + + const saved = makeNotebook([{ text: "SELECT 999" }], undefined, savedUri); + controller.handleNotebookOpened(saved); + + // No transfer — the untitled notebook keeps its manager. + expect(controller.connections.has(savedUri.toString())).to.be.false; + expect(controller.connections.get(untitledUri.toString())).to.equal( + mockNotebookConnMgr, + ); + }); + + test("does not adopt into a notebook with different content; disposes after TTL", async () => { + const clock = sandbox.useFakeTimers(); + const untitled = await connectUntitledNotebook("SELECT 1"); + + controller.handleNotebookClosed(untitled); + + const unrelated = makeNotebook([{ text: "SELECT 999" }], undefined, savedUri); + controller.handleNotebookOpened(unrelated); + expect(controller.connections.has(savedUri.toString())).to.be.false; + + // Nothing adopted the parked manager — it must be disposed after the TTL + // so the orphaned connection is cleaned up. + await clock.tickAsync(11000); + expect(mockNotebookConnMgr.dispose).to.have.been.called; + + clock.restore(); + }); + + test("disposes immediately when a file-based notebook closes", async () => { + const notebook = makeNotebook([{ text: "SELECT 1" }], undefined, savedUri); + await mockController.executeHandler(notebook.getCells(), notebook, mockController); + + controller.handleNotebookClosed(notebook); + + expect(mockNotebookConnMgr.dispose).to.have.been.called; + expect(controller.connections.has(savedUri.toString())).to.be.false; + }); + + test("disposes immediately when a disconnected untitled notebook closes", async () => { + const untitled = await connectUntitledNotebook(); + mockNotebookConnMgr.isConnected.returns(false); + + controller.handleNotebookClosed(untitled); + + expect(mockNotebookConnMgr.dispose).to.have.been.called; + }); + + test("re-parking a reused untitled URI disposes the superseded manager and keeps the new one adoptable", async () => { + // VS Code recycles untitled names: two different notebooks can be + // parked under the same untitled URI within the TTL window. + const mgrA = sandbox.createStubInstance(NotebookConnectionManager); + mgrA.isConnected.returns(true); + const untitledA = makeNotebook([{ text: "SELECT 1" }], undefined, untitledUri); + controller.connections.set(untitledUri.toString(), mgrA); + controller.handleNotebookClosed(untitledA); + + const mgrB = sandbox.createStubInstance(NotebookConnectionManager); + mgrB.isConnected.returns(true); + const untitledB = makeNotebook([{ text: "SELECT 2" }], undefined, untitledUri); + controller.connections.set(untitledUri.toString(), mgrB); + controller.handleNotebookClosed(untitledB); + + // The superseded first manager is disposed; the second survives. + expect(mgrA.dispose).to.have.been.called; + expect(mgrB.dispose).to.not.have.been.called; + + // The second manager is still adoptable by its saved notebook. + const saved = makeNotebook([{ text: "SELECT 2" }], undefined, savedUri); + controller.handleNotebookOpened(saved); + expect(controller.connections.get(savedUri.toString())).to.equal(mgrB); + }); + + test("controller dispose cleans up parked managers", async () => { + const untitled = await connectUntitledNotebook(); + controller.handleNotebookClosed(untitled); + expect(mockNotebookConnMgr.dispose).to.not.have.been.called; + + controller.dispose(); + + expect(mockNotebookConnMgr.dispose).to.have.been.called; + }); + }); + + suite("late cell document registration", () => { + test("registers a cell document that opens while the notebook is connected", async () => { + const notebook = makeNotebook([{ text: "SELECT 1" }]); + openNotebookDocuments.push(notebook); + + // Execute once so a connected manager exists for the notebook. + await mockController.executeHandler(notebook.getCells(), notebook, mockController); + mockNotebookConnMgr.connectCellForIntellisense.resetHistory(); + + const cellDoc = notebook.getCells()[0].document; + didOpenTextDocumentHandler(cellDoc); + + expect(mockNotebookConnMgr.connectCellForIntellisense).to.have.been.calledWith( + cellDoc.uri.toString(), + ); + }); + + test("ignores non-cell documents", async () => { + const notebook = makeNotebook([{ text: "SELECT 1" }]); + openNotebookDocuments.push(notebook); + await mockController.executeHandler(notebook.getCells(), notebook, mockController); + mockNotebookConnMgr.connectCellForIntellisense.resetHistory(); + + didOpenTextDocumentHandler({ + uri: vscode.Uri.parse("file:///c:/temp/query.sql"), + languageId: "sql", + } as vscode.TextDocument); + + expect(mockNotebookConnMgr.connectCellForIntellisense).to.not.have.been.called; + }); + + test("skips cell documents whose notebook has no manager", () => { + const notebook = makeNotebook([{ text: "SELECT 1" }]); + openNotebookDocuments.push(notebook); + + const cellDoc = notebook.getCells()[0].document; + didOpenTextDocumentHandler(cellDoc); + + expect(mockNotebookConnMgr.connectCellForIntellisense).to.not.have.been.called; + }); + + test("skips cell documents when the notebook is not connected", async () => { + const notebook = makeNotebook([{ text: "SELECT 1" }]); + openNotebookDocuments.push(notebook); + await mockController.executeHandler(notebook.getCells(), notebook, mockController); + + mockNotebookConnMgr.isConnected.returns(false); + mockNotebookConnMgr.connectCellForIntellisense.resetHistory(); + + const cellDoc = notebook.getCells()[0].document; + didOpenTextDocumentHandler(cellDoc); + + expect(mockNotebookConnMgr.connectCellForIntellisense).to.not.have.been.called; + }); + }); + suite("connection metadata persistence", () => { test("saves metadata to workspaceState after connectWith", async () => { const mockNotebook = makeNotebook();