Kernel post-initialization stability improvements

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<boolean>
+    ): 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<boolean> }
     ) {
         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');
             }

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<boolean>
     ) {
         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<boolean>) {
+        return new CodeExecution(code, extensionId, validate);
     }
     public async start(session: IKernelSession) {
         this.session = session;

src/kernels/execution/types.ts

@@ -29,6 +29,7 @@ export interface ICodeExecution {
     executionId: string;
     code: string;
     result: Promise<void>;
+    validate?: () => Promise<boolean>;
     onRequestSent: Event<void>;
     onRequestAcknowledged: Event<void>;
     onDidEmitOutput: Event<NotebookCellOutput>;

src/kernels/kernel.ts

@@ -196,6 +196,10 @@ abstract class BaseKernel implements IBaseKernel {
     public get restarting() {
         return this._restartPromise || Promise.resolve();
     }
+    private _postInitializingDeferred = createDeferred<void>();
+    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;
         });

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<void>;
             executionAcknowledged: EventEmitter<void>;
         },
-        token: CancellationToken
+        token: CancellationToken,
+        validate?: () => Promise<boolean>
     ): AsyncGenerator<NotebookCellOutput, void, unknown> {
         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(

src/kernels/types.ts

@@ -357,6 +357,7 @@ export interface IBaseKernel extends IAsyncDisposable {
     readonly onRestarted: Event<void>;
     readonly onPostInitialized: Event<void>;
     readonly restarting: Promise<void>;
+    readonly postInitializing: Promise<void>;
     readonly status: KernelMessage.Status;
     readonly disposed: boolean;
     readonly disposing: boolean;
@@ -459,7 +460,8 @@ export interface INotebookKernelExecution {
             started: EventEmitter<void>;
             executionAcknowledged: EventEmitter<void>;
         },
-        token: CancellationToken
+        token: CancellationToken,
+        validate?: () => Promise<boolean>
     ): AsyncGenerator<NotebookCellOutput, void, unknown>;
     /**
      * Given the cell execution message Id and the like , this will resume the execution of a cell from a detached state.

src/standalone/api/kernels/kernel.ts

@@ -315,7 +315,6 @@ class WrappedKernelPerExtension extends DisposableBase implements Kernel {
     }
 
     async *executeCode(code: string, token: CancellationToken): AsyncGenerator<Output, void, unknown> {
-        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<string, (data?: string) => Promise<string | undefined>>,
         token: CancellationToken
     ): AsyncGenerator<Output, void, unknown> {
-        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,18 @@ class WrappedKernelPerExtension extends DisposableBase implements Kernel {
         );
 
         try {
-            for await (const output of kernelExecution.executeCode(code, this.extensionId, events, token)) {
+            let hasAccess: boolean = false;
+            for await (const output of kernelExecution.executeCode(code, this.extensionId, events, token, async () => {
+                try {
+                    // Validate access before execution
+                    await this.checkAccess();
+                    hasAccess = true;
+                    return true;
+                } catch (e) {
+                    hasAccess = false;
+                    return false;
+                }
+            })) {
                 trackDisplayDataForExtension(this.extensionId, this.kernel.session, output);
                 output.items.forEach((output) => mimeTypes.add(output.mime));
                 if (handlers && hasChatOutput(output)) {
@@ -452,6 +461,9 @@ class WrappedKernelPerExtension extends DisposableBase implements Kernel {
                     yield output;
                 }
             }
+            if (!hasAccess) {
+                throw new KernelAPIAccessRevoked();
+            }
         } finally {
             dispose(disposables);
         }