diff --git a/src/kernels/execution/cellExecutionQueue.ts b/src/kernels/execution/cellExecutionQueue.ts index baf0f2097e4..eaa58f23db6 100644 --- a/src/kernels/execution/cellExecutionQueue.ts +++ b/src/kernels/execution/cellExecutionQueue.ts @@ -65,8 +65,13 @@ export class CellExecutionQueue implements Disposable { /** * Queue the code for execution & start processing it immediately. */ - public queueCode(code: string, extensionId: string, token: CancellationToken): ICodeExecution { - const item = this.enqueue({ code, extensionId, token }); + public queueCode( + code: string, + extensionId: string, + token: CancellationToken, + validate?: () => Promise + ): ICodeExecution { + const item = this.enqueue({ code, extensionId, token, validate }); return item as ICodeExecution; } private enqueue( @@ -79,7 +84,7 @@ export class CellExecutionQueue implements Disposable { }; codeOverride?: string; } - | { code: string; extensionId: string; token: CancellationToken } + | { code: string; extensionId: string; token: CancellationToken; validate?: () => Promise } ) { let executionItem: ICellExecution | ICodeExecution; if ('cell' in options) { @@ -96,8 +101,8 @@ export class CellExecutionQueue implements Disposable { traceCellMessage(cell, 'User queued cell for execution'); } else { - const { code, extensionId, token } = options; - const codeExecution = CodeExecution.fromCode(code, extensionId); + const { code, extensionId, token, validate } = options; + const codeExecution = CodeExecution.fromCode(code, extensionId, validate); executionItem = codeExecution; this.disposables.push(codeExecution); this.queueOfItemsToExecute.push(codeExecution); @@ -205,6 +210,17 @@ export class CellExecutionQueue implements Disposable { // This way we don't accidentally end up queueing the same cell again (we know its in the queue). const itemToExecute = this.queueOfItemsToExecute[0]; this.lastCellExecution = itemToExecute; + + // Perform any async validations here to check if the code should be executed. + // We want to respect the order in which executes are called, but in some cases + // async validations need to be processed so we defer them until the queue is being run. + if (itemToExecute.type === 'code' && itemToExecute.validate) { + const isValid = await itemToExecute.validate(); + if (!isValid) { + continue; + } + } + if (itemToExecute.type === 'cell') { traceCellMessage(itemToExecute.cell, 'Before Execute individual cell'); } diff --git a/src/kernels/execution/codeExecution.ts b/src/kernels/execution/codeExecution.ts index 033f1d1103b..0ce8e84846b 100644 --- a/src/kernels/execution/codeExecution.ts +++ b/src/kernels/execution/codeExecution.ts @@ -61,7 +61,8 @@ export class CodeExecution implements ICodeExecution, IDisposable { public readonly executionId: string; private constructor( public readonly code: string, - public readonly extensionId: string + public readonly extensionId: string, + public readonly validate?: () => Promise ) { let executionId = extensionIdsPerExtension.get(extensionId) || 0; executionId += 1; @@ -70,8 +71,8 @@ export class CodeExecution implements ICodeExecution, IDisposable { this.disposables.push(this._onDidEmitOutput); } - public static fromCode(code: string, extensionId: string) { - return new CodeExecution(code, extensionId); + public static fromCode(code: string, extensionId: string, validate?: () => Promise) { + return new CodeExecution(code, extensionId, validate); } public async start(session: IKernelSession) { this.session = session; diff --git a/src/kernels/execution/types.ts b/src/kernels/execution/types.ts index 2a8a763d1f1..8bdc4daaf2e 100644 --- a/src/kernels/execution/types.ts +++ b/src/kernels/execution/types.ts @@ -29,6 +29,7 @@ export interface ICodeExecution { executionId: string; code: string; result: Promise; + validate?: () => Promise; onRequestSent: Event; onRequestAcknowledged: Event; onDidEmitOutput: Event; diff --git a/src/kernels/kernel.ts b/src/kernels/kernel.ts index a63893ad1fb..ecb72ddfd17 100644 --- a/src/kernels/kernel.ts +++ b/src/kernels/kernel.ts @@ -196,6 +196,10 @@ abstract class BaseKernel implements IBaseKernel { public get restarting() { return this._restartPromise || Promise.resolve(); } + private _postInitializingDeferred = createDeferred(); + public get postInitializing() { + return this._postInitializingDeferred.promise; + } constructor( public readonly id: string, public readonly uri: Uri, @@ -258,9 +262,10 @@ abstract class BaseKernel implements IBaseKernel { // If we started and the UI is no longer disabled (ie., a user executed a cell) // then we can signal that the kernel was created and can be used by third-party extensions. // We also only want to fire off a single event here. - if (!options?.disableUI && !this._postInitializedOnStart) { + if (!this.startupUI.disableUI && !this._postInitializedOnStart) { this._onPostInitialized.fire(); this._postInitializedOnStart = true; + this._postInitializingDeferred.resolve(); } return result; }); diff --git a/src/kernels/kernelExecution.ts b/src/kernels/kernelExecution.ts index 772ddb2f05a..c53d73d02f7 100644 --- a/src/kernels/kernelExecution.ts +++ b/src/kernels/kernelExecution.ts @@ -166,6 +166,10 @@ export class NotebookKernelExecution implements INotebookKernelExecution { const sessionPromise = this.kernel.restarting.then(() => this.kernel.start(new DisplayOptions(false))); traceCellMessage(cell, `NotebookKernelExecution.executeCell (3), ${getDisplayPath(cell.notebook.uri)}`); + + // Wait for the kernel to complete post initialization before queueing the cell in case + // we need to allow extensions to run code before the initial user-triggered execution + await this.kernel.postInitializing; const executionQueue = this.getOrCreateCellExecutionQueue(cell.notebook, sessionPromise); executionQueue.queueCell(cell, codeOverride); let success = true; @@ -193,7 +197,8 @@ export class NotebookKernelExecution implements INotebookKernelExecution { started: EventEmitter; executionAcknowledged: EventEmitter; }, - token: CancellationToken + token: CancellationToken, + validate?: () => Promise ): AsyncGenerator { const stopWatch = new StopWatch(); // If we're restarting, wait for it to finish @@ -208,7 +213,7 @@ export class NotebookKernelExecution implements INotebookKernelExecution { result = CodeExecution.fromCode(code, extensionId); void sessionPromise.then((session) => result.start(session)); } else { - result = executionQueue.queueCode(code, extensionId, token); + result = executionQueue.queueCode(code, extensionId, token, validate); } if (extensionId !== JVSC_EXTENSION_ID) { logger.trace( diff --git a/src/kernels/types.ts b/src/kernels/types.ts index b9bc32cf8a4..83baff3ac41 100644 --- a/src/kernels/types.ts +++ b/src/kernels/types.ts @@ -357,6 +357,7 @@ export interface IBaseKernel extends IAsyncDisposable { readonly onRestarted: Event; readonly onPostInitialized: Event; readonly restarting: Promise; + readonly postInitializing: Promise; readonly status: KernelMessage.Status; readonly disposed: boolean; readonly disposing: boolean; @@ -459,7 +460,8 @@ export interface INotebookKernelExecution { started: EventEmitter; executionAcknowledged: EventEmitter; }, - token: CancellationToken + token: CancellationToken, + validate?: () => Promise ): AsyncGenerator; /** * Given the cell execution message Id and the like , this will resume the execution of a cell from a detached state. diff --git a/src/standalone/api/kernels/kernel.ts b/src/standalone/api/kernels/kernel.ts index efb1cb962d7..e2cfef32dc3 100644 --- a/src/standalone/api/kernels/kernel.ts +++ b/src/standalone/api/kernels/kernel.ts @@ -315,7 +315,6 @@ class WrappedKernelPerExtension extends DisposableBase implements Kernel { } async *executeCode(code: string, token: CancellationToken): AsyncGenerator { - await this.checkAccess(); for await (const output of this.executeCodeInternal(code, undefined, token)) { yield output; } @@ -325,7 +324,6 @@ class WrappedKernelPerExtension extends DisposableBase implements Kernel { handlers: Record Promise>, token: CancellationToken ): AsyncGenerator { - await this.checkAccess(); const allowedList = ['ms-vscode.dscopilot-agent', JVSC_EXTENSION_ID]; if (!allowedList.includes(this.extensionId.toLowerCase())) { throw new Error(`Proposed API is not supported for extension ${this.extensionId}`); @@ -434,7 +432,15 @@ class WrappedKernelPerExtension extends DisposableBase implements Kernel { ); try { - for await (const output of kernelExecution.executeCode(code, this.extensionId, events, token)) { + for await (const output of kernelExecution.executeCode(code, this.extensionId, events, token, async () => { + try { + // Validate access before execution + await this.checkAccess(); + return true; + } catch (e) { + return false; + } + })) { trackDisplayDataForExtension(this.extensionId, this.kernel.session, output); output.items.forEach((output) => mimeTypes.add(output.mime)); if (handlers && hasChatOutput(output)) {