From e796403e40c9306a737b3857682834c7e7a9a08e Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Wed, 25 Mar 2026 16:08:32 +0000 Subject: [PATCH 01/26] =?UTF-8?q?feat(scratchpad):=20Phase=201=20=E2=80=94?= =?UTF-8?q?=20worker=20thread=20scaffold?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Create workerTypes.ts with typed IPC message protocol (MainToWorkerMessage, WorkerToMainMessage) - Create scratchpadWorker.ts with init/eval/shutdown/tokenRequest handlers - Add scratchpadWorker webpack entry point to webpack.config.ext.js - Worker lazy-imports @mongosh/* packages (same pattern as current evaluator) - Supports both SCRAM and Entra ID auth (OIDC via IPC token callback) - Worker logs lifecycle events to main thread via 'log' messages Step 6.2 WI-1, Phase 1 of 3. --- src/documentdb/scratchpad/scratchpadWorker.ts | 304 ++++++++++++++++++ src/documentdb/scratchpad/workerTypes.ts | 124 +++++++ webpack.config.ext.js | 1 + 3 files changed, 429 insertions(+) create mode 100644 src/documentdb/scratchpad/scratchpadWorker.ts create mode 100644 src/documentdb/scratchpad/workerTypes.ts diff --git a/src/documentdb/scratchpad/scratchpadWorker.ts b/src/documentdb/scratchpad/scratchpadWorker.ts new file mode 100644 index 00000000..c86ba2c9 --- /dev/null +++ b/src/documentdb/scratchpad/scratchpadWorker.ts @@ -0,0 +1,304 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +/** + * Worker thread entry point for scratchpad code evaluation. + * + * This file runs in a Node.js worker_thread, isolated from the extension host. + * It owns its own MongoClient instance (authenticated via credentials passed from + * the main thread at init time) and evaluates user code through the @mongosh pipeline. + * + * Communication with the main thread is via postMessage (structured clone). + * See workerTypes.ts for the message protocol. + */ + +import { randomUUID } from 'crypto'; +import { parentPort } from 'worker_threads'; +import type { MainToWorkerMessage, WorkerToMainMessage } from './workerTypes'; + +if (!parentPort) { + throw new Error('scratchpadWorker.ts must be run as a worker_thread'); +} + +// ─── Types for lazily-imported @mongosh packages ───────────────────────────── + +type MongoClientType = import('mongodb').MongoClient; + +// ─── Worker state ──────────────────────────────────────────────────────────── + +let mongoClient: MongoClientType | undefined; +let currentDatabaseName: string | undefined; + +/** + * Cache for pending Entra ID token requests from the OIDC callback. + * The OIDC_CALLBACK in the worker sends a tokenRequest to the main thread + * and awaits the response via this map. + */ +const pendingTokenRequests = new Map void; reject: (err: Error) => void }>(); + +// ─── Logging helper ────────────────────────────────────────────────────────── + +function log(level: 'info' | 'warn' | 'error' | 'debug', message: string): void { + const msg: WorkerToMainMessage = { type: 'log', level, message }; + parentPort!.postMessage(msg); +} + +// ─── Message handler ───────────────────────────────────────────────────────── + +parentPort.on('message', (msg: MainToWorkerMessage) => { + switch (msg.type) { + case 'init': + void handleInit(msg).catch((err: unknown) => { + const errorMessage = err instanceof Error ? err.message : String(err); + const response: WorkerToMainMessage = { + type: 'initResult', + requestId: msg.requestId, + success: false, + error: errorMessage, + }; + parentPort!.postMessage(response); + }); + break; + + case 'eval': + void handleEval(msg).catch((err: unknown) => { + const errorMessage = err instanceof Error ? err.message : String(err); + const stack = err instanceof Error ? err.stack : undefined; + const response: WorkerToMainMessage = { + type: 'evalError', + requestId: msg.requestId, + error: errorMessage, + stack, + }; + parentPort!.postMessage(response); + }); + break; + + case 'shutdown': + void handleShutdown(msg); + break; + + case 'tokenResponse': + handleTokenResponse(msg); + break; + + case 'tokenError': + handleTokenError(msg); + break; + } +}); + +// ─── Init handler ──────────────────────────────────────────────────────────── + +async function handleInit( + msg: Extract, +): Promise { + log('info', `Initializing worker (auth: ${msg.authMechanism}, db: ${msg.databaseName})`); + + // Lazy-import MongoDB driver + const { MongoClient } = await import('mongodb'); + + // Build MongoClient options from the serializable subset + const options: import('mongodb').MongoClientOptions = { + ...msg.clientOptions, + }; + + // For Entra ID, configure OIDC callback that requests tokens via IPC + if (msg.authMechanism === 'MicrosoftEntraID') { + options.authMechanism = 'MONGODB-OIDC'; + options.tls = true; + options.authMechanismProperties = { + ALLOWED_HOSTS: ['*.azure.com'], + OIDC_CALLBACK: async (): Promise<{ accessToken: string; expiresInSeconds: number }> => { + const requestId = randomUUID(); + const tokenPromise = new Promise((resolve, reject) => { + pendingTokenRequests.set(requestId, { resolve, reject }); + }); + const tokenRequest: WorkerToMainMessage = { + type: 'tokenRequest', + requestId, + scopes: ['https://ossrdbms-aad.database.windows.net/.default'], + tenantId: msg.tenantId, + }; + parentPort!.postMessage(tokenRequest); + const accessToken = await tokenPromise; + return { accessToken, expiresInSeconds: 0 }; + }, + }; + } + + // Create and connect MongoClient + mongoClient = new MongoClient(msg.connectionString, options); + await mongoClient.connect(); + currentDatabaseName = msg.databaseName; + + log('info', 'Worker initialized — MongoClient connected'); + + const response: WorkerToMainMessage = { + type: 'initResult', + requestId: msg.requestId, + success: true, + }; + parentPort!.postMessage(response); +} + +// ─── Eval handler ──────────────────────────────────────────────────────────── + +async function handleEval( + msg: Extract, +): Promise { + if (!mongoClient) { + throw new Error('Worker not initialized — call init first'); + } + + log('debug', `Evaluating code (${msg.code.length} chars, db: ${msg.databaseName})`); + + // Lazy-import @mongosh packages + const { EventEmitter } = await import('events'); + const vm = await import('vm'); + const { NodeDriverServiceProvider } = await import('@mongosh/service-provider-node-driver'); + const { ShellInstanceState } = await import('@mongosh/shell-api'); + const { ShellEvaluator } = await import('@mongosh/shell-evaluator'); + + const startTime = Date.now(); + + // Create fresh shell context per execution (no variable leakage between runs) + const bus = new EventEmitter(); + const serviceProvider = new NodeDriverServiceProvider(mongoClient, bus, { + productDocsLink: 'https://github.com/microsoft/vscode-documentdb', + productName: 'DocumentDB for VS Code Scratchpad', + }); + const instanceState = new ShellInstanceState(serviceProvider, bus); + const evaluator = new ShellEvaluator(instanceState); + + // Set up eval context with shell globals (db, ObjectId, ISODate, etc.) + const context = {}; + instanceState.setCtx(context); + + // The eval function using vm.runInContext for @mongosh + // eslint-disable-next-line @typescript-eslint/require-await + const customEvalFn = async (code: string, ctx: object): Promise => { + const vmContext = vm.isContext(ctx) ? ctx : vm.createContext(ctx); + return vm.runInContext(code, vmContext) as unknown; + }; + + // Switch database if different from current + if (msg.databaseName !== currentDatabaseName) { + await evaluator.customEval( + customEvalFn, + `use(${JSON.stringify(msg.databaseName)})`, + context, + 'scratchpad', + ); + currentDatabaseName = msg.databaseName; + } else { + // Pre-select the target database (fresh context each time) + await evaluator.customEval( + customEvalFn, + `use(${JSON.stringify(msg.databaseName)})`, + context, + 'scratchpad', + ); + } + + // Evaluate user code + const result = await evaluator.customEval(customEvalFn, msg.code, context, 'scratchpad'); + const durationMs = Date.now() - startTime; + + // result is a ShellResult { type, printable, rawValue, source? } + const shellResult = result as { + type: string | null; + printable: unknown; + source?: { namespace?: { db: string; collection: string } }; + }; + + // Serialize the printable value via EJSON for safe transfer + let printableStr: string; + try { + const { EJSON } = await import('bson'); + printableStr = EJSON.stringify(shellResult.printable, { relaxed: true }, 2); + } catch { + // Fallback: try JSON, then plain string + try { + printableStr = JSON.stringify(shellResult.printable, null, 2); + } catch { + printableStr = String(shellResult.printable); + } + } + + log('debug', `Evaluation complete (${durationMs}ms, type: ${shellResult.type ?? 'null'})`); + + const response: WorkerToMainMessage = { + type: 'evalResult', + requestId: msg.requestId, + result: { + type: shellResult.type, + printable: printableStr, + durationMs, + source: shellResult.source?.namespace + ? { + namespace: { + db: shellResult.source.namespace.db, + collection: shellResult.source.namespace.collection, + }, + } + : undefined, + }, + }; + parentPort!.postMessage(response); +} + +// ─── Shutdown handler ──────────────────────────────────────────────────────── + +async function handleShutdown( + msg: Extract, +): Promise { + log('info', 'Shutting down worker — closing MongoClient'); + + try { + if (mongoClient) { + await mongoClient.close(); + mongoClient = undefined; + } + } catch (err: unknown) { + const errorMessage = err instanceof Error ? err.message : String(err); + log('warn', `Error closing MongoClient during shutdown: ${errorMessage}`); + } + + const response: WorkerToMainMessage = { + type: 'shutdownComplete', + requestId: msg.requestId, + }; + parentPort!.postMessage(response); +} + +// ─── Token response/error handlers (Entra ID) ─────────────────────────────── + +function handleTokenResponse(msg: Extract): void { + const pending = pendingTokenRequests.get(msg.requestId); + if (pending) { + pending.resolve(msg.accessToken); + pendingTokenRequests.delete(msg.requestId); + } +} + +function handleTokenError(msg: Extract): void { + const pending = pendingTokenRequests.get(msg.requestId); + if (pending) { + pending.reject(new Error(msg.error)); + pendingTokenRequests.delete(msg.requestId); + } +} + +// ─── Uncaught exception handler ────────────────────────────────────────────── + +process.on('uncaughtException', (error: Error) => { + log('error', `Uncaught exception in worker: ${error.message}\n${error.stack ?? ''}`); +}); + +process.on('unhandledRejection', (reason: unknown) => { + const message = reason instanceof Error ? reason.message : String(reason); + log('error', `Unhandled rejection in worker: ${message}`); +}); diff --git a/src/documentdb/scratchpad/workerTypes.ts b/src/documentdb/scratchpad/workerTypes.ts new file mode 100644 index 00000000..2004c462 --- /dev/null +++ b/src/documentdb/scratchpad/workerTypes.ts @@ -0,0 +1,124 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +/** + * IPC message types for the scratchpad worker thread. + * + * This file is shared between the main thread (ScratchpadEvaluator) and + * the worker thread (scratchpadWorker). It must have zero runtime dependencies — + * only TypeScript types and string literal unions. + * + * Communication uses Node.js `worker_threads` `postMessage()` with the + * structured clone algorithm. Functions cannot be sent — this is why + * Entra ID OIDC tokens must be requested via IPC (tokenRequest/tokenResponse). + */ + +// ─── Serializable subset of MongoClientOptions ────────────────────────────── + +/** + * Only the MongoClientOptions fields that can survive structured clone. + * Function-valued options (like OIDC_CALLBACK) are stripped before sending + * and reconstructed on the worker side. + */ +export interface SerializableMongoClientOptions { + readonly serverSelectionTimeoutMS?: number; + readonly tlsAllowInvalidCertificates?: boolean; + readonly appName?: string; + readonly tls?: boolean; +} + +// ─── Serializable execution result ────────────────────────────────────────── + +/** + * The subset of ExecutionResult that can be sent via postMessage. + * BSON types are serialized to EJSON strings by the worker before sending. + */ +export interface SerializableExecutionResult { + readonly type: string | null; + /** EJSON-serialized printable value */ + readonly printable: string; + readonly durationMs: number; + readonly source?: { + readonly namespace?: { + readonly db: string; + readonly collection: string; + }; + }; +} + +// ─── Main → Worker messages ───────────────────────────────────────────────── + +export type MainToWorkerMessage = + | { + readonly type: 'init'; + readonly requestId: string; + /** Connection string (with embedded credentials for SCRAM, without for Entra ID) */ + readonly connectionString: string; + readonly clientOptions: SerializableMongoClientOptions; + readonly databaseName: string; + readonly authMechanism: 'NativeAuth' | 'MicrosoftEntraID'; + /** Tenant ID for Entra ID clusters */ + readonly tenantId?: string; + /** Display batch size from documentDB.mongoShell.batchSize setting */ + readonly displayBatchSize: number; + } + | { + readonly type: 'eval'; + readonly requestId: string; + /** JavaScript code to evaluate */ + readonly code: string; + /** Target database name (may differ from init if user switched databases) */ + readonly databaseName: string; + } + | { + readonly type: 'shutdown'; + readonly requestId: string; + } + | { + readonly type: 'tokenResponse'; + readonly requestId: string; + readonly accessToken: string; + } + | { + readonly type: 'tokenError'; + readonly requestId: string; + readonly error: string; + }; + +// ─── Worker → Main messages ───────────────────────────────────────────────── + +export type WorkerToMainMessage = + | { + readonly type: 'initResult'; + readonly requestId: string; + readonly success: boolean; + readonly error?: string; + } + | { + readonly type: 'evalResult'; + readonly requestId: string; + readonly result: SerializableExecutionResult; + } + | { + readonly type: 'evalError'; + readonly requestId: string; + readonly error: string; + readonly stack?: string; + } + | { + readonly type: 'shutdownComplete'; + readonly requestId: string; + } + | { + readonly type: 'tokenRequest'; + readonly requestId: string; + readonly scopes: readonly string[]; + readonly tenantId?: string; + } + | { + readonly type: 'log'; + readonly level: 'info' | 'warn' | 'error' | 'debug'; + readonly message: string; + }; diff --git a/webpack.config.ext.js b/webpack.config.ext.js index 82a70551..de547515 100644 --- a/webpack.config.ext.js +++ b/webpack.config.ext.js @@ -22,6 +22,7 @@ module.exports = (env, { mode }) => { node: { __filename: false, __dirname: false }, entry: { main: './main.ts', + scratchpadWorker: './src/documentdb/scratchpad/scratchpadWorker.ts', }, output: { path: path.resolve(__dirname, 'dist'), From 088fed010d911cf8ac3af34fae294831756a7fdc Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Wed, 25 Mar 2026 16:16:20 +0000 Subject: [PATCH 02/26] =?UTF-8?q?feat(scratchpad):=20Phase=202=20=E2=80=94?= =?UTF-8?q?=20main=20thread=20integration?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rewrite ScratchpadEvaluator to route all execution through the worker thread: - ScratchpadEvaluator now manages worker lifecycle (spawn, kill, shutdown, dispose) - Worker state machine: idle → spawning → ready → executing → ready (or terminated) - Request/response correlation via requestId UUID map - Timeout enforced via worker.terminate() — actually stops infinite loops - Help command stays in main thread (static text, no @mongosh needed) - Cluster switch detection: kills worker and respawns with new credentials - Entra ID OIDC token requests delegated via IPC to main thread - Worker logging routed to ext.outputChannel - executeScratchpadCode.ts: cancellable progress notification In-process eval path is fully replaced — no feature flag. Step 6.2 WI-1, Phase 2 of 3. --- .../scratchpad/executeScratchpadCode.ts | 11 +- .../scratchpad/ScratchpadEvaluator.ts | 472 +++++++++++++++--- src/documentdb/scratchpad/scratchpadWorker.ts | 35 +- 3 files changed, 407 insertions(+), 111 deletions(-) diff --git a/src/commands/scratchpad/executeScratchpadCode.ts b/src/commands/scratchpad/executeScratchpadCode.ts index 2a305f21..b3bb8186 100644 --- a/src/commands/scratchpad/executeScratchpadCode.ts +++ b/src/commands/scratchpad/executeScratchpadCode.ts @@ -41,10 +41,15 @@ export async function executeScratchpadCode(code: string): Promise { await vscode.window.withProgress( { location: vscode.ProgressLocation.Notification, - title: l10n.t('Running scratchpad query…'), - cancellable: false, + title: l10n.t('Running scratchpad…'), + cancellable: true, }, - async () => { + async (_progress, token) => { + // Cancel kills the worker — user can re-run to respawn + token.onCancellationRequested(() => { + evaluator?.killWorker(); + }); + const startTime = Date.now(); try { const result = await evaluator!.evaluate(connection, code); diff --git a/src/documentdb/scratchpad/ScratchpadEvaluator.ts b/src/documentdb/scratchpad/ScratchpadEvaluator.ts index 012def87..3dab0357 100644 --- a/src/documentdb/scratchpad/ScratchpadEvaluator.ts +++ b/src/documentdb/scratchpad/ScratchpadEvaluator.ts @@ -3,25 +3,50 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ -import * as l10n from '@vscode/l10n'; -import { EventEmitter } from 'events'; -import * as vm from 'vm'; +import { randomUUID } from 'crypto'; +import * as path from 'path'; import * as vscode from 'vscode'; +import { Worker } from 'worker_threads'; import { ext } from '../../extensionVariables'; -import { ClustersClient } from '../ClustersClient'; +import { CredentialCache } from '../CredentialCache'; import { type ExecutionResult, type ScratchpadConnection } from './types'; +import { + type MainToWorkerMessage, + type SerializableExecutionResult, + type SerializableMongoClientOptions, + type WorkerToMainMessage, +} from './workerTypes'; + +/** Worker lifecycle states */ +type WorkerState = 'idle' | 'spawning' | 'ready' | 'executing'; /** - * Evaluates scratchpad code in-process using the `@mongosh` pipeline, - * reusing the existing authenticated `MongoClient` from `ClustersClient`. + * Evaluates scratchpad code in a persistent worker thread. * - * A fresh `ShellInstanceState` + `ShellEvaluator` is created per execution - * to avoid variable leakage between runs. Only the `MongoClient` is reused. + * The worker owns its own `MongoClient` (authenticated via credentials from + * `CredentialCache`) and stays alive between runs. This provides: + * - Infinite loop safety (main thread can kill the worker) + * - MongoClient isolation from the Collection View + * - Zero re-auth overhead after the first run * - * Dependencies (`@mongosh/*`) are lazy-imported to avoid loading ~2-4 MB - * of Babel + async-rewriter at extension activation time. + * The public API is unchanged from the in-process evaluator: + * `evaluate(connection, code) → Promise` */ -export class ScratchpadEvaluator { +export class ScratchpadEvaluator implements vscode.Disposable { + private _worker: Worker | undefined; + private _workerState: WorkerState = 'idle'; + /** Which cluster the live worker is connected to (to detect cluster switches) */ + private _workerClusterId: string | undefined; + + /** Pending request correlation map: requestId → { resolve, reject } */ + private _pendingRequests = new Map< + string, + { + resolve: (value: unknown) => void; + reject: (error: Error) => void; + } + >(); + /** * Evaluate user code against the connected database. * @@ -30,79 +55,377 @@ export class ScratchpadEvaluator { * @returns Formatted execution result with type, printable value, and timing. */ async evaluate(connection: ScratchpadConnection, code: string): Promise { - // Lazy-import @mongosh packages to avoid loading at activation time - const { NodeDriverServiceProvider } = await import('@mongosh/service-provider-node-driver'); - const { ShellInstanceState } = await import('@mongosh/shell-api'); - const { ShellEvaluator } = await import('@mongosh/shell-evaluator'); - - // Reuse the existing authenticated MongoClient - const client = await ClustersClient.getClient(connection.clusterId); - const mongoClient = client.getMongoClient(); - - // Create fresh shell context per execution (no variable leakage) - const bus = new EventEmitter(); - const serviceProvider = new NodeDriverServiceProvider(mongoClient, bus, { - productDocsLink: 'https://github.com/microsoft/vscode-documentdb', - productName: 'DocumentDB for VS Code Scratchpad', - }); - const instanceState = new ShellInstanceState(serviceProvider, bus); - const evaluator = new ShellEvaluator(instanceState); + // Intercept scratchpad-specific commands before they reach the worker + const trimmed = code.trim(); + const helpResult = this.tryHandleHelp(trimmed); + if (helpResult) { + return { ...helpResult, durationMs: 0 }; + } + + // Ensure worker is alive and connected to the right cluster + await this.ensureWorker(connection); + + // Send eval message and await result + const timeoutSec = vscode.workspace.getConfiguration().get(ext.settingsKeys.shellTimeout) ?? 30; + const timeoutMs = timeoutSec * 1000; + + const result = await this.sendEval(connection, code, timeoutMs); + return result; + } - // Set up eval context with shell globals (db, ObjectId, ISODate, etc.) - const context = {}; - instanceState.setCtx(context); + /** + * Gracefully shut down the worker: close MongoClient, then terminate thread. + * Returns after the worker has confirmed shutdown or after a timeout. + */ + async shutdown(): Promise { + if (!this._worker || this._workerState === 'idle') { + return; + } - // Pre-select the target database - await evaluator.customEval( - customEvalFn, - `use(${JSON.stringify(connection.databaseName)})`, - context, - 'scratchpad', - ); + try { + await this.sendRequest({ type: 'shutdown', requestId: '' }, 5000); + } catch { + // Shutdown timed out or failed — force-kill + } - // Execute with timeout - const timeoutMs = (vscode.workspace.getConfiguration().get(ext.settingsKeys.shellTimeout) ?? 30) * 1000; + this.terminateWorker(); + } - const startTime = Date.now(); + /** + * Force-terminate the worker thread immediately. + * Used for infinite loop recovery (timeout) and cancellation. + */ + killWorker(): void { + this.terminateWorker(); + } - // Intercept scratchpad-specific commands before they reach @mongosh - const trimmed = code.trim(); - const helpResult = this.tryHandleHelp(trimmed); - if (helpResult) { - return { ...helpResult, durationMs: Date.now() - startTime }; + dispose(): void { + this.terminateWorker(); + } + + // ─── Private: Worker lifecycle ─────────────────────────────────────────── + + /** + * Ensure a worker is alive and connected to the correct cluster. + * Spawns a new worker if needed (lazy), or kills and respawns if the + * cluster has changed. + */ + private async ensureWorker(connection: ScratchpadConnection): Promise { + // If worker is alive but connected to a different cluster, shut it down + if (this._worker && this._workerClusterId !== connection.clusterId) { + this.terminateWorker(); } - const evalPromise = evaluator.customEval(customEvalFn, code, context, 'scratchpad'); - const timeoutPromise = new Promise((_resolve, reject) => { - setTimeout(() => reject(new Error(l10n.t('Execution timed out'))), timeoutMs); + // If no worker exists, spawn one + if (!this._worker || this._workerState === 'idle') { + await this.spawnWorker(connection); + } + } + + /** + * Spawn a new worker thread and send the init message. + */ + private async spawnWorker(connection: ScratchpadConnection): Promise { + this._workerState = 'spawning'; + + // Resolve worker script path (same directory as the main bundle in dist/) + const workerPath = path.join(__dirname, 'scratchpadWorker.js'); + this._worker = new Worker(workerPath); + this._workerClusterId = connection.clusterId; + + // Listen for messages from the worker + this._worker.on('message', (msg: WorkerToMainMessage) => { + this.handleWorkerMessage(msg); + }); + + // Listen for worker exit (crash or termination) + this._worker.on('exit', (exitCode: number) => { + ext.outputChannel.appendLine(`[Scratchpad Worker] Worker exited with code ${String(exitCode)}`); + this.handleWorkerExit(); + }); + + this._worker.on('error', (error: Error) => { + ext.outputChannel.appendLine(`[Scratchpad Worker] ERROR: ${error.message}`); }); - const result = await Promise.race([evalPromise, timeoutPromise]); - const durationMs = Date.now() - startTime; + // Build init message from cached credentials + const initMsg = this.buildInitMessage(connection); + + // Send init and wait for acknowledgment + await this.sendRequest(initMsg, 30000); + this._workerState = 'ready'; + } + + /** + * Build the init message from CredentialCache data. + */ + private buildInitMessage(connection: ScratchpadConnection): MainToWorkerMessage & { type: 'init' } { + const credentials = CredentialCache.getCredentials(connection.clusterId); + if (!credentials) { + throw new Error(`No credentials found for cluster ${connection.clusterId}`); + } - // customEval already runs toShellResult() internally via resultHandler, - // so `result` is already a ShellResult { type, printable, rawValue, source? }. - const shellResult = result as { - type: string | null; - printable: unknown; - source?: { namespace?: { db: string; collection: string } }; + const authMechanism = credentials.authMechanism ?? 'NativeAuth'; + + // Build connection string + let connectionString: string; + if (authMechanism === 'NativeAuth') { + connectionString = CredentialCache.getConnectionStringWithPassword(connection.clusterId); + } else { + // Entra ID: use connection string without embedded credentials + connectionString = credentials.connectionString; + } + + // Build serializable MongoClientOptions + const clientOptions: SerializableMongoClientOptions = { + serverSelectionTimeoutMS: credentials.emulatorConfiguration?.isEmulator ? 4000 : undefined, + tlsAllowInvalidCertificates: + credentials.emulatorConfiguration?.isEmulator && + credentials.emulatorConfiguration?.disableEmulatorSecurity + ? true + : undefined, }; return { - type: shellResult.type, - printable: shellResult.printable, - durationMs, - source: shellResult.source?.namespace - ? { - namespace: { - db: shellResult.source.namespace.db, - collection: shellResult.source.namespace.collection, - }, - } - : undefined, + type: 'init', + requestId: '', + connectionString, + clientOptions, + databaseName: connection.databaseName, + authMechanism: authMechanism as 'NativeAuth' | 'MicrosoftEntraID', + tenantId: credentials.entraIdConfig?.tenantId, + displayBatchSize: 50, + }; + } + + /** + * Send an eval message to the worker and await the result. + */ + private async sendEval( + connection: ScratchpadConnection, + code: string, + timeoutMs: number, + ): Promise { + this._workerState = 'executing'; + + const evalMsg: MainToWorkerMessage = { + type: 'eval', + requestId: '', + code, + databaseName: connection.databaseName, }; + + try { + const result = await this.sendRequest<{ result: SerializableExecutionResult }>(evalMsg, timeoutMs); + + // Deserialize the result — printable is an EJSON string from the worker + const serResult = result.result; + let printable: unknown; + try { + printable = JSON.parse(serResult.printable) as unknown; + } catch { + printable = serResult.printable; + } + + return { + type: serResult.type, + printable, + durationMs: serResult.durationMs, + source: serResult.source, + }; + } finally { + if (this._workerState === 'executing') { + this._workerState = 'ready'; + } + } + } + + // ─── Private: IPC request/response ─────────────────────────────────────── + + /** + * Send a message to the worker and return a promise that resolves + * when the corresponding response arrives. + */ + private sendRequest(msg: MainToWorkerMessage, timeoutMs: number): Promise { + if (!this._worker) { + return Promise.reject(new Error('Worker is not running')); + } + + const requestId = randomUUID(); + const msgWithId = { ...msg, requestId }; + + return new Promise((resolve, reject) => { + this._pendingRequests.set(requestId, { + resolve: resolve as (value: unknown) => void, + reject, + }); + + // Timeout — kills the worker for safety (infinite loop protection) + const timer = setTimeout(() => { + const pending = this._pendingRequests.get(requestId); + if (pending) { + this._pendingRequests.delete(requestId); + this.killWorker(); + pending.reject( + new Error(`Execution timed out after ${String(Math.round(timeoutMs / 1000))} seconds`), + ); + } + }, timeoutMs); + + // Store timer reference on the pending entry so we can clear it + const entry = this._pendingRequests.get(requestId)!; + const originalResolve = entry.resolve; + const originalReject = entry.reject; + entry.resolve = (value: unknown) => { + clearTimeout(timer); + originalResolve(value); + }; + entry.reject = (error: Error) => { + clearTimeout(timer); + originalReject(error); + }; + + this._worker!.postMessage(msgWithId); + }); + } + + /** + * Handle an incoming message from the worker. + */ + private handleWorkerMessage(msg: WorkerToMainMessage): void { + switch (msg.type) { + case 'initResult': { + const pending = this._pendingRequests.get(msg.requestId); + if (pending) { + this._pendingRequests.delete(msg.requestId); + if (msg.success) { + pending.resolve(undefined); + } else { + pending.reject(new Error(msg.error ?? 'Worker init failed')); + } + } + break; + } + + case 'evalResult': { + const pending = this._pendingRequests.get(msg.requestId); + if (pending) { + this._pendingRequests.delete(msg.requestId); + pending.resolve(msg); + } + break; + } + + case 'evalError': { + const pending = this._pendingRequests.get(msg.requestId); + if (pending) { + this._pendingRequests.delete(msg.requestId); + const error = new Error(msg.error); + if (msg.stack) { + error.stack = msg.stack; + } + pending.reject(error); + } + break; + } + + case 'shutdownComplete': { + const pending = this._pendingRequests.get(msg.requestId); + if (pending) { + this._pendingRequests.delete(msg.requestId); + pending.resolve(undefined); + } + break; + } + + case 'tokenRequest': { + // Entra ID: worker needs an OIDC token — delegate to main thread VS Code API + void this.handleTokenRequest(msg); + break; + } + + case 'log': { + const prefix = '[Scratchpad Worker]'; + switch (msg.level) { + case 'error': + ext.outputChannel.appendLine(`${prefix} ERROR: ${msg.message}`); + break; + case 'warn': + ext.outputChannel.appendLine(`${prefix} WARN: ${msg.message}`); + break; + default: + ext.outputChannel.appendLine(`${prefix} ${msg.message}`); + break; + } + break; + } + } + } + + /** + * Handle a token request from the worker (Entra ID OIDC). + * Calls VS Code's auth API on the main thread and sends the token back. + */ + private async handleTokenRequest(msg: Extract): Promise { + try { + const { getSessionFromVSCode } = await import( + // eslint-disable-next-line import/no-internal-modules + '@microsoft/vscode-azext-azureauth/out/src/getSessionFromVSCode' + ); + const session = await getSessionFromVSCode(msg.scopes as string[], msg.tenantId, { createIfNone: true }); + + if (!session) { + throw new Error('Failed to obtain Entra ID session'); + } + + const response: MainToWorkerMessage = { + type: 'tokenResponse', + requestId: msg.requestId, + accessToken: session.accessToken, + }; + this._worker?.postMessage(response); + } catch (error: unknown) { + const errorMessage = error instanceof Error ? error.message : String(error); + const response: MainToWorkerMessage = { + type: 'tokenError', + requestId: msg.requestId, + error: errorMessage, + }; + this._worker?.postMessage(response); + } + } + + // ─── Private: Worker cleanup ───────────────────────────────────────────── + + private terminateWorker(): void { + if (this._worker) { + void this._worker.terminate(); + this._worker = undefined; + } + this._workerState = 'idle'; + this._workerClusterId = undefined; + + // Reject all pending requests + for (const [, entry] of this._pendingRequests) { + entry.reject(new Error('Worker terminated')); + } + this._pendingRequests.clear(); + } + + private handleWorkerExit(): void { + this._worker = undefined; + this._workerState = 'idle'; + this._workerClusterId = undefined; + + // Reject any still-pending requests + for (const [, entry] of this._pendingRequests) { + entry.reject(new Error('Worker exited unexpectedly')); + } + this._pendingRequests.clear(); } + // ─── Help command interception ─────────────────────────────────────────── + /** * Handle `help` command with scratchpad-specific output. * Returns undefined if the input is not a help command. @@ -178,16 +501,3 @@ export class ScratchpadEvaluator { return { type: 'Help', printable: helpText }; } } - -/** - * The eval function passed to `ShellEvaluator.customEval()`. - * Called by @mongosh with (rewrittenCode, context, filename). - * - * Uses `vm.runInContext` so that all properties on the context object - * (including getters like `db`) are accessible as globals in the code. - */ -// eslint-disable-next-line @typescript-eslint/require-await -async function customEvalFn(code: string, context: object): Promise { - const vmContext = vm.isContext(context) ? context : vm.createContext(context); - return vm.runInContext(code, vmContext) as unknown; -} diff --git a/src/documentdb/scratchpad/scratchpadWorker.ts b/src/documentdb/scratchpad/scratchpadWorker.ts index c86ba2c9..492dd5ee 100644 --- a/src/documentdb/scratchpad/scratchpadWorker.ts +++ b/src/documentdb/scratchpad/scratchpadWorker.ts @@ -15,17 +15,14 @@ */ import { randomUUID } from 'crypto'; +import { type MongoClientOptions, type MongoClient as MongoClientType } from 'mongodb'; import { parentPort } from 'worker_threads'; -import type { MainToWorkerMessage, WorkerToMainMessage } from './workerTypes'; +import { type MainToWorkerMessage, type WorkerToMainMessage } from './workerTypes'; if (!parentPort) { throw new Error('scratchpadWorker.ts must be run as a worker_thread'); } -// ─── Types for lazily-imported @mongosh packages ───────────────────────────── - -type MongoClientType = import('mongodb').MongoClient; - // ─── Worker state ──────────────────────────────────────────────────────────── let mongoClient: MongoClientType | undefined; @@ -92,16 +89,14 @@ parentPort.on('message', (msg: MainToWorkerMessage) => { // ─── Init handler ──────────────────────────────────────────────────────────── -async function handleInit( - msg: Extract, -): Promise { +async function handleInit(msg: Extract): Promise { log('info', `Initializing worker (auth: ${msg.authMechanism}, db: ${msg.databaseName})`); // Lazy-import MongoDB driver const { MongoClient } = await import('mongodb'); // Build MongoClient options from the serializable subset - const options: import('mongodb').MongoClientOptions = { + const options: MongoClientOptions = { ...msg.clientOptions, }; @@ -146,9 +141,7 @@ async function handleInit( // ─── Eval handler ──────────────────────────────────────────────────────────── -async function handleEval( - msg: Extract, -): Promise { +async function handleEval(msg: Extract): Promise { if (!mongoClient) { throw new Error('Worker not initialized — call init first'); } @@ -186,21 +179,11 @@ async function handleEval( // Switch database if different from current if (msg.databaseName !== currentDatabaseName) { - await evaluator.customEval( - customEvalFn, - `use(${JSON.stringify(msg.databaseName)})`, - context, - 'scratchpad', - ); + await evaluator.customEval(customEvalFn, `use(${JSON.stringify(msg.databaseName)})`, context, 'scratchpad'); currentDatabaseName = msg.databaseName; } else { // Pre-select the target database (fresh context each time) - await evaluator.customEval( - customEvalFn, - `use(${JSON.stringify(msg.databaseName)})`, - context, - 'scratchpad', - ); + await evaluator.customEval(customEvalFn, `use(${JSON.stringify(msg.databaseName)})`, context, 'scratchpad'); } // Evaluate user code @@ -252,9 +235,7 @@ async function handleEval( // ─── Shutdown handler ──────────────────────────────────────────────────────── -async function handleShutdown( - msg: Extract, -): Promise { +async function handleShutdown(msg: Extract): Promise { log('info', 'Shutting down worker — closing MongoClient'); try { From c82ed50e0a94c403218dddf600524fad008dad1c Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Wed, 25 Mar 2026 16:18:25 +0000 Subject: [PATCH 03/26] =?UTF-8?q?feat(scratchpad):=20Phase=203=20=E2=80=94?= =?UTF-8?q?=20worker=20disposal=20+=20lifecycle=20wiring?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Export disposeEvaluator() from executeScratchpadCode.ts for clean worker shutdown - Wire evaluator disposal into extension deactivation via ext.context.subscriptions - Worker thread is properly terminated when the extension deactivates Completes the SCRAM auth + kill/respawn wiring (credential passthrough was already implemented in Phase 2's buildInitMessage). Step 6.2 WI-1, Phase 3 of 3. --- src/commands/scratchpad/executeScratchpadCode.ts | 9 +++++++++ src/documentdb/ClustersExtension.ts | 4 ++++ 2 files changed, 13 insertions(+) diff --git a/src/commands/scratchpad/executeScratchpadCode.ts b/src/commands/scratchpad/executeScratchpadCode.ts index b3bb8186..fe91bc24 100644 --- a/src/commands/scratchpad/executeScratchpadCode.ts +++ b/src/commands/scratchpad/executeScratchpadCode.ts @@ -16,6 +16,15 @@ import { type ExecutionResult, type ScratchpadConnection } from '../../documentd /** Shared evaluator instance — lazily created, reused across runs. */ let evaluator: ScratchpadEvaluator | undefined; +/** + * Dispose the shared evaluator instance (kills the worker thread). + * Called during extension deactivation. + */ +export function disposeEvaluator(): void { + evaluator?.dispose(); + evaluator = undefined; +} + /** * Executes scratchpad code and displays the result in a read-only side panel. * Used by both `runAll` and `runSelected` commands. diff --git a/src/documentdb/ClustersExtension.ts b/src/documentdb/ClustersExtension.ts index 64cdc6ba..400e6e85 100644 --- a/src/documentdb/ClustersExtension.ts +++ b/src/documentdb/ClustersExtension.ts @@ -57,6 +57,7 @@ import { retryAuthentication } from '../commands/retryAuthentication/retryAuthen import { revealView } from '../commands/revealView/revealView'; import { clearSchemaCache } from '../commands/schemaStore/clearSchemaCache'; import { connectDatabase } from '../commands/scratchpad/connectDatabase'; +import { disposeEvaluator } from '../commands/scratchpad/executeScratchpadCode'; import { newScratchpad } from '../commands/scratchpad/newScratchpad'; import { runAll } from '../commands/scratchpad/runAll'; import { runSelected } from '../commands/scratchpad/runSelected'; @@ -204,6 +205,9 @@ export class ClustersExtension implements vscode.Disposable { const scratchpadService = ScratchpadService.getInstance(); ext.context.subscriptions.push(scratchpadService); + // Register evaluator disposal for clean worker shutdown on deactivation + ext.context.subscriptions.push({ dispose: disposeEvaluator }); + // Register CodeLens provider for scratchpad files const codeLensProvider = new ScratchpadCodeLensProvider(); ext.context.subscriptions.push(codeLensProvider); From c415bf731a4044667a5077d2efd1f499d916ed2e Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Wed, 25 Mar 2026 16:50:11 +0000 Subject: [PATCH 04/26] =?UTF-8?q?feat(scratchpad):=20Phase=205=20=E2=80=94?= =?UTF-8?q?=20SchemaStore=20doc=20cap=20+=20connection=20sync?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SchemaStore integration: - Cap schema feeding at 100 documents (randomly sampled via Fisher-Yates) - Prevents unbounded IPC/memory usage for large result sets Connection state synchronization: - Shutdown worker when scratchpad connection is cleared (disconnect) - Shutdown worker when the last .documentdb editor tab closes - Worker respawns lazily on next Run Export shutdownEvaluator() for graceful worker cleanup. Step 6.2 WI-2, Phase 5. --- .../scratchpad/executeScratchpadCode.ts | 43 +++++++++++++++++-- src/documentdb/ClustersExtension.ts | 33 +++++++++++++- 2 files changed, 72 insertions(+), 4 deletions(-) diff --git a/src/commands/scratchpad/executeScratchpadCode.ts b/src/commands/scratchpad/executeScratchpadCode.ts index fe91bc24..e0c21d10 100644 --- a/src/commands/scratchpad/executeScratchpadCode.ts +++ b/src/commands/scratchpad/executeScratchpadCode.ts @@ -25,6 +25,15 @@ export function disposeEvaluator(): void { evaluator = undefined; } +/** + * Gracefully shut down the evaluator's worker thread (closes MongoClient). + * Called when the scratchpad connection is cleared or when all scratchpad editors close. + * The worker will be re-spawned lazily on the next Run. + */ +export function shutdownEvaluator(): void { + void evaluator?.shutdown(); +} + /** * Executes scratchpad code and displays the result in a read-only side panel. * Used by both `runAll` and `runSelected` commands. @@ -97,9 +106,17 @@ export async function executeScratchpadCode(code: string): Promise { ); } +/** + * Maximum number of documents to feed to SchemaStore per execution. + * If the result set is larger, a random sample of this size is used. + */ +const SCHEMA_DOC_CAP = 100; + /** * Very conservative schema feeding: only feed 'Cursor' and 'Document' result types. * Extracts documents from the printable result and adds them to SchemaStore. + * + * Caps at {@link SCHEMA_DOC_CAP} documents (randomly sampled if more). */ function feedResultToSchemaStore(result: ExecutionResult, connection: ScratchpadConnection): void { // Only feed known document-producing result types @@ -122,12 +139,32 @@ function feedResultToSchemaStore(result: ExecutionResult, connection: Scratchpad // Filter to actual document objects with _id (not primitives, not nested arrays, // not projection results with _id: 0 which have artificial shapes) - const docs = items.filter( + let docs = items.filter( (d): d is WithId => d !== null && d !== undefined && typeof d === 'object' && !Array.isArray(d) && '_id' in d, ); - if (docs.length > 0) { - SchemaStore.getInstance().addDocuments(connection.clusterId, ns.db, ns.collection, docs); + if (docs.length === 0) { + return; + } + + // Cap at SCHEMA_DOC_CAP documents — randomly sample if more + if (docs.length > SCHEMA_DOC_CAP) { + docs = randomSample(docs, SCHEMA_DOC_CAP); + } + + SchemaStore.getInstance().addDocuments(connection.clusterId, ns.db, ns.collection, docs); +} + +/** + * Fisher–Yates shuffle-based random sample of `count` items from `array`. + * Returns a new array of length `count` with randomly selected items. + */ +function randomSample(array: T[], count: number): T[] { + const copy = [...array]; + for (let i = copy.length - 1; i > 0; i--) { + const j = Math.floor(Math.random() * (i + 1)); + [copy[i], copy[j]] = [copy[j], copy[i]]; } + return copy.slice(0, count); } diff --git a/src/documentdb/ClustersExtension.ts b/src/documentdb/ClustersExtension.ts index 400e6e85..b3f3a485 100644 --- a/src/documentdb/ClustersExtension.ts +++ b/src/documentdb/ClustersExtension.ts @@ -57,7 +57,7 @@ import { retryAuthentication } from '../commands/retryAuthentication/retryAuthen import { revealView } from '../commands/revealView/revealView'; import { clearSchemaCache } from '../commands/schemaStore/clearSchemaCache'; import { connectDatabase } from '../commands/scratchpad/connectDatabase'; -import { disposeEvaluator } from '../commands/scratchpad/executeScratchpadCode'; +import { disposeEvaluator, shutdownEvaluator } from '../commands/scratchpad/executeScratchpadCode'; import { newScratchpad } from '../commands/scratchpad/newScratchpad'; import { runAll } from '../commands/scratchpad/runAll'; import { runSelected } from '../commands/scratchpad/runSelected'; @@ -208,6 +208,37 @@ export class ClustersExtension implements vscode.Disposable { // Register evaluator disposal for clean worker shutdown on deactivation ext.context.subscriptions.push({ dispose: disposeEvaluator }); + // Shut down the scratchpad worker when connection is cleared + ext.context.subscriptions.push( + scratchpadService.onDidChangeState(() => { + if (!scratchpadService.isConnected()) { + shutdownEvaluator(); + } + }), + ); + + // Shut down the scratchpad worker when the last .documentdb editor closes + ext.context.subscriptions.push( + vscode.workspace.onDidCloseTextDocument((doc) => { + if (doc.languageId === SCRATCHPAD_LANGUAGE_ID) { + // Check if any other scratchpad editors remain open + const hasOpenScratchpad = vscode.window.tabGroups.all.some((group) => + group.tabs.some((tab) => { + const input = tab.input; + return ( + input instanceof vscode.TabInputText && + (input.uri.path.endsWith('.documentdb') || + input.uri.path.endsWith('.documentdb.js')) + ); + }), + ); + if (!hasOpenScratchpad) { + shutdownEvaluator(); + } + } + }), + ); + // Register CodeLens provider for scratchpad files const codeLensProvider = new ScratchpadCodeLensProvider(); ext.context.subscriptions.push(codeLensProvider); From bd7c3a8ac13b6242fd3a4ed4036239212d806f9a Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Wed, 25 Mar 2026 17:12:55 +0000 Subject: [PATCH 05/26] fix(scratchpad): use EJSON.parse for BSON type fidelity in schema analysis MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace JSON.parse with EJSON.parse when deserializing worker eval results. This preserves BSON types (ObjectId, Date, Decimal128) so that SchemaAnalyzer correctly identifies field types for autocompletion. With JSON.parse, BSON types became plain objects ({'$oid': '...'}) causing SchemaAnalyzer to create wrong field paths (e.g. '_id.$oid' instead of '_id') and wrong types (object instead of objectid). Benchmarked on 100 documents with ~100 fields each: - JSON.parse: 2.2ms parse, broken types (wrong paths + types) - EJSON.parse: 9.9ms parse, correct types (only Int32/Long→Double) - Both are negligible vs actual query time (100-5000ms) Int32 and Long are still collapsed to Double (JavaScript number) — this is a fundamental EJSON limitation, not a serialization bug. --- src/documentdb/scratchpad/ScratchpadEvaluator.ts | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/src/documentdb/scratchpad/ScratchpadEvaluator.ts b/src/documentdb/scratchpad/ScratchpadEvaluator.ts index 3dab0357..74aafcb6 100644 --- a/src/documentdb/scratchpad/ScratchpadEvaluator.ts +++ b/src/documentdb/scratchpad/ScratchpadEvaluator.ts @@ -218,13 +218,23 @@ export class ScratchpadEvaluator implements vscode.Disposable { try { const result = await this.sendRequest<{ result: SerializableExecutionResult }>(evalMsg, timeoutMs); - // Deserialize the result — printable is an EJSON string from the worker + // Deserialize the result — printable is an EJSON string from the worker. + // Use EJSON.parse to reconstruct BSON types (ObjectId, Date, Decimal128, etc.) + // so that SchemaAnalyzer correctly identifies field types. + // Int32 and Long are still collapsed to Double (JavaScript number) — this is + // a fundamental EJSON limitation, not a bug. const serResult = result.result; let printable: unknown; try { - printable = JSON.parse(serResult.printable) as unknown; + const { EJSON } = await import('bson'); + printable = EJSON.parse(serResult.printable); } catch { - printable = serResult.printable; + // Fallback to JSON.parse if EJSON fails, then raw string + try { + printable = JSON.parse(serResult.printable) as unknown; + } catch { + printable = serResult.printable; + } } return { From 366386087b4c7f89a666404a5a6227549c7fbbf9 Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Wed, 25 Mar 2026 17:19:44 +0000 Subject: [PATCH 06/26] feat(scratchpad): phased progress notifications MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Show distinct progress messages during scratchpad execution: - 'Initializing scratchpad runtime…' — worker thread being created - 'Authenticating with {clusterName}…' — MongoClient connecting + auth - 'Running query…' — user code being evaluated On subsequent runs (worker already alive), only 'Running query…' is shown. The progress notification remains cancellable (Cancel kills worker). Added onProgress callback parameter to ScratchpadEvaluator.evaluate(). --- l10n/bundle.l10n.json | 5 +++- .../scratchpad/executeScratchpadCode.ts | 6 +++-- .../scratchpad/ScratchpadEvaluator.ts | 26 +++++++++++++++---- 3 files changed, 29 insertions(+), 8 deletions(-) diff --git a/l10n/bundle.l10n.json b/l10n/bundle.l10n.json index 9ad22905..4e9bbb9c 100644 --- a/l10n/bundle.l10n.json +++ b/l10n/bundle.l10n.json @@ -185,6 +185,7 @@ "Authenticate to Connect with Your DocumentDB Cluster": "Authenticate to Connect with Your DocumentDB Cluster", "Authenticate using a username and password": "Authenticate using a username and password", "Authenticate using Microsoft Entra ID (Azure AD)": "Authenticate using Microsoft Entra ID (Azure AD)", + "Authenticating with {0}…": "Authenticating with {0}…", "Authentication configuration is missing for \"{cluster}\".": "Authentication configuration is missing for \"{cluster}\".", "Authentication data (primary connection string) is missing for \"{cluster}\".": "Authentication data (primary connection string) is missing for \"{cluster}\".", "Authentication data (properties.connectionString) is missing for \"{cluster}\".": "Authentication data (properties.connectionString) is missing for \"{cluster}\".", @@ -599,6 +600,7 @@ "Indexes": "Indexes", "Info from the webview: ": "Info from the webview: ", "Information was confusing": "Information was confusing", + "Initializing scratchpad runtime…": "Initializing scratchpad runtime…", "Initializing task...": "Initializing task...", "Inserted {0} document(s).": "Inserted {0} document(s).", "Install Azure Account Extension...": "Install Azure Account Extension...", @@ -834,7 +836,8 @@ "Run All": "Run All", "Run the entire file ({0}+Shift+Enter)": "Run the entire file ({0}+Shift+Enter)", "Run this block ({0}+Enter)": "Run this block ({0}+Enter)", - "Running scratchpad query…": "Running scratchpad query…", + "Running query…": "Running query…", + "Running scratchpad…": "Running scratchpad…", "Running…": "Running…", "Save": "Save", "Save credentials for future connections.": "Save credentials for future connections.", diff --git a/src/commands/scratchpad/executeScratchpadCode.ts b/src/commands/scratchpad/executeScratchpadCode.ts index e0c21d10..8427c80b 100644 --- a/src/commands/scratchpad/executeScratchpadCode.ts +++ b/src/commands/scratchpad/executeScratchpadCode.ts @@ -62,7 +62,7 @@ export async function executeScratchpadCode(code: string): Promise { title: l10n.t('Running scratchpad…'), cancellable: true, }, - async (_progress, token) => { + async (progress, token) => { // Cancel kills the worker — user can re-run to respawn token.onCancellationRequested(() => { evaluator?.killWorker(); @@ -70,7 +70,9 @@ export async function executeScratchpadCode(code: string): Promise { const startTime = Date.now(); try { - const result = await evaluator!.evaluate(connection, code); + const result = await evaluator!.evaluate(connection, code, (message) => { + progress.report({ message }); + }); const formattedOutput = formatResult(result, code, connection); // Feed document results to SchemaStore for cross-tab schema sharing diff --git a/src/documentdb/scratchpad/ScratchpadEvaluator.ts b/src/documentdb/scratchpad/ScratchpadEvaluator.ts index 74aafcb6..a0e7d7fb 100644 --- a/src/documentdb/scratchpad/ScratchpadEvaluator.ts +++ b/src/documentdb/scratchpad/ScratchpadEvaluator.ts @@ -3,6 +3,7 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ +import * as l10n from '@vscode/l10n'; import { randomUUID } from 'crypto'; import * as path from 'path'; import * as vscode from 'vscode'; @@ -52,9 +53,14 @@ export class ScratchpadEvaluator implements vscode.Disposable { * * @param connection - Active scratchpad connection (clusterId + databaseName). * @param code - JavaScript code string to evaluate. + * @param onProgress - Optional callback for phased progress reporting. * @returns Formatted execution result with type, printable value, and timing. */ - async evaluate(connection: ScratchpadConnection, code: string): Promise { + async evaluate( + connection: ScratchpadConnection, + code: string, + onProgress?: (message: string) => void, + ): Promise { // Intercept scratchpad-specific commands before they reach the worker const trimmed = code.trim(); const helpResult = this.tryHandleHelp(trimmed); @@ -63,9 +69,15 @@ export class ScratchpadEvaluator implements vscode.Disposable { } // Ensure worker is alive and connected to the right cluster - await this.ensureWorker(connection); + const needsSpawn = + !this._worker || this._workerState === 'idle' || this._workerClusterId !== connection.clusterId; + if (needsSpawn) { + onProgress?.(l10n.t('Initializing scratchpad runtime…')); + } + await this.ensureWorker(connection, onProgress); // Send eval message and await result + onProgress?.(l10n.t('Running query…')); const timeoutSec = vscode.workspace.getConfiguration().get(ext.settingsKeys.shellTimeout) ?? 30; const timeoutMs = timeoutSec * 1000; @@ -110,7 +122,10 @@ export class ScratchpadEvaluator implements vscode.Disposable { * Spawns a new worker if needed (lazy), or kills and respawns if the * cluster has changed. */ - private async ensureWorker(connection: ScratchpadConnection): Promise { + private async ensureWorker( + connection: ScratchpadConnection, + onProgress?: (message: string) => void, + ): Promise { // If worker is alive but connected to a different cluster, shut it down if (this._worker && this._workerClusterId !== connection.clusterId) { this.terminateWorker(); @@ -118,14 +133,14 @@ export class ScratchpadEvaluator implements vscode.Disposable { // If no worker exists, spawn one if (!this._worker || this._workerState === 'idle') { - await this.spawnWorker(connection); + await this.spawnWorker(connection, onProgress); } } /** * Spawn a new worker thread and send the init message. */ - private async spawnWorker(connection: ScratchpadConnection): Promise { + private async spawnWorker(connection: ScratchpadConnection, onProgress?: (message: string) => void): Promise { this._workerState = 'spawning'; // Resolve worker script path (same directory as the main bundle in dist/) @@ -152,6 +167,7 @@ export class ScratchpadEvaluator implements vscode.Disposable { const initMsg = this.buildInitMessage(connection); // Send init and wait for acknowledgment + onProgress?.(l10n.t('Authenticating with {0}…', connection.clusterDisplayName)); await this.sendRequest(initMsg, 30000); this._workerState = 'ready'; } From e88d39073664a02f5614eb903443b6c568ce9212 Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Wed, 25 Mar 2026 19:12:52 +0000 Subject: [PATCH 07/26] fix(scratchpad): proper log levels, progress UX, cancel handling MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Log levels: - Worker init/shutdown → debug (lifecycle, not user-visible) - Eval start/end → trace (verbose diagnostic) - Errors/uncaught exceptions → error - MongoClient close failure → warn - Worker exit, connection clear, editors close → debug - Route worker IPC log messages to matching LogOutputChannel methods (trace/debug/info/warn/error) instead of appendLine for all Progress UX: - Title changed to 'DocumentDB Scratchpad' (static bold prefix) - Phase messages show as: 'Initializing…', 'Authenticating…', 'Running query…' - Removed cluster name from authenticating phase (redundant in context) Cancel handling: - Suppress error notification when user explicitly cancels execution - Track cancelled state to avoid showing 'Worker terminated' error panel --- .../scratchpad/executeScratchpadCode.ts | 11 ++++++++++- src/documentdb/ClustersExtension.ts | 2 ++ .../scratchpad/ScratchpadEvaluator.ts | 17 ++++++++++------- src/documentdb/scratchpad/scratchpadWorker.ts | 12 ++++++------ src/documentdb/scratchpad/workerTypes.ts | 2 +- 5 files changed, 29 insertions(+), 15 deletions(-) diff --git a/src/commands/scratchpad/executeScratchpadCode.ts b/src/commands/scratchpad/executeScratchpadCode.ts index 8427c80b..024e55bf 100644 --- a/src/commands/scratchpad/executeScratchpadCode.ts +++ b/src/commands/scratchpad/executeScratchpadCode.ts @@ -59,12 +59,16 @@ export async function executeScratchpadCode(code: string): Promise { await vscode.window.withProgress( { location: vscode.ProgressLocation.Notification, - title: l10n.t('Running scratchpad…'), + title: l10n.t('DocumentDB Scratchpad'), cancellable: true, }, async (progress, token) => { + // Track whether user cancelled — to suppress the error notification + let cancelled = false; + // Cancel kills the worker — user can re-run to respawn token.onCancellationRequested(() => { + cancelled = true; evaluator?.killWorker(); }); @@ -87,6 +91,11 @@ export async function executeScratchpadCode(code: string): Promise { { viewColumn: vscode.ViewColumn.Beside, preserveFocus: true }, ); } catch (error: unknown) { + // Don't show error UI when the user explicitly cancelled + if (cancelled) { + return; + } + const durationMs = Date.now() - startTime; const errorMessage = error instanceof Error ? error.message : String(error); const formattedOutput = formatError(error, code, durationMs, connection); diff --git a/src/documentdb/ClustersExtension.ts b/src/documentdb/ClustersExtension.ts index b3f3a485..0dcd8247 100644 --- a/src/documentdb/ClustersExtension.ts +++ b/src/documentdb/ClustersExtension.ts @@ -212,6 +212,7 @@ export class ClustersExtension implements vscode.Disposable { ext.context.subscriptions.push( scratchpadService.onDidChangeState(() => { if (!scratchpadService.isConnected()) { + ext.outputChannel.debug('[Scratchpad] Connection cleared — shutting down worker'); shutdownEvaluator(); } }), @@ -233,6 +234,7 @@ export class ClustersExtension implements vscode.Disposable { }), ); if (!hasOpenScratchpad) { + ext.outputChannel.debug('[Scratchpad] All editors closed — shutting down worker'); shutdownEvaluator(); } } diff --git a/src/documentdb/scratchpad/ScratchpadEvaluator.ts b/src/documentdb/scratchpad/ScratchpadEvaluator.ts index a0e7d7fb..5b2070c3 100644 --- a/src/documentdb/scratchpad/ScratchpadEvaluator.ts +++ b/src/documentdb/scratchpad/ScratchpadEvaluator.ts @@ -72,7 +72,7 @@ export class ScratchpadEvaluator implements vscode.Disposable { const needsSpawn = !this._worker || this._workerState === 'idle' || this._workerClusterId !== connection.clusterId; if (needsSpawn) { - onProgress?.(l10n.t('Initializing scratchpad runtime…')); + onProgress?.(l10n.t('Initializing…')); } await this.ensureWorker(connection, onProgress); @@ -155,19 +155,19 @@ export class ScratchpadEvaluator implements vscode.Disposable { // Listen for worker exit (crash or termination) this._worker.on('exit', (exitCode: number) => { - ext.outputChannel.appendLine(`[Scratchpad Worker] Worker exited with code ${String(exitCode)}`); + ext.outputChannel.debug(`[Scratchpad Worker] Worker exited with code ${String(exitCode)}`); this.handleWorkerExit(); }); this._worker.on('error', (error: Error) => { - ext.outputChannel.appendLine(`[Scratchpad Worker] ERROR: ${error.message}`); + ext.outputChannel.error(`[Scratchpad Worker] ${error.message}`); }); // Build init message from cached credentials const initMsg = this.buildInitMessage(connection); // Send init and wait for acknowledgment - onProgress?.(l10n.t('Authenticating with {0}…', connection.clusterDisplayName)); + onProgress?.(l10n.t('Authenticating…')); await this.sendRequest(initMsg, 30000); this._workerState = 'ready'; } @@ -374,13 +374,16 @@ export class ScratchpadEvaluator implements vscode.Disposable { const prefix = '[Scratchpad Worker]'; switch (msg.level) { case 'error': - ext.outputChannel.appendLine(`${prefix} ERROR: ${msg.message}`); + ext.outputChannel.error(`${prefix} ${msg.message}`); break; case 'warn': - ext.outputChannel.appendLine(`${prefix} WARN: ${msg.message}`); + ext.outputChannel.warn(`${prefix} ${msg.message}`); + break; + case 'debug': + ext.outputChannel.debug(`${prefix} ${msg.message}`); break; default: - ext.outputChannel.appendLine(`${prefix} ${msg.message}`); + ext.outputChannel.trace(`${prefix} ${msg.message}`); break; } break; diff --git a/src/documentdb/scratchpad/scratchpadWorker.ts b/src/documentdb/scratchpad/scratchpadWorker.ts index 492dd5ee..60337e07 100644 --- a/src/documentdb/scratchpad/scratchpadWorker.ts +++ b/src/documentdb/scratchpad/scratchpadWorker.ts @@ -37,7 +37,7 @@ const pendingTokenRequests = new Map void; // ─── Logging helper ────────────────────────────────────────────────────────── -function log(level: 'info' | 'warn' | 'error' | 'debug', message: string): void { +function log(level: 'trace' | 'debug' | 'info' | 'warn' | 'error', message: string): void { const msg: WorkerToMainMessage = { type: 'log', level, message }; parentPort!.postMessage(msg); } @@ -90,7 +90,7 @@ parentPort.on('message', (msg: MainToWorkerMessage) => { // ─── Init handler ──────────────────────────────────────────────────────────── async function handleInit(msg: Extract): Promise { - log('info', `Initializing worker (auth: ${msg.authMechanism}, db: ${msg.databaseName})`); + log('debug', `Initializing worker (auth: ${msg.authMechanism}, db: ${msg.databaseName})`); // Lazy-import MongoDB driver const { MongoClient } = await import('mongodb'); @@ -129,7 +129,7 @@ async function handleInit(msg: Extract): await mongoClient.connect(); currentDatabaseName = msg.databaseName; - log('info', 'Worker initialized — MongoClient connected'); + log('debug', 'Worker initialized — MongoClient connected'); const response: WorkerToMainMessage = { type: 'initResult', @@ -146,7 +146,7 @@ async function handleEval(msg: Extract): throw new Error('Worker not initialized — call init first'); } - log('debug', `Evaluating code (${msg.code.length} chars, db: ${msg.databaseName})`); + log('trace', `Evaluating code (${msg.code.length} chars, db: ${msg.databaseName})`); // Lazy-import @mongosh packages const { EventEmitter } = await import('events'); @@ -211,7 +211,7 @@ async function handleEval(msg: Extract): } } - log('debug', `Evaluation complete (${durationMs}ms, type: ${shellResult.type ?? 'null'})`); + log('trace', `Evaluation complete (${durationMs}ms, type: ${shellResult.type ?? 'null'})`); const response: WorkerToMainMessage = { type: 'evalResult', @@ -236,7 +236,7 @@ async function handleEval(msg: Extract): // ─── Shutdown handler ──────────────────────────────────────────────────────── async function handleShutdown(msg: Extract): Promise { - log('info', 'Shutting down worker — closing MongoClient'); + log('debug', 'Shutting down worker — closing MongoClient'); try { if (mongoClient) { diff --git a/src/documentdb/scratchpad/workerTypes.ts b/src/documentdb/scratchpad/workerTypes.ts index 2004c462..7018fc9d 100644 --- a/src/documentdb/scratchpad/workerTypes.ts +++ b/src/documentdb/scratchpad/workerTypes.ts @@ -119,6 +119,6 @@ export type WorkerToMainMessage = } | { readonly type: 'log'; - readonly level: 'info' | 'warn' | 'error' | 'debug'; + readonly level: 'trace' | 'debug' | 'info' | 'warn' | 'error'; readonly message: string; }; From c95c106063c78b2d49b1d8d26e965a9203623b32 Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Wed, 25 Mar 2026 19:23:12 +0000 Subject: [PATCH 08/26] fix(scratchpad): fix editor close detection + add schema stats command Fix regression: worker not shutting down when scratchpad editors close. - Switched from onDidCloseTextDocument to tabGroups.onDidChangeTabs - onDidCloseTextDocument fires before tab state updates (race condition) - onDidChangeTabs fires after tabs are removed, state is consistent Add 'Show Schema Store Stats' diagnostics command: - Shows collection count, document count, field count in output channel - Per-collection breakdown with key, doc count, field count - Available via Command Palette: 'DocumentDB: Show Schema Store Stats' --- package.json | 6 ++ .../schemaStore/showSchemaStoreStats.ts | 35 ++++++++++++ src/documentdb/ClustersExtension.ts | 55 +++++++++++++------ 3 files changed, 80 insertions(+), 16 deletions(-) create mode 100644 src/commands/schemaStore/showSchemaStoreStats.ts diff --git a/package.json b/package.json index e491869b..fcbda2d5 100644 --- a/package.json +++ b/package.json @@ -551,6 +551,12 @@ "category": "DocumentDB", "command": "vscode-documentdb.command.clearSchemaCache", "title": "Clear Schema Cache" + }, + { + "//": "[Diagnostics] Show Schema Store Stats", + "category": "DocumentDB", + "command": "vscode-documentdb.command.showSchemaStoreStats", + "title": "Show Schema Store Stats" } ], "submenus": [ diff --git a/src/commands/schemaStore/showSchemaStoreStats.ts b/src/commands/schemaStore/showSchemaStoreStats.ts new file mode 100644 index 00000000..0fc41df3 --- /dev/null +++ b/src/commands/schemaStore/showSchemaStoreStats.ts @@ -0,0 +1,35 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { type IActionContext } from '@microsoft/vscode-azext-utils'; +import { SchemaStore } from '../../documentdb/SchemaStore'; +import { ext } from '../../extensionVariables'; + +/** + * Command handler: Show SchemaStore statistics in the output channel. + * Displays collection count, document count, field count, and per-collection breakdown. + */ +export function showSchemaStoreStats(_context: IActionContext): void { + const store = SchemaStore.getInstance(); + const stats = store.getStats(); + + ext.outputChannel.appendLog( + `[SchemaStore] Stats: ${String(stats.collectionCount)} collections, ` + + `${String(stats.totalDocuments)} documents analyzed, ` + + `${String(stats.totalFields)} fields discovered`, + ); + + if (stats.collections.length > 0) { + for (const c of stats.collections) { + ext.outputChannel.appendLog( + `[SchemaStore] ${c.key}: ${String(c.documentCount)} docs, ${String(c.fieldCount)} fields`, + ); + } + } else { + ext.outputChannel.appendLog('[SchemaStore] (empty — no schemas cached)'); + } + + ext.outputChannel.show(); +} diff --git a/src/documentdb/ClustersExtension.ts b/src/documentdb/ClustersExtension.ts index 0dcd8247..f6554267 100644 --- a/src/documentdb/ClustersExtension.ts +++ b/src/documentdb/ClustersExtension.ts @@ -56,6 +56,7 @@ import { removeDiscoveryRegistry } from '../commands/removeDiscoveryRegistry/rem import { retryAuthentication } from '../commands/retryAuthentication/retryAuthentication'; import { revealView } from '../commands/revealView/revealView'; import { clearSchemaCache } from '../commands/schemaStore/clearSchemaCache'; +import { showSchemaStoreStats } from '../commands/schemaStore/showSchemaStoreStats'; import { connectDatabase } from '../commands/scratchpad/connectDatabase'; import { disposeEvaluator, shutdownEvaluator } from '../commands/scratchpad/executeScratchpadCode'; import { newScratchpad } from '../commands/scratchpad/newScratchpad'; @@ -220,23 +221,40 @@ export class ClustersExtension implements vscode.Disposable { // Shut down the scratchpad worker when the last .documentdb editor closes ext.context.subscriptions.push( - vscode.workspace.onDidCloseTextDocument((doc) => { - if (doc.languageId === SCRATCHPAD_LANGUAGE_ID) { - // Check if any other scratchpad editors remain open - const hasOpenScratchpad = vscode.window.tabGroups.all.some((group) => - group.tabs.some((tab) => { - const input = tab.input; - return ( - input instanceof vscode.TabInputText && - (input.uri.path.endsWith('.documentdb') || - input.uri.path.endsWith('.documentdb.js')) - ); - }), + vscode.window.tabGroups.onDidChangeTabs((event) => { + // Only react when tabs are closed + if (event.closed.length === 0) { + return; + } + + // Check if any closed tab was a scratchpad + const closedScratchpad = event.closed.some((tab) => { + const input = tab.input; + return ( + input instanceof vscode.TabInputText && + (input.uri.path.endsWith('.documentdb') || input.uri.path.endsWith('.documentdb.js')) ); - if (!hasOpenScratchpad) { - ext.outputChannel.debug('[Scratchpad] All editors closed — shutting down worker'); - shutdownEvaluator(); - } + }); + + if (!closedScratchpad) { + return; + } + + // Check if any scratchpad tabs remain open + const hasOpenScratchpad = vscode.window.tabGroups.all.some((group) => + group.tabs.some((tab) => { + const input = tab.input; + return ( + input instanceof vscode.TabInputText && + (input.uri.path.endsWith('.documentdb') || + input.uri.path.endsWith('.documentdb.js')) + ); + }), + ); + + if (!hasOpenScratchpad) { + ext.outputChannel.debug('[Scratchpad] All editors closed — shutting down worker'); + shutdownEvaluator(); } }), ); @@ -270,6 +288,11 @@ export class ClustersExtension implements vscode.Disposable { registerCommand('vscode-documentdb.command.clearSchemaCache', withCommandCorrelation(clearSchemaCache)); + registerCommand( + 'vscode-documentdb.command.showSchemaStoreStats', + withCommandCorrelation(showSchemaStoreStats), + ); + //// General Commands: registerCommandWithTreeNodeUnwrapping( From a1fe8923794a41c99d340aa5eddbf52945cf1d36 Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Wed, 25 Mar 2026 19:33:22 +0000 Subject: [PATCH 09/26] fix(scratchpad): normalize CursorIterationResult to plain Array before EJSON serialization CursorIterationResult from @mongosh is an Array subclass with extra properties (cursorHasMore, documents). EJSON.serialize treats it as a plain object and includes those properties, producing: { cursorHasMore: true, documents: [...] } instead of just: [doc1, doc2, ...] This caused: 1. Output showed cursor wrapper object instead of document array 2. resultFormatter showed 'Result: Cursor' instead of 'N documents returned' 3. SchemaStore never received documents (no _id at top level of wrapper object) Fix: Array.from(shellResult.printable) before EJSON.stringify normalizes Array subclasses to plain Arrays, preserving correct serialization. --- src/documentdb/scratchpad/scratchpadWorker.ts | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/src/documentdb/scratchpad/scratchpadWorker.ts b/src/documentdb/scratchpad/scratchpadWorker.ts index 60337e07..79e29d33 100644 --- a/src/documentdb/scratchpad/scratchpadWorker.ts +++ b/src/documentdb/scratchpad/scratchpadWorker.ts @@ -197,17 +197,25 @@ async function handleEval(msg: Extract): source?: { namespace?: { db: string; collection: string } }; }; - // Serialize the printable value via EJSON for safe transfer + // Serialize the printable value via EJSON for safe transfer. + // CursorIterationResult is an Array subclass with extra properties (cursorHasMore, + // documents). EJSON.serialize treats it as a plain object and includes those + // properties, producing { cursorHasMore: true, documents: [...] } instead of [...]. + // Convert to a plain Array first to preserve correct array serialization. + const printableValue = Array.isArray(shellResult.printable) + ? Array.from(shellResult.printable as unknown[]) + : shellResult.printable; + let printableStr: string; try { const { EJSON } = await import('bson'); - printableStr = EJSON.stringify(shellResult.printable, { relaxed: true }, 2); + printableStr = EJSON.stringify(printableValue, { relaxed: true }, 2); } catch { // Fallback: try JSON, then plain string try { - printableStr = JSON.stringify(shellResult.printable, null, 2); + printableStr = JSON.stringify(printableValue, null, 2); } catch { - printableStr = String(shellResult.printable); + printableStr = String(printableValue); } } From 4909cb26877cab68acd094ed8003b0b637bfe230 Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Wed, 25 Mar 2026 20:27:44 +0000 Subject: [PATCH 10/26] fix(scratchpad): unwrap CursorIterationResult { documents } wrapper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit @mongosh's CursorIterationResult extends ShellApiValueClass (not Array). Its asPrintable() returns { ...this } which produces: { cursorHasMore: true, documents: [doc1, doc2, ...] } This wrapper object was passed through as-is, causing: 1. Output showed { cursorHasMore, documents: [...] } instead of just [...] 2. Header showed 'Result: Cursor' instead of 'N documents returned' 3. SchemaStore never received documents (wrapper has no _id field) Fix: Add unwrapCursorResult() helper that extracts the documents array from the { documents: [...] } wrapper. Applied in both: - resultFormatter.ts — for display formatting and document count - executeScratchpadCode.ts — for SchemaStore feeding --- .../scratchpad/executeScratchpadCode.ts | 17 +++++++++-- src/documentdb/scratchpad/resultFormatter.ts | 29 +++++++++++++++++-- 2 files changed, 41 insertions(+), 5 deletions(-) diff --git a/src/commands/scratchpad/executeScratchpadCode.ts b/src/commands/scratchpad/executeScratchpadCode.ts index 024e55bf..2e53d0bf 100644 --- a/src/commands/scratchpad/executeScratchpadCode.ts +++ b/src/commands/scratchpad/executeScratchpadCode.ts @@ -145,8 +145,21 @@ function feedResultToSchemaStore(result: ExecutionResult, connection: Scratchpad return; } - // Normalize to array - const items: unknown[] = Array.isArray(printable) ? printable : [printable]; + // CursorIterationResult from @mongosh wraps documents in { cursorHasMore, documents }. + // Unwrap to get the actual document array. + let items: unknown[]; + if ( + typeof printable === 'object' && + !Array.isArray(printable) && + 'documents' in printable && + Array.isArray((printable as { documents: unknown }).documents) + ) { + items = (printable as { documents: unknown[] }).documents; + } else if (Array.isArray(printable)) { + items = printable; + } else { + items = [printable]; + } // Filter to actual document objects with _id (not primitives, not nested arrays, // not projection results with _id: 0 which have artificial shapes) diff --git a/src/documentdb/scratchpad/resultFormatter.ts b/src/documentdb/scratchpad/resultFormatter.ts index 7d270e98..3e69237d 100644 --- a/src/documentdb/scratchpad/resultFormatter.ts +++ b/src/documentdb/scratchpad/resultFormatter.ts @@ -31,7 +31,9 @@ export function formatResult(result: ExecutionResult, code: string, connection: // Result metadata if (result.type) { - const printable = result.printable; + // CursorIterationResult from @mongosh wraps documents in { cursorHasMore, documents }. + // Unwrap to get the actual document array for display and counting. + const printable = unwrapCursorResult(result.printable); if (Array.isArray(printable)) { lines.push(`// ${l10n.t('{0} documents returned', printable.length)}`); } else { @@ -43,8 +45,8 @@ export function formatResult(result: ExecutionResult, code: string, connection: lines.push('// ─────────────────────────'); lines.push(''); - // Result value - lines.push(formatPrintable(result.printable)); + // Result value — unwrap cursor wrapper for clean output + lines.push(formatPrintable(unwrapCursorResult(result.printable))); return lines.join('\n'); } @@ -81,6 +83,27 @@ export function formatError( return lines.join('\n'); } +/** + * Unwrap CursorIterationResult from @mongosh. + * + * @mongosh's `asPrintable()` on CursorIterationResult produces `{ cursorHasMore, documents }` + * instead of a plain array. This function extracts the `documents` array for clean display + * and schema feeding. + */ +function unwrapCursorResult(printable: unknown): unknown { + if ( + printable !== null && + printable !== undefined && + typeof printable === 'object' && + !Array.isArray(printable) && + 'documents' in printable && + Array.isArray((printable as { documents: unknown }).documents) + ) { + return (printable as { documents: unknown[] }).documents; + } + return printable; +} + function formatPrintable(printable: unknown): string { if (printable === undefined) { return 'undefined'; From 6a825a66cb9bdc6e6d1ce003008feba6e741d9a1 Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Wed, 25 Mar 2026 20:32:31 +0000 Subject: [PATCH 11/26] fix(scratchpad): show authoritative result type in output header MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Use @mongosh's result type instead of guessing from array shape: - Cursor results: 'Result: Cursor (20 documents)' — type + batch count - Other typed results: 'Result: Document', 'Result: string', etc. - No type: no header line (plain JS values) Previous behavior tried to detect 'documents' by checking Array.isArray which was fragile and didn't communicate what kind of result it was. --- l10n/bundle.l10n.json | 8 ++++---- src/documentdb/scratchpad/resultFormatter.ts | 11 ++++++----- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/l10n/bundle.l10n.json b/l10n/bundle.l10n.json index 4e9bbb9c..279b80cf 100644 --- a/l10n/bundle.l10n.json +++ b/l10n/bundle.l10n.json @@ -94,7 +94,6 @@ "{0} completed successfully": "{0} completed successfully", "{0} connections": "{0} connections", "{0} created": "{0} created", - "{0} documents returned": "{0} documents returned", "{0} failed: {1}": "{0} failed: {1}", "{0} file(s) were ignored because they do not match the \"*.json\" pattern.": "{0} file(s) were ignored because they do not match the \"*.json\" pattern.", "{0} inserted": "{0} inserted", @@ -185,7 +184,7 @@ "Authenticate to Connect with Your DocumentDB Cluster": "Authenticate to Connect with Your DocumentDB Cluster", "Authenticate using a username and password": "Authenticate using a username and password", "Authenticate using Microsoft Entra ID (Azure AD)": "Authenticate using Microsoft Entra ID (Azure AD)", - "Authenticating with {0}…": "Authenticating with {0}…", + "Authenticating…": "Authenticating…", "Authentication configuration is missing for \"{cluster}\".": "Authentication configuration is missing for \"{cluster}\".", "Authentication data (primary connection string) is missing for \"{cluster}\".": "Authentication data (primary connection string) is missing for \"{cluster}\".", "Authentication data (properties.connectionString) is missing for \"{cluster}\".": "Authentication data (properties.connectionString) is missing for \"{cluster}\".", @@ -368,6 +367,7 @@ "DocumentDB for VS Code is not signed in to Azure": "DocumentDB for VS Code is not signed in to Azure", "DocumentDB Local": "DocumentDB Local", "DocumentDB Performance Tips": "DocumentDB Performance Tips", + "DocumentDB Scratchpad": "DocumentDB Scratchpad", "DocumentDB Scratchpad connected to {0}": "DocumentDB Scratchpad connected to {0}", "DocumentDB Scratchpad connected to {0}/{1}": "DocumentDB Scratchpad connected to {0}/{1}", "Documents": "Documents", @@ -600,8 +600,8 @@ "Indexes": "Indexes", "Info from the webview: ": "Info from the webview: ", "Information was confusing": "Information was confusing", - "Initializing scratchpad runtime…": "Initializing scratchpad runtime…", "Initializing task...": "Initializing task...", + "Initializing…": "Initializing…", "Inserted {0} document(s).": "Inserted {0} document(s).", "Install Azure Account Extension...": "Install Azure Account Extension...", "Internal error: connectionString must be defined.": "Internal error: connectionString must be defined.", @@ -825,6 +825,7 @@ "Report an issue": "Report an issue", "Resource group \"{0}\" already exists in subscription \"{1}\".": "Resource group \"{0}\" already exists in subscription \"{1}\".", "Result: {0}": "Result: {0}", + "Result: Cursor ({0} documents)": "Result: Cursor ({0} documents)", "Results found": "Results found", "Retry": "Retry", "Reusing active connection for \"{cluster}\".": "Reusing active connection for \"{cluster}\".", @@ -837,7 +838,6 @@ "Run the entire file ({0}+Shift+Enter)": "Run the entire file ({0}+Shift+Enter)", "Run this block ({0}+Enter)": "Run this block ({0}+Enter)", "Running query…": "Running query…", - "Running scratchpad…": "Running scratchpad…", "Running…": "Running…", "Save": "Save", "Save credentials for future connections.": "Save credentials for future connections.", diff --git a/src/documentdb/scratchpad/resultFormatter.ts b/src/documentdb/scratchpad/resultFormatter.ts index 3e69237d..9a2501fb 100644 --- a/src/documentdb/scratchpad/resultFormatter.ts +++ b/src/documentdb/scratchpad/resultFormatter.ts @@ -30,13 +30,14 @@ export function formatResult(result: ExecutionResult, code: string, connection: } // Result metadata + // Result metadata — state what we know from @mongosh, don't guess if (result.type) { - // CursorIterationResult from @mongosh wraps documents in { cursorHasMore, documents }. - // Unwrap to get the actual document array for display and counting. - const printable = unwrapCursorResult(result.printable); - if (Array.isArray(printable)) { - lines.push(`// ${l10n.t('{0} documents returned', printable.length)}`); + const unwrapped = unwrapCursorResult(result.printable); + if (result.type === 'Cursor' && Array.isArray(unwrapped)) { + // Cursor with a known batch: "Result: Cursor (20 documents)" + lines.push(`// ${l10n.t('Result: Cursor ({0} documents)', unwrapped.length)}`); } else { + // Any other typed result: "Result: Document", "Result: string", etc. lines.push(`// ${l10n.t('Result: {0}', result.type)}`); } } From 2f4d90accfc3c596fe469f47bad069bcfa0882fe Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Wed, 25 Mar 2026 20:38:23 +0000 Subject: [PATCH 12/26] fix(scratchpad): show result metadata for untyped results (.toArray, count) - .toArray() returns type=null with an Array: show 'N results' - .count() returns type=null with a number: no special header (value shown) - Cursor: 'Result: Cursor (N documents)' (unchanged) - Typed results: 'Result: Document', etc. (unchanged) --- l10n/bundle.l10n.json | 1 + src/documentdb/scratchpad/resultFormatter.ts | 19 ++++++++++--------- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/l10n/bundle.l10n.json b/l10n/bundle.l10n.json index 279b80cf..0b8a78ce 100644 --- a/l10n/bundle.l10n.json +++ b/l10n/bundle.l10n.json @@ -100,6 +100,7 @@ "{0} item(s) already exist in the destination. Check the Output panel for details.": "{0} item(s) already exist in the destination. Check the Output panel for details.", "{0} processed": "{0} processed", "{0} replaced": "{0} replaced", + "{0} results": "{0} results", "{0} skipped": "{0} skipped", "{0} subfolders": "{0} subfolders", "{0} task(s) are using connections being moved. Check the Output panel for details.": "{0} task(s) are using connections being moved. Check the Output panel for details.", diff --git a/src/documentdb/scratchpad/resultFormatter.ts b/src/documentdb/scratchpad/resultFormatter.ts index 9a2501fb..ef1b95ee 100644 --- a/src/documentdb/scratchpad/resultFormatter.ts +++ b/src/documentdb/scratchpad/resultFormatter.ts @@ -31,15 +31,16 @@ export function formatResult(result: ExecutionResult, code: string, connection: // Result metadata // Result metadata — state what we know from @mongosh, don't guess - if (result.type) { - const unwrapped = unwrapCursorResult(result.printable); - if (result.type === 'Cursor' && Array.isArray(unwrapped)) { - // Cursor with a known batch: "Result: Cursor (20 documents)" - lines.push(`// ${l10n.t('Result: Cursor ({0} documents)', unwrapped.length)}`); - } else { - // Any other typed result: "Result: Document", "Result: string", etc. - lines.push(`// ${l10n.t('Result: {0}', result.type)}`); - } + const unwrapped = unwrapCursorResult(result.printable); + if (result.type === 'Cursor' && Array.isArray(unwrapped)) { + // Cursor with a known batch: "Result: Cursor (20 documents)" + lines.push(`// ${l10n.t('Result: Cursor ({0} documents)', unwrapped.length)}`); + } else if (result.type) { + // Typed result: "Result: Document", "Result: string", etc. + lines.push(`// ${l10n.t('Result: {0}', result.type)}`); + } else if (Array.isArray(unwrapped)) { + // Untyped array (e.g. .toArray()) — show count + lines.push(`// ${l10n.t('{0} results', unwrapped.length)}`); } lines.push(`// ${l10n.t('Executed in {0}ms', result.durationMs)}`); From 479257b64d389f31188fef2f6c7266a491214518 Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Thu, 26 Mar 2026 07:45:54 +0000 Subject: [PATCH 13/26] fix(scratchpad): polish output headers, logging, schema stats display 1. Untyped array results (e.g. .toArray()): 'Result: Array (5 elements)' 2. Worker eval log: include line count '(3 lines, 51 chars, db: demo_data)' 3. Schema stats: show 'db/collection' instead of internal clusterId::db::coll --- l10n/bundle.l10n.json | 2 +- .../schemaStore/showSchemaStoreStats.ts | 5 ++- src/documentdb/scratchpad/resultFormatter.ts | 4 +-- src/documentdb/scratchpad/scratchpadWorker.ts | 32 +++++++++++++------ 4 files changed, 30 insertions(+), 13 deletions(-) diff --git a/l10n/bundle.l10n.json b/l10n/bundle.l10n.json index 0b8a78ce..d73b749d 100644 --- a/l10n/bundle.l10n.json +++ b/l10n/bundle.l10n.json @@ -100,7 +100,6 @@ "{0} item(s) already exist in the destination. Check the Output panel for details.": "{0} item(s) already exist in the destination. Check the Output panel for details.", "{0} processed": "{0} processed", "{0} replaced": "{0} replaced", - "{0} results": "{0} results", "{0} skipped": "{0} skipped", "{0} subfolders": "{0} subfolders", "{0} task(s) are using connections being moved. Check the Output panel for details.": "{0} task(s) are using connections being moved. Check the Output panel for details.", @@ -826,6 +825,7 @@ "Report an issue": "Report an issue", "Resource group \"{0}\" already exists in subscription \"{1}\".": "Resource group \"{0}\" already exists in subscription \"{1}\".", "Result: {0}": "Result: {0}", + "Result: Array ({0} elements)": "Result: Array ({0} elements)", "Result: Cursor ({0} documents)": "Result: Cursor ({0} documents)", "Results found": "Results found", "Retry": "Retry", diff --git a/src/commands/schemaStore/showSchemaStoreStats.ts b/src/commands/schemaStore/showSchemaStoreStats.ts index 0fc41df3..f038f34d 100644 --- a/src/commands/schemaStore/showSchemaStoreStats.ts +++ b/src/commands/schemaStore/showSchemaStoreStats.ts @@ -23,8 +23,11 @@ export function showSchemaStoreStats(_context: IActionContext): void { if (stats.collections.length > 0) { for (const c of stats.collections) { + // Key format is "clusterId::db::collection" — show only db/collection + const parts = c.key.split('::'); + const displayKey = parts.length >= 3 ? `${parts[1]}/${parts[2]}` : c.key; ext.outputChannel.appendLog( - `[SchemaStore] ${c.key}: ${String(c.documentCount)} docs, ${String(c.fieldCount)} fields`, + `[SchemaStore] ${displayKey}: ${String(c.documentCount)} docs, ${String(c.fieldCount)} fields`, ); } } else { diff --git a/src/documentdb/scratchpad/resultFormatter.ts b/src/documentdb/scratchpad/resultFormatter.ts index ef1b95ee..13dcad67 100644 --- a/src/documentdb/scratchpad/resultFormatter.ts +++ b/src/documentdb/scratchpad/resultFormatter.ts @@ -39,8 +39,8 @@ export function formatResult(result: ExecutionResult, code: string, connection: // Typed result: "Result: Document", "Result: string", etc. lines.push(`// ${l10n.t('Result: {0}', result.type)}`); } else if (Array.isArray(unwrapped)) { - // Untyped array (e.g. .toArray()) — show count - lines.push(`// ${l10n.t('{0} results', unwrapped.length)}`); + // Untyped array (e.g. .toArray()) — show type and count + lines.push(`// ${l10n.t('Result: Array ({0} elements)', unwrapped.length)}`); } lines.push(`// ${l10n.t('Executed in {0}ms', result.durationMs)}`); diff --git a/src/documentdb/scratchpad/scratchpadWorker.ts b/src/documentdb/scratchpad/scratchpadWorker.ts index 79e29d33..f653867a 100644 --- a/src/documentdb/scratchpad/scratchpadWorker.ts +++ b/src/documentdb/scratchpad/scratchpadWorker.ts @@ -146,7 +146,11 @@ async function handleEval(msg: Extract): throw new Error('Worker not initialized — call init first'); } - log('trace', `Evaluating code (${msg.code.length} chars, db: ${msg.databaseName})`); + const lineCount = msg.code.split('\n').length; + log( + 'trace', + `Evaluating code (${String(lineCount)} lines, ${String(msg.code.length)} chars, db: ${msg.databaseName})`, + ); // Lazy-import @mongosh packages const { EventEmitter } = await import('events'); @@ -197,14 +201,24 @@ async function handleEval(msg: Extract): source?: { namespace?: { db: string; collection: string } }; }; - // Serialize the printable value via EJSON for safe transfer. - // CursorIterationResult is an Array subclass with extra properties (cursorHasMore, - // documents). EJSON.serialize treats it as a plain object and includes those - // properties, producing { cursorHasMore: true, documents: [...] } instead of [...]. - // Convert to a plain Array first to preserve correct array serialization. - const printableValue = Array.isArray(shellResult.printable) - ? Array.from(shellResult.printable as unknown[]) - : shellResult.printable; + // Normalize the printable value for IPC transfer. + // @mongosh's ShellEvaluator wraps cursor results as { cursorHasMore, documents } + // when running in a worker context. Extract the documents array so that the + // main thread receives a clean array (matching the in-process behavior where + // printable was a CursorIterationResult array). + let printableValue: unknown = shellResult.printable; + if ( + shellResult.type === 'Cursor' && + typeof shellResult.printable === 'object' && + shellResult.printable !== null && + 'documents' in shellResult.printable && + Array.isArray((shellResult.printable as { documents?: unknown }).documents) + ) { + printableValue = (shellResult.printable as { documents: unknown[] }).documents; + } else if (Array.isArray(shellResult.printable)) { + // Array subclass (CursorIterationResult) — normalize to plain Array + printableValue = Array.from(shellResult.printable as unknown[]); + } let printableStr: string; try { From ef2cc9be4330dcddd1a07a99c8a49f3c4844c847 Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Thu, 26 Mar 2026 13:10:04 +0000 Subject: [PATCH 14/26] feat(scratchpad): modal connect dialog + multi-statement documentation - Connect instruction dialog is now modal with title/detail separation - Scratchpad template includes note: 'only the last result is displayed' - Help text Tips section includes same note - Wording differs between template and help (not identical) --- l10n/bundle.l10n.json | 2 +- src/commands/scratchpad/connectDatabase.ts | 11 ++++++----- src/commands/scratchpad/newScratchpad.ts | 1 + src/documentdb/scratchpad/ScratchpadEvaluator.ts | 1 + 4 files changed, 9 insertions(+), 6 deletions(-) diff --git a/l10n/bundle.l10n.json b/l10n/bundle.l10n.json index d73b749d..7aae3291 100644 --- a/l10n/bundle.l10n.json +++ b/l10n/bundle.l10n.json @@ -831,6 +831,7 @@ "Retry": "Retry", "Reusing active connection for \"{cluster}\".": "Reusing active connection for \"{cluster}\".", "Revisit connection details and try again.": "Revisit connection details and try again.", + "Right-click a database or collection in the DocumentDB panel and select \"Connect Scratchpad to this database\".": "Right-click a database or collection in the DocumentDB panel and select \"Connect Scratchpad to this database\".", "Role assignment \"{0}\" created for the {2} resource \"{1}\".": "Role assignment \"{0}\" created for the {2} resource \"{1}\".", "Role Assignment {0} created for {1}": "Role Assignment {0} created for {1}", "Role Assignment {0} failed for {1}": "Role Assignment {0} failed for {1}", @@ -1013,7 +1014,6 @@ "This will also delete {0}.": "This will also delete {0}.", "This will prevent the query planner from using this index.": "This will prevent the query planner from using this index.", "To connect to Azure resources, you need to sign in to Azure accounts.": "To connect to Azure resources, you need to sign in to Azure accounts.", - "To connect, right-click a database or collection in the DocumentDB panel and select \"Connect Scratchpad to this database\".": "To connect, right-click a database or collection in the DocumentDB panel and select \"Connect Scratchpad to this database\".", "TODO: Share the steps needed to reliably reproduce the problem. Please include actual and expected results.": "TODO: Share the steps needed to reliably reproduce the problem. Please include actual and expected results.", "Total documents to import: {0}": "Total documents to import: {0}", "Total time taken to execute the query on the server": "Total time taken to execute the query on the server", diff --git a/src/commands/scratchpad/connectDatabase.ts b/src/commands/scratchpad/connectDatabase.ts index 263955c3..c2fb03fa 100644 --- a/src/commands/scratchpad/connectDatabase.ts +++ b/src/commands/scratchpad/connectDatabase.ts @@ -27,11 +27,12 @@ export async function connectDatabase(_context: IActionContext, node?: DatabaseI l10n.t('DocumentDB Scratchpad connected to {0}/{1}', node.cluster.name, node.databaseInfo.name), ); } else { - // No tree context — show instructions (per plan §2.3) - void vscode.window.showInformationMessage( - l10n.t( - 'To connect, right-click a database or collection in the DocumentDB panel and select "Connect Scratchpad to this database".', + // No tree context — show instructions as modal dialog + void vscode.window.showInformationMessage(l10n.t('No database connected'), { + modal: true, + detail: l10n.t( + 'Right-click a database or collection in the DocumentDB panel and select "Connect Scratchpad to this database".', ), - ); + }); } } diff --git a/src/commands/scratchpad/newScratchpad.ts b/src/commands/scratchpad/newScratchpad.ts index 7121a0c3..d3fcdd8a 100644 --- a/src/commands/scratchpad/newScratchpad.ts +++ b/src/commands/scratchpad/newScratchpad.ts @@ -39,6 +39,7 @@ export async function newScratchpad(_context: IActionContext, node?: DatabaseIte headerComment, '// Use Ctrl+Enter (Cmd+Enter) to run the current block', '// Use Ctrl+Shift+Enter (Cmd+Shift+Enter) to run the entire file', + '// Note: when running multiple statements, only the last result is displayed', '', `db.getCollection('${collectionName}').find({ })`, '', diff --git a/src/documentdb/scratchpad/ScratchpadEvaluator.ts b/src/documentdb/scratchpad/ScratchpadEvaluator.ts index 5b2070c3..b648554d 100644 --- a/src/documentdb/scratchpad/ScratchpadEvaluator.ts +++ b/src/documentdb/scratchpad/ScratchpadEvaluator.ts @@ -524,6 +524,7 @@ export class ScratchpadEvaluator implements vscode.Disposable { 'Tips:', ' • Separate code blocks with blank lines', ' • Variables persist within a block but not between separate runs', + ' • When running multiple statements, only the last result is shown', ' • Use .toArray() to get all results (default: first 20 documents)', ].join('\n'); From e2db9febae81b03cb0dc2fe403dd58ea3c076405 Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Thu, 26 Mar 2026 13:44:10 +0000 Subject: [PATCH 15/26] fix(scratchpad): use generic 'client' terminology in logs and comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace 'MongoClient' with 'client' or 'database client' in: - Worker log messages visible in the output channel - JSDoc comments describing worker behavior - Code comments in worker and evaluator Type references (MongoClient, MongoClientOptions) are unchanged — these are the actual driver API names. The extension is a DocumentDB tool using the MongoDB API wire protocol. User-facing text should not reference MongoDB implementation details. --- src/documentdb/scratchpad/ScratchpadEvaluator.ts | 6 +++--- src/documentdb/scratchpad/scratchpadWorker.ts | 12 ++++++------ 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/documentdb/scratchpad/ScratchpadEvaluator.ts b/src/documentdb/scratchpad/ScratchpadEvaluator.ts index b648554d..bb5cdc8a 100644 --- a/src/documentdb/scratchpad/ScratchpadEvaluator.ts +++ b/src/documentdb/scratchpad/ScratchpadEvaluator.ts @@ -24,10 +24,10 @@ type WorkerState = 'idle' | 'spawning' | 'ready' | 'executing'; /** * Evaluates scratchpad code in a persistent worker thread. * - * The worker owns its own `MongoClient` (authenticated via credentials from + * The worker owns its own database client (authenticated via credentials from * `CredentialCache`) and stays alive between runs. This provides: * - Infinite loop safety (main thread can kill the worker) - * - MongoClient isolation from the Collection View + * - Client isolation from the Collection View * - Zero re-auth overhead after the first run * * The public API is unchanged from the in-process evaluator: @@ -86,7 +86,7 @@ export class ScratchpadEvaluator implements vscode.Disposable { } /** - * Gracefully shut down the worker: close MongoClient, then terminate thread. + * Gracefully shut down the worker: close the database client, then terminate thread. * Returns after the worker has confirmed shutdown or after a timeout. */ async shutdown(): Promise { diff --git a/src/documentdb/scratchpad/scratchpadWorker.ts b/src/documentdb/scratchpad/scratchpadWorker.ts index f653867a..4b9833f8 100644 --- a/src/documentdb/scratchpad/scratchpadWorker.ts +++ b/src/documentdb/scratchpad/scratchpadWorker.ts @@ -7,7 +7,7 @@ * Worker thread entry point for scratchpad code evaluation. * * This file runs in a Node.js worker_thread, isolated from the extension host. - * It owns its own MongoClient instance (authenticated via credentials passed from + * It owns its own database client instance (authenticated via credentials passed from * the main thread at init time) and evaluates user code through the @mongosh pipeline. * * Communication with the main thread is via postMessage (structured clone). @@ -95,7 +95,7 @@ async function handleInit(msg: Extract): // Lazy-import MongoDB driver const { MongoClient } = await import('mongodb'); - // Build MongoClient options from the serializable subset + // Build client options from the serializable subset const options: MongoClientOptions = { ...msg.clientOptions, }; @@ -124,12 +124,12 @@ async function handleInit(msg: Extract): }; } - // Create and connect MongoClient + // Create and connect the database client mongoClient = new MongoClient(msg.connectionString, options); await mongoClient.connect(); currentDatabaseName = msg.databaseName; - log('debug', 'Worker initialized — MongoClient connected'); + log('debug', 'Worker initialized — client connected'); const response: WorkerToMainMessage = { type: 'initResult', @@ -258,7 +258,7 @@ async function handleEval(msg: Extract): // ─── Shutdown handler ──────────────────────────────────────────────────────── async function handleShutdown(msg: Extract): Promise { - log('debug', 'Shutting down worker — closing MongoClient'); + log('debug', 'Shutting down worker — closing client'); try { if (mongoClient) { @@ -267,7 +267,7 @@ async function handleShutdown(msg: Extract Date: Thu, 26 Mar 2026 15:55:18 +0000 Subject: [PATCH 16/26] fix(scratchpad): clean up worker on init failure to prevent stranded state If buildInitMessage() throws or the worker reports initResult { success: false }, spawnWorker() now calls terminateWorker() before rethrowing. This returns the evaluator to 'idle' state so the next evaluate() call can respawn a fresh worker instead of being stuck in the 'spawning' state indefinitely. --- .../scratchpad/ScratchpadEvaluator.ts | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/src/documentdb/scratchpad/ScratchpadEvaluator.ts b/src/documentdb/scratchpad/ScratchpadEvaluator.ts index bb5cdc8a..7cdb73e5 100644 --- a/src/documentdb/scratchpad/ScratchpadEvaluator.ts +++ b/src/documentdb/scratchpad/ScratchpadEvaluator.ts @@ -163,13 +163,20 @@ export class ScratchpadEvaluator implements vscode.Disposable { ext.outputChannel.error(`[Scratchpad Worker] ${error.message}`); }); - // Build init message from cached credentials - const initMsg = this.buildInitMessage(connection); + // Build init message from cached credentials and send to worker. + // If init fails (bad credentials, unreachable host, etc.), tear down + // the worker so the next evaluate() call can respawn cleanly. + try { + const initMsg = this.buildInitMessage(connection); - // Send init and wait for acknowledgment - onProgress?.(l10n.t('Authenticating…')); - await this.sendRequest(initMsg, 30000); - this._workerState = 'ready'; + // Send init and wait for acknowledgment + onProgress?.(l10n.t('Authenticating…')); + await this.sendRequest(initMsg, 30000); + this._workerState = 'ready'; + } catch (error) { + this.terminateWorker(); + throw error; + } } /** From 12a02e114909446d427569536ceda1ee20619114 Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Thu, 26 Mar 2026 15:55:59 +0000 Subject: [PATCH 17/26] fix(scratchpad): use canonical EJSON for worker IPC to preserve BSON numeric types MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Switch from relaxed to canonical EJSON (relaxed: false) for the worker-to-main IPC payload. Canonical EJSON preserves Int32, Long, Double, and Decimal128 type wrappers so that EJSON.parse on the main thread reconstructs actual BSON instances. This allows SchemaAnalyzer.inferType() to correctly distinguish numeric subtypes, which feeds into type-aware operator ranking in completions. Also drops the space/indent parameter from EJSON.stringify since the IPC payload is never displayed to users — reducing transfer size. Fixes the incorrect comment that claimed Int32/Long collapse was a fundamental EJSON limitation (it was caused by using relaxed mode). --- src/documentdb/scratchpad/ScratchpadEvaluator.ts | 10 ++++------ src/documentdb/scratchpad/scratchpadWorker.ts | 4 ++-- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/src/documentdb/scratchpad/ScratchpadEvaluator.ts b/src/documentdb/scratchpad/ScratchpadEvaluator.ts index 7cdb73e5..f822f280 100644 --- a/src/documentdb/scratchpad/ScratchpadEvaluator.ts +++ b/src/documentdb/scratchpad/ScratchpadEvaluator.ts @@ -241,16 +241,14 @@ export class ScratchpadEvaluator implements vscode.Disposable { try { const result = await this.sendRequest<{ result: SerializableExecutionResult }>(evalMsg, timeoutMs); - // Deserialize the result — printable is an EJSON string from the worker. - // Use EJSON.parse to reconstruct BSON types (ObjectId, Date, Decimal128, etc.) - // so that SchemaAnalyzer correctly identifies field types. - // Int32 and Long are still collapsed to Double (JavaScript number) — this is - // a fundamental EJSON limitation, not a bug. + // Deserialize the result — printable is a canonical EJSON string from the worker. + // Canonical EJSON preserves all BSON types (ObjectId, Date, Decimal128, Int32, + // Long, Double, etc.) so that SchemaAnalyzer correctly identifies field types. const serResult = result.result; let printable: unknown; try { const { EJSON } = await import('bson'); - printable = EJSON.parse(serResult.printable); + printable = EJSON.parse(serResult.printable, { relaxed: false }); } catch { // Fallback to JSON.parse if EJSON fails, then raw string try { diff --git a/src/documentdb/scratchpad/scratchpadWorker.ts b/src/documentdb/scratchpad/scratchpadWorker.ts index 4b9833f8..b33dc08e 100644 --- a/src/documentdb/scratchpad/scratchpadWorker.ts +++ b/src/documentdb/scratchpad/scratchpadWorker.ts @@ -223,11 +223,11 @@ async function handleEval(msg: Extract): let printableStr: string; try { const { EJSON } = await import('bson'); - printableStr = EJSON.stringify(printableValue, { relaxed: true }, 2); + printableStr = EJSON.stringify(printableValue, { relaxed: false }); } catch { // Fallback: try JSON, then plain string try { - printableStr = JSON.stringify(printableValue, null, 2); + printableStr = JSON.stringify(printableValue); } catch { printableStr = String(printableValue); } From 450443ef1d741a5f8f6297dda05528920336a173 Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Thu, 26 Mar 2026 15:56:31 +0000 Subject: [PATCH 18/26] chore(scratchpad): add TODO comments for unwired displayBatchSize field MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mark displayBatchSize in workerTypes.ts and ScratchpadEvaluator.ts with TODO(F11) comments noting the field is sent but not yet read by the worker. References future-work.md §F11 for the plan to wire documentDB.mongoShell.batchSize. --- src/documentdb/scratchpad/ScratchpadEvaluator.ts | 1 + src/documentdb/scratchpad/workerTypes.ts | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/documentdb/scratchpad/ScratchpadEvaluator.ts b/src/documentdb/scratchpad/ScratchpadEvaluator.ts index f822f280..c6b994c8 100644 --- a/src/documentdb/scratchpad/ScratchpadEvaluator.ts +++ b/src/documentdb/scratchpad/ScratchpadEvaluator.ts @@ -217,6 +217,7 @@ export class ScratchpadEvaluator implements vscode.Disposable { databaseName: connection.databaseName, authMechanism: authMechanism as 'NativeAuth' | 'MicrosoftEntraID', tenantId: credentials.entraIdConfig?.tenantId, + // TODO(F11): Read from documentDB.mongoShell.batchSize setting and wire in worker displayBatchSize: 50, }; } diff --git a/src/documentdb/scratchpad/workerTypes.ts b/src/documentdb/scratchpad/workerTypes.ts index 7018fc9d..af13b0e5 100644 --- a/src/documentdb/scratchpad/workerTypes.ts +++ b/src/documentdb/scratchpad/workerTypes.ts @@ -61,7 +61,8 @@ export type MainToWorkerMessage = readonly authMechanism: 'NativeAuth' | 'MicrosoftEntraID'; /** Tenant ID for Entra ID clusters */ readonly tenantId?: string; - /** Display batch size from documentDB.mongoShell.batchSize setting */ + // TODO(F11): Wire displayBatchSize end-to-end — currently sent but not read by the worker. + // See future-work.md §F11 for the plan to honor documentDB.mongoShell.batchSize. readonly displayBatchSize: number; } | { From 10976e1b416a12e1e9615cea87a6c36593a7b74a Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Thu, 26 Mar 2026 15:57:09 +0000 Subject: [PATCH 19/26] fix(scratchpad): localize user-visible error messages in ScratchpadEvaluator Wrap the three error strings most likely to reach users in l10n.t(): - 'No credentials found for cluster {0}' - 'Worker is not running' - 'Execution timed out after {0} seconds' These flow through to vscode.window.showErrorMessage via the catch in executeScratchpadCode.ts, so non-English users now see translated details. --- l10n/bundle.l10n.json | 3 +++ src/documentdb/scratchpad/ScratchpadEvaluator.ts | 6 +++--- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/l10n/bundle.l10n.json b/l10n/bundle.l10n.json index 7aae3291..bd96f787 100644 --- a/l10n/bundle.l10n.json +++ b/l10n/bundle.l10n.json @@ -438,6 +438,7 @@ "Execution Strategy": "Execution Strategy", "Execution Time": "Execution Time", "Execution timed out": "Execution timed out", + "Execution timed out after {0} seconds": "Execution timed out after {0} seconds", "Execution timed out.": "Execution timed out.", "Exit": "Exit", "Explain(aggregate) completed [{durationMs}ms]": "Explain(aggregate) completed [{durationMs}ms]", @@ -712,6 +713,7 @@ "No collection has been marked for copy. Please use \"Copy Collection...\" first to select a source collection.": "No collection has been marked for copy. Please use \"Copy Collection...\" first to select a source collection.", "No collection selected.": "No collection selected.", "No Connectivity": "No Connectivity", + "No credentials found for cluster {0}": "No credentials found for cluster {0}", "No credentials found for id {clusterId}": "No credentials found for id {clusterId}", "No credentials found for the selected cluster.": "No credentials found for the selected cluster.", "No database connected": "No database connected", @@ -1083,6 +1085,7 @@ "What's New": "What's New", "Where to save the exported documents?": "Where to save the exported documents?", "with Popover": "with Popover", + "Worker is not running": "Worker is not running", "Working...": "Working...", "Working…": "Working…", "Would you like to open the Collection View?": "Would you like to open the Collection View?", diff --git a/src/documentdb/scratchpad/ScratchpadEvaluator.ts b/src/documentdb/scratchpad/ScratchpadEvaluator.ts index c6b994c8..dfa88910 100644 --- a/src/documentdb/scratchpad/ScratchpadEvaluator.ts +++ b/src/documentdb/scratchpad/ScratchpadEvaluator.ts @@ -185,7 +185,7 @@ export class ScratchpadEvaluator implements vscode.Disposable { private buildInitMessage(connection: ScratchpadConnection): MainToWorkerMessage & { type: 'init' } { const credentials = CredentialCache.getCredentials(connection.clusterId); if (!credentials) { - throw new Error(`No credentials found for cluster ${connection.clusterId}`); + throw new Error(l10n.t('No credentials found for cluster {0}', connection.clusterId)); } const authMechanism = credentials.authMechanism ?? 'NativeAuth'; @@ -280,7 +280,7 @@ export class ScratchpadEvaluator implements vscode.Disposable { */ private sendRequest(msg: MainToWorkerMessage, timeoutMs: number): Promise { if (!this._worker) { - return Promise.reject(new Error('Worker is not running')); + return Promise.reject(new Error(l10n.t('Worker is not running'))); } const requestId = randomUUID(); @@ -299,7 +299,7 @@ export class ScratchpadEvaluator implements vscode.Disposable { this._pendingRequests.delete(requestId); this.killWorker(); pending.reject( - new Error(`Execution timed out after ${String(Math.round(timeoutMs / 1000))} seconds`), + new Error(l10n.t('Execution timed out after {0} seconds', String(Math.round(timeoutMs / 1000)))), ); } }, timeoutMs); From 7eba1648bccca4e3b21b9697824baaf26cd8488f Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Thu, 26 Mar 2026 17:07:15 +0000 Subject: [PATCH 20/26] chore: l10n --- src/documentdb/scratchpad/ScratchpadEvaluator.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/documentdb/scratchpad/ScratchpadEvaluator.ts b/src/documentdb/scratchpad/ScratchpadEvaluator.ts index dfa88910..899dc1af 100644 --- a/src/documentdb/scratchpad/ScratchpadEvaluator.ts +++ b/src/documentdb/scratchpad/ScratchpadEvaluator.ts @@ -299,7 +299,9 @@ export class ScratchpadEvaluator implements vscode.Disposable { this._pendingRequests.delete(requestId); this.killWorker(); pending.reject( - new Error(l10n.t('Execution timed out after {0} seconds', String(Math.round(timeoutMs / 1000)))), + new Error( + l10n.t('Execution timed out after {0} seconds', String(Math.round(timeoutMs / 1000))), + ), ); } }, timeoutMs); From 137ddc15044ad1ac045f5c07c53b35f9d21a4e13 Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Fri, 27 Mar 2026 15:02:25 +0000 Subject: [PATCH 21/26] feat(scratchpad): enhance telemetry for scratchpad execution and worker lifecycle --- .../scratchpad/executeScratchpadCode.ts | 184 ++++++++++++------ .../scratchpad/ScratchpadEvaluator.ts | 45 +++++ src/documentdb/utils/getClusterMetadata.ts | 11 ++ 3 files changed, 185 insertions(+), 55 deletions(-) diff --git a/src/commands/scratchpad/executeScratchpadCode.ts b/src/commands/scratchpad/executeScratchpadCode.ts index 2e53d0bf..0c7d04e8 100644 --- a/src/commands/scratchpad/executeScratchpadCode.ts +++ b/src/commands/scratchpad/executeScratchpadCode.ts @@ -3,15 +3,22 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ -import { openReadOnlyContent } from '@microsoft/vscode-azext-utils'; +import { + UserCancelledError, + callWithTelemetryAndErrorHandling, + openReadOnlyContent, +} from '@microsoft/vscode-azext-utils'; import * as l10n from '@vscode/l10n'; import { type Document, type WithId } from 'mongodb'; import * as vscode from 'vscode'; +import { CredentialCache } from '../../documentdb/CredentialCache'; import { SchemaStore } from '../../documentdb/SchemaStore'; import { ScratchpadEvaluator } from '../../documentdb/scratchpad/ScratchpadEvaluator'; import { ScratchpadService } from '../../documentdb/scratchpad/ScratchpadService'; import { formatError, formatResult } from '../../documentdb/scratchpad/resultFormatter'; import { type ExecutionResult, type ScratchpadConnection } from '../../documentdb/scratchpad/types'; +import { getHostsFromConnectionString } from '../../documentdb/utils/connectionStringHelpers'; +import { addDomainInfoToProperties } from '../../documentdb/utils/getClusterMetadata'; /** Shared evaluator instance — lazily created, reused across runs. */ let evaluator: ScratchpadEvaluator | undefined; @@ -56,65 +63,132 @@ export async function executeScratchpadCode(code: string): Promise { service.setExecuting(true); - await vscode.window.withProgress( - { - location: vscode.ProgressLocation.Notification, - title: l10n.t('DocumentDB Scratchpad'), - cancellable: true, - }, - async (progress, token) => { - // Track whether user cancelled — to suppress the error notification - let cancelled = false; - - // Cancel kills the worker — user can re-run to respawn - token.onCancellationRequested(() => { - cancelled = true; - evaluator?.killWorker(); - }); - - const startTime = Date.now(); - try { - const result = await evaluator!.evaluate(connection, code, (message) => { - progress.report({ message }); + // callWithTelemetryAndErrorHandling automatically tracks: + // - duration (measured from callback start to end) + // - result: 'Succeeded' | 'Failed' | 'Canceled' + // - error / errorMessage (from thrown errors) + // We add custom properties for scratchpad-specific analytics. + await callWithTelemetryAndErrorHandling('scratchpad.execute', async (context) => { + context.errorHandling.suppressDisplay = true; // we show our own error UI + context.errorHandling.rethrow = false; + + // ── Pre-execution telemetry (known before eval) ────────────── + context.telemetry.properties.sessionId = evaluator!.sessionId ?? 'none'; + context.telemetry.properties.sessionEvalCount = String(evaluator!.sessionEvalCount); + context.telemetry.properties.authMethod = evaluator!.sessionAuthMethod ?? 'unknown'; + context.telemetry.measurements.codeLineCount = code.split('\n').length; + + // Domain info — privacy-safe hashed host data for platform analytics + const domainProps: Record = {}; + collectDomainTelemetry(connection, domainProps); + Object.assign(context.telemetry.properties, domainProps); + + await vscode.window.withProgress( + { + location: vscode.ProgressLocation.Notification, + title: l10n.t('DocumentDB Scratchpad'), + cancellable: true, + }, + async (progress, token) => { + let cancelled = false; + + token.onCancellationRequested(() => { + cancelled = true; + evaluator?.killWorker(); }); - const formattedOutput = formatResult(result, code, connection); - - // Feed document results to SchemaStore for cross-tab schema sharing - feedResultToSchemaStore(result, connection); - - const resultLabel = l10n.t('{0}/{1} — Results', connection.clusterDisplayName, connection.databaseName); - - await openReadOnlyContent( - { label: resultLabel, fullId: `scratchpad-results-${Date.now()}` }, - formattedOutput, - '.jsonc', - { viewColumn: vscode.ViewColumn.Beside, preserveFocus: true }, - ); - } catch (error: unknown) { - // Don't show error UI when the user explicitly cancelled - if (cancelled) { - return; - } - const durationMs = Date.now() - startTime; - const errorMessage = error instanceof Error ? error.message : String(error); - const formattedOutput = formatError(error, code, durationMs, connection); + try { + const result = await evaluator!.evaluate(connection, code, (message) => { + progress.report({ message }); + }); - const errorLabel = l10n.t('{0}/{1} — Error', connection.clusterDisplayName, connection.databaseName); + // ── Post-execution telemetry (known after success) ──── + context.telemetry.properties.resultType = result.type ?? 'null'; + context.telemetry.measurements.initDurationMs = evaluator!.lastInitDurationMs; + // sessionId/sessionEvalCount may have changed after evaluate (if worker was spawned) + context.telemetry.properties.sessionId = evaluator!.sessionId ?? 'none'; + context.telemetry.properties.sessionEvalCount = String(evaluator!.sessionEvalCount); + context.telemetry.properties.authMethod = evaluator!.sessionAuthMethod ?? 'unknown'; - await openReadOnlyContent( - { label: errorLabel, fullId: `scratchpad-error-${Date.now()}` }, - formattedOutput, - '.jsonc', - { viewColumn: vscode.ViewColumn.Beside, preserveFocus: true }, - ); + const formattedOutput = formatResult(result, code, connection); + feedResultToSchemaStore(result, connection); - void vscode.window.showErrorMessage(l10n.t('Scratchpad execution failed: {0}', errorMessage)); - } finally { - service.setExecuting(false); - } - }, - ); + const resultLabel = l10n.t( + '{0}/{1} — Results', + connection.clusterDisplayName, + connection.databaseName, + ); + + await openReadOnlyContent( + { label: resultLabel, fullId: `scratchpad-results-${Date.now()}` }, + formattedOutput, + '.jsonc', + { viewColumn: vscode.ViewColumn.Beside, preserveFocus: true }, + ); + + // result: 'Succeeded' is set automatically by the framework + } catch (error: unknown) { + // Update session telemetry even on failure (worker may have spawned before failing) + context.telemetry.properties.sessionId = evaluator!.sessionId ?? 'none'; + context.telemetry.properties.sessionEvalCount = String(evaluator!.sessionEvalCount); + context.telemetry.properties.authMethod = evaluator!.sessionAuthMethod ?? 'unknown'; + context.telemetry.measurements.initDurationMs = evaluator!.lastInitDurationMs; + + if (cancelled) { + // Throw UserCancelledError so framework marks result as 'Canceled' + throw new UserCancelledError('scratchpad.execute'); + } + + // Show our own error UI before re-throwing + const errorMessage = error instanceof Error ? error.message : String(error); + const formattedOutput = formatError(error, code, 0, connection); + + const errorLabel = l10n.t( + '{0}/{1} — Error', + connection.clusterDisplayName, + connection.databaseName, + ); + + await openReadOnlyContent( + { label: errorLabel, fullId: `scratchpad-error-${Date.now()}` }, + formattedOutput, + '.jsonc', + { viewColumn: vscode.ViewColumn.Beside, preserveFocus: true }, + ); + + void vscode.window.showErrorMessage(l10n.t('Scratchpad execution failed: {0}', errorMessage)); + + // Re-throw so framework automatically captures result: 'Failed', + // error, and errorMessage in telemetry + throw error; + } finally { + service.setExecuting(false); + } + }, + ); + }); +} + +// ─── Domain telemetry ──────────────────────────────────────────────────────── + +/** + * Collects domain info from the scratchpad connection's cached credentials. + * Reuses the same hashing logic as the connection metadata telemetry. + */ +function collectDomainTelemetry( + connection: ScratchpadConnection, + properties: Record, +): void { + try { + const credentials = CredentialCache.getCredentials(connection.clusterId); + if (!credentials?.connectionString) { + return; + } + const hosts = getHostsFromConnectionString(credentials.connectionString); + addDomainInfoToProperties(hosts, properties); + } catch { + // Domain info is best-effort — don't fail telemetry if parsing fails + } } /** diff --git a/src/documentdb/scratchpad/ScratchpadEvaluator.ts b/src/documentdb/scratchpad/ScratchpadEvaluator.ts index 899dc1af..8b7296cb 100644 --- a/src/documentdb/scratchpad/ScratchpadEvaluator.ts +++ b/src/documentdb/scratchpad/ScratchpadEvaluator.ts @@ -39,6 +39,19 @@ export class ScratchpadEvaluator implements vscode.Disposable { /** Which cluster the live worker is connected to (to detect cluster switches) */ private _workerClusterId: string | undefined; + /** + * Telemetry session ID — generated on worker spawn, stable across evals within the + * same worker lifecycle. Resets when the worker is terminated/respawned. + * Used to correlate multiple scratchpad runs within a single "session". + */ + private _sessionId: string | undefined; + + /** Number of eval calls completed during this worker session (for usage tracking). */ + private _sessionEvalCount: number = 0; + + /** Auth mechanism used for the current worker session (for telemetry). */ + private _sessionAuthMethod: string | undefined; + /** Pending request correlation map: requestId → { resolve, reject } */ private _pendingRequests = new Map< string, @@ -48,6 +61,23 @@ export class ScratchpadEvaluator implements vscode.Disposable { } >(); + /** Telemetry accessors — read by the command layer for telemetry properties. */ + get sessionId(): string | undefined { + return this._sessionId; + } + get sessionEvalCount(): number { + return this._sessionEvalCount; + } + get sessionAuthMethod(): string | undefined { + return this._sessionAuthMethod; + } + + /** Duration of the last worker init (spawn + auth), in ms. 0 if worker was already alive. */ + private _lastInitDurationMs: number = 0; + get lastInitDurationMs(): number { + return this._lastInitDurationMs; + } + /** * Evaluate user code against the connected database. * @@ -74,7 +104,10 @@ export class ScratchpadEvaluator implements vscode.Disposable { if (needsSpawn) { onProgress?.(l10n.t('Initializing…')); } + + const initStartTime = Date.now(); await this.ensureWorker(connection, onProgress); + this._lastInitDurationMs = needsSpawn ? Date.now() - initStartTime : 0; // Send eval message and await result onProgress?.(l10n.t('Running query…')); @@ -169,6 +202,11 @@ export class ScratchpadEvaluator implements vscode.Disposable { try { const initMsg = this.buildInitMessage(connection); + // Start a new telemetry session for this worker lifecycle + this._sessionId = randomUUID(); + this._sessionEvalCount = 0; + this._sessionAuthMethod = initMsg.authMechanism; + // Send init and wait for acknowledgment onProgress?.(l10n.t('Authenticating…')); await this.sendRequest(initMsg, 30000); @@ -231,6 +269,7 @@ export class ScratchpadEvaluator implements vscode.Disposable { timeoutMs: number, ): Promise { this._workerState = 'executing'; + this._sessionEvalCount++; const evalMsg: MainToWorkerMessage = { type: 'eval', @@ -441,6 +480,9 @@ export class ScratchpadEvaluator implements vscode.Disposable { } this._workerState = 'idle'; this._workerClusterId = undefined; + this._sessionId = undefined; + this._sessionEvalCount = 0; + this._sessionAuthMethod = undefined; // Reject all pending requests for (const [, entry] of this._pendingRequests) { @@ -453,6 +495,9 @@ export class ScratchpadEvaluator implements vscode.Disposable { this._worker = undefined; this._workerState = 'idle'; this._workerClusterId = undefined; + this._sessionId = undefined; + this._sessionEvalCount = 0; + this._sessionAuthMethod = undefined; // Reject any still-pending requests for (const [, entry] of this._pendingRequests) { diff --git a/src/documentdb/utils/getClusterMetadata.ts b/src/documentdb/utils/getClusterMetadata.ts index 36430136..1e433f7a 100644 --- a/src/documentdb/utils/getClusterMetadata.ts +++ b/src/documentdb/utils/getClusterMetadata.ts @@ -108,6 +108,17 @@ async function fetchHostInfo(adminDb: Admin, result: ClusterMetadata): Promise): void { for (const [index, host] of hosts.entries()) { const telemetrySuffix = index > 0 ? `_h${index}` : ''; try { From f30c7b6af6d1a9b929b909fa96c7b02f46ea5292 Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Fri, 27 Mar 2026 15:15:55 +0000 Subject: [PATCH 22/26] feat(scratchpad): add runMode parameter to executeScratchpadCode for improved telemetry --- src/commands/scratchpad/executeScratchpadCode.ts | 5 ++++- src/commands/scratchpad/runAll.ts | 2 +- src/commands/scratchpad/runSelected.ts | 2 +- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/src/commands/scratchpad/executeScratchpadCode.ts b/src/commands/scratchpad/executeScratchpadCode.ts index 0c7d04e8..0c934e6c 100644 --- a/src/commands/scratchpad/executeScratchpadCode.ts +++ b/src/commands/scratchpad/executeScratchpadCode.ts @@ -45,7 +45,9 @@ export function shutdownEvaluator(): void { * Executes scratchpad code and displays the result in a read-only side panel. * Used by both `runAll` and `runSelected` commands. */ -export async function executeScratchpadCode(code: string): Promise { +export type ScratchpadRunMode = 'runAll' | 'runSelected'; + +export async function executeScratchpadCode(code: string, runMode: ScratchpadRunMode): Promise { const service = ScratchpadService.getInstance(); const connection = service.getConnection(); if (!connection) { @@ -76,6 +78,7 @@ export async function executeScratchpadCode(code: string): Promise { context.telemetry.properties.sessionId = evaluator!.sessionId ?? 'none'; context.telemetry.properties.sessionEvalCount = String(evaluator!.sessionEvalCount); context.telemetry.properties.authMethod = evaluator!.sessionAuthMethod ?? 'unknown'; + context.telemetry.properties.runMode = runMode; context.telemetry.measurements.codeLineCount = code.split('\n').length; // Domain info — privacy-safe hashed host data for platform analytics diff --git a/src/commands/scratchpad/runAll.ts b/src/commands/scratchpad/runAll.ts index 65016eaf..d9781730 100644 --- a/src/commands/scratchpad/runAll.ts +++ b/src/commands/scratchpad/runAll.ts @@ -32,5 +32,5 @@ export async function runAll(_context: IActionContext): Promise { return; } - await executeScratchpadCode(code); + await executeScratchpadCode(code, 'runAll'); } diff --git a/src/commands/scratchpad/runSelected.ts b/src/commands/scratchpad/runSelected.ts index 7f62b7d2..0cf30c06 100644 --- a/src/commands/scratchpad/runSelected.ts +++ b/src/commands/scratchpad/runSelected.ts @@ -68,5 +68,5 @@ export async function runSelected(_context: IActionContext, startLine?: number, return; } - await executeScratchpadCode(codeToRun); + await executeScratchpadCode(codeToRun, 'runSelected'); } From 986e2b8f698d02a45d82ae11ee3fe3fc23e21df8 Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Fri, 27 Mar 2026 16:29:23 +0000 Subject: [PATCH 23/26] fix(scratchpad): tighten cursor unwrap to require full CursorIterationResult shape The unwrapCursorResult() check and feedResultToSchemaStore() unwrap now require both 'cursorHasMore' (boolean) and 'documents' (array) before unwrapping. Previously, any object with a 'documents' array field would be unwrapped, which could false-positive on user documents with a 'documents' field. --- src/commands/scratchpad/executeScratchpadCode.ts | 5 ++++- src/documentdb/scratchpad/resultFormatter.ts | 9 ++++++--- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/src/commands/scratchpad/executeScratchpadCode.ts b/src/commands/scratchpad/executeScratchpadCode.ts index 0c934e6c..74b0bcd6 100644 --- a/src/commands/scratchpad/executeScratchpadCode.ts +++ b/src/commands/scratchpad/executeScratchpadCode.ts @@ -223,11 +223,14 @@ function feedResultToSchemaStore(result: ExecutionResult, connection: Scratchpad } // CursorIterationResult from @mongosh wraps documents in { cursorHasMore, documents }. - // Unwrap to get the actual document array. + // Only unwrap when the full wrapper shape is present to avoid false positives + // on user documents that happen to have a `documents` field. let items: unknown[]; if ( typeof printable === 'object' && !Array.isArray(printable) && + 'cursorHasMore' in printable && + typeof (printable as Record).cursorHasMore === 'boolean' && 'documents' in printable && Array.isArray((printable as { documents: unknown }).documents) ) { diff --git a/src/documentdb/scratchpad/resultFormatter.ts b/src/documentdb/scratchpad/resultFormatter.ts index 13dcad67..2b21e368 100644 --- a/src/documentdb/scratchpad/resultFormatter.ts +++ b/src/documentdb/scratchpad/resultFormatter.ts @@ -88,9 +88,10 @@ export function formatError( /** * Unwrap CursorIterationResult from @mongosh. * - * @mongosh's `asPrintable()` on CursorIterationResult produces `{ cursorHasMore, documents }` - * instead of a plain array. This function extracts the `documents` array for clean display - * and schema feeding. + * @mongosh's `asPrintable()` on CursorIterationResult produces + * `{ cursorHasMore: boolean, documents: unknown[] }` instead of a plain array. + * Only unwraps when the full wrapper shape is present to avoid + * false positives on user documents that happen to have a `documents` field. */ function unwrapCursorResult(printable: unknown): unknown { if ( @@ -98,6 +99,8 @@ function unwrapCursorResult(printable: unknown): unknown { printable !== undefined && typeof printable === 'object' && !Array.isArray(printable) && + 'cursorHasMore' in printable && + typeof (printable as Record).cursorHasMore === 'boolean' && 'documents' in printable && Array.isArray((printable as { documents: unknown }).documents) ) { From 497ff533086bfb42c99d394a427e298e0c9b87f2 Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Fri, 27 Mar 2026 16:30:27 +0000 Subject: [PATCH 24/26] fix(scratchpad): use generic 'Operation timed out' message for sendRequest timeout The timeout in sendRequest() is used for init, eval, and shutdown. The previous message 'Execution timed out' was misleading when init hangs. Changed to 'Operation timed out after {0} seconds' which is accurate for all callers. --- l10n/bundle.l10n.json | 2 +- src/documentdb/scratchpad/ScratchpadEvaluator.ts | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/l10n/bundle.l10n.json b/l10n/bundle.l10n.json index bd96f787..ec52f624 100644 --- a/l10n/bundle.l10n.json +++ b/l10n/bundle.l10n.json @@ -438,7 +438,6 @@ "Execution Strategy": "Execution Strategy", "Execution Time": "Execution Time", "Execution timed out": "Execution timed out", - "Execution timed out after {0} seconds": "Execution timed out after {0} seconds", "Execution timed out.": "Execution timed out.", "Exit": "Exit", "Explain(aggregate) completed [{durationMs}ms]": "Explain(aggregate) completed [{durationMs}ms]", @@ -747,6 +746,7 @@ "Open the VS Code Marketplace to learn more about \"{0}\"": "Open the VS Code Marketplace to learn more about \"{0}\"", "Opening DocumentDB connection…": "Opening DocumentDB connection…", "Operation cancelled.": "Operation cancelled.", + "Operation timed out after {0} seconds": "Operation timed out after {0} seconds", "Optimization Opportunities": "Optimization Opportunities", "Optimize Index Strategy": "Optimize Index Strategy", "Optimizing the index on {0} can improve query performance by better matching the query pattern.": "Optimizing the index on {0} can improve query performance by better matching the query pattern.", diff --git a/src/documentdb/scratchpad/ScratchpadEvaluator.ts b/src/documentdb/scratchpad/ScratchpadEvaluator.ts index 8b7296cb..fccabcb9 100644 --- a/src/documentdb/scratchpad/ScratchpadEvaluator.ts +++ b/src/documentdb/scratchpad/ScratchpadEvaluator.ts @@ -339,7 +339,7 @@ export class ScratchpadEvaluator implements vscode.Disposable { this.killWorker(); pending.reject( new Error( - l10n.t('Execution timed out after {0} seconds', String(Math.round(timeoutMs / 1000))), + l10n.t('Operation timed out after {0} seconds', String(Math.round(timeoutMs / 1000))), ), ); } From ed120462c4b1fea10577061792b7f16229b62465 Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Fri, 27 Mar 2026 16:31:22 +0000 Subject: [PATCH 25/26] fix(scratchpad): pass actual duration to formatError instead of hardcoded 0 Capture startTime before evaluate() and compute elapsed time on failure. The error output panel now shows the real duration instead of 'Executed in 0ms'. --- src/commands/scratchpad/executeScratchpadCode.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/commands/scratchpad/executeScratchpadCode.ts b/src/commands/scratchpad/executeScratchpadCode.ts index 74b0bcd6..7c8400ca 100644 --- a/src/commands/scratchpad/executeScratchpadCode.ts +++ b/src/commands/scratchpad/executeScratchpadCode.ts @@ -100,6 +100,7 @@ export async function executeScratchpadCode(code: string, runMode: ScratchpadRun evaluator?.killWorker(); }); + const startTime = Date.now(); try { const result = await evaluator!.evaluate(connection, code, (message) => { progress.report({ message }); @@ -144,7 +145,8 @@ export async function executeScratchpadCode(code: string, runMode: ScratchpadRun // Show our own error UI before re-throwing const errorMessage = error instanceof Error ? error.message : String(error); - const formattedOutput = formatError(error, code, 0, connection); + const durationMs = Date.now() - startTime; + const formattedOutput = formatError(error, code, durationMs, connection); const errorLabel = l10n.t( '{0}/{1} — Error', From f9cfe22db6a6af89fb8a0cd7b98525ab9392bcdb Mon Sep 17 00:00:00 2001 From: Tomasz Naumowicz Date: Fri, 27 Mar 2026 17:01:21 +0000 Subject: [PATCH 26/26] perf(scratchpad): use partial Fisher-Yates for randomSample Only perform count (100) swaps instead of shuffling the entire array. Same output distribution, simpler loop bounds. --- src/commands/scratchpad/executeScratchpadCode.ts | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/commands/scratchpad/executeScratchpadCode.ts b/src/commands/scratchpad/executeScratchpadCode.ts index 7c8400ca..98bb02a2 100644 --- a/src/commands/scratchpad/executeScratchpadCode.ts +++ b/src/commands/scratchpad/executeScratchpadCode.ts @@ -263,13 +263,13 @@ function feedResultToSchemaStore(result: ExecutionResult, connection: Scratchpad } /** - * Fisher–Yates shuffle-based random sample of `count` items from `array`. - * Returns a new array of length `count` with randomly selected items. + * Partial Fisher–Yates random sample of `count` items from `array`. + * Only performs `count` swaps instead of shuffling the entire array. */ function randomSample(array: T[], count: number): T[] { const copy = [...array]; - for (let i = copy.length - 1; i > 0; i--) { - const j = Math.floor(Math.random() * (i + 1)); + for (let i = 0; i < count; i++) { + const j = i + Math.floor(Math.random() * (copy.length - i)); [copy[i], copy[j]] = [copy[j], copy[i]]; } return copy.slice(0, count);