fixes

src/notebooks/execution/cellExecution.ts

@@ -33,17 +33,17 @@ import { IJupyterSession, KernelConnectionMetadata, NotebookCellRunState } from
 import { getInteractiveCellMetadata } from '../../interactive-window/helpers';
 import { isCancellationError } from '../../platform/common/cancellation';
 import { activeNotebookCellExecution, CellExecutionMessageHandler } from './cellExecutionMessageHandler';
-import { CellExecutionMessageHandlerFactory } from './cellExecutionMessageHandlerFactory';
+import { CellExecutionMessageHandlerService } from './cellExecutionMessageHandlerService';
 
 export class CellExecutionFactory {
     constructor(
         private readonly controller: NotebookController,
-        private readonly factory: CellExecutionMessageHandlerFactory
+        private readonly requestListener: CellExecutionMessageHandlerService
     ) {}
 
     public create(cell: NotebookCell, metadata: Readonly<KernelConnectionMetadata>) {
         // eslint-disable-next-line @typescript-eslint/no-use-before-define
-        return CellExecution.fromCell(cell, metadata, this.controller, this.factory);
+        return CellExecution.fromCell(cell, metadata, this.controller, this.requestListener);
     }
 }
 
@@ -85,7 +85,7 @@ export class CellExecution implements IDisposable {
         public readonly cell: NotebookCell,
         private readonly kernelConnection: Readonly<KernelConnectionMetadata>,
         private readonly controller: NotebookController,
-        private readonly factory: CellExecutionMessageHandlerFactory
+        private readonly requestListener: CellExecutionMessageHandlerService
     ) {
         workspace.onDidCloseTextDocument(
             (e) => {
@@ -129,9 +129,9 @@ export class CellExecution implements IDisposable {
         cell: NotebookCell,
         metadata: Readonly<KernelConnectionMetadata>,
         controller: NotebookController,
-        factory: CellExecutionMessageHandlerFactory
+        requestListener: CellExecutionMessageHandlerService
     ) {
-        return new CellExecution(cell, metadata, controller, factory);
+        return new CellExecution(cell, metadata, controller, requestListener);
     }
     public async start(session: IJupyterSession) {
         if (this.cancelHandled) {
@@ -354,20 +354,16 @@ export class CellExecution implements IDisposable {
             traceError(`Cell execution failed without request, for cell Index ${this.cell.index}`, ex);
             return this.completedWithErrors(ex);
         }
-        this.cellExecutionHandler = this.factory.create(this.cell, {
+        this.cellExecutionHandler = this.requestListener.registerListener(this.cell, {
             kernel: session.kernel!,
             cellExecution: this.execution!,
-            request: this.request
-        });
-        this.cellExecutionHandler.onErrorHandlingExecuteRequestIOPubMessage(
-            ({ error }) => {
+            request: this.request,
+            onErrorHandlingExecuteRequestIOPubMessage: (error) => {
                 traceError(`Cell (index = ${this.cell.index}) execution completed with errors (2).`, error);
                 // If not a restart error, then tell the subscriber
                 this.completedWithErrors(error);
-            },
-            this,
-            this.disposables
-        );
+            }
+        });
 
         // WARNING: Do not dispose `request`.
         // Even after request.done & execute_reply is sent we could have more messages coming from iopub.

src/notebooks/execution/cellExecutionMessageHandlerFactory.ts → src/notebooks/execution/cellExecutionMessageHandlerService.ts

@@ -10,7 +10,7 @@ import { IDisposable, IExtensionContext } from '../../platform/common/types';
 import { CellOutputDisplayIdTracker } from './cellDisplayIdTracker';
 import { CellExecutionMessageHandler } from './cellExecutionMessageHandler';
 
-export class CellExecutionMessageHandlerFactory {
+export class CellExecutionMessageHandlerService {
     private readonly disposables: IDisposable[] = [];
     private notebook?: NotebookDocument;
     private readonly messageHandlers = new WeakMap<NotebookCell, CellExecutionMessageHandler>();
@@ -40,12 +40,13 @@ export class CellExecutionMessageHandlerFactory {
             this.notebook.getCells().forEach((cell) => this.messageHandlers.get(cell)?.dispose());
         }
     }
-    public create(
+    public registerListener(
         cell: NotebookCell,
         options: {
             kernel: Kernel.IKernelConnection;
             request: Kernel.IShellFuture<KernelMessage.IExecuteRequestMsg, KernelMessage.IExecuteReplyMsg>;
             cellExecution: NotebookCellExecution;
+            onErrorHandlingExecuteRequestIOPubMessage: (error: Error) => void;
         }
     ): CellExecutionMessageHandler {
         this.notebook = cell.notebook;

src/notebooks/execution/kernelExecution.ts

@@ -27,7 +27,7 @@ import {
 import { traceCellMessage } from '../helpers';
 import { getDisplayPath } from '../../platform/common/platform/fs-paths';
 import { getAssociatedNotebookDocument } from '../controllers/kernelSelector';
-import { CellExecutionMessageHandlerFactory } from './cellExecutionMessageHandlerFactory';
+import { CellExecutionMessageHandlerService } from './cellExecutionMessageHandlerService';
 
 /**
  * Separate class that deals just with kernel execution.
@@ -50,14 +50,14 @@ export class KernelExecution implements IDisposable {
         context: IExtensionContext,
         formatters: ITracebackFormatter[]
     ) {
-        const factory = new CellExecutionMessageHandlerFactory(
+        const requestListener = new CellExecutionMessageHandlerService(
             appShell,
             controller,
             outputTracker,
             context,
             formatters
         );
-        this.executionFactory = new CellExecutionFactory(controller, factory);
+        this.executionFactory = new CellExecutionFactory(controller, requestListener);
     }
 
     public get onPreExecute() {