Misc

src/kernels/raw/session/rawKernel.node.ts

@@ -83,9 +83,7 @@ export class RawKernel implements Kernel.IKernelConnection {
     public clone(
         options?: Pick<Kernel.IKernelConnection.IOptions, 'clientId' | 'username' | 'handleComms'>
     ): Kernel.IKernelConnection {
-        return createRawKernel(this.kernelProcess, options?.clientId || this.clientId, (msg) =>
-            this.anyMessage.emit(msg)
-        );
+        return createRawKernel(this.kernelProcess, options?.clientId || this.clientId);
     }
 
     public async shutdown(): Promise<void> {
@@ -271,25 +269,21 @@ export class RawKernel implements Kernel.IKernelConnection {
 
 let nonSerializingKernel: typeof import('@jupyterlab/services/lib/kernel/default');
 
-export function createRawKernel(
-    kernelProcess: IKernelProcess,
-    clientId: string,
-    onAnyMessage: (msg: Kernel.IAnyMessageArgs) => void
-): RawKernel {
+export function createRawKernel(kernelProcess: IKernelProcess, clientId: string): RawKernel {
     const jupyterLab = require('@jupyterlab/services') as typeof import('@jupyterlab/services'); // NOSONAR
     const jupyterLabSerialize =
         require('@jupyterlab/services/lib/kernel/serialize') as typeof import('@jupyterlab/services/lib/kernel/serialize'); // NOSONAR
 
     // Dummy websocket we give to the underlying real kernel
     let socketInstance: any;
+    let rawKernel: RawKernel;
     class RawSocketWrapper extends RawSocket {
         constructor() {
-            super(
-                kernelProcess.connection,
-                jupyterLabSerialize.serialize,
-                jupyterLabSerialize.deserialize,
-                onAnyMessage
-            );
+            super(kernelProcess.connection, jupyterLabSerialize.serialize, jupyterLabSerialize.deserialize, (msg) => {
+                if (rawKernel) {
+                    rawKernel?.anyMessage.emit({ direction: 'send', msg });
+                }
+            });
             socketInstance = this;
         }
     }
@@ -320,5 +314,6 @@ export function createRawKernel(
     });
 
     // Use this real kernel in result.
-    return new RawKernel(realKernel, socketInstance, kernelProcess);
+    rawKernel = new RawKernel(realKernel, socketInstance, kernelProcess);
+    return rawKernel;
 }

src/kernels/raw/session/rawSession.node.ts

@@ -75,9 +75,7 @@ export class RawSession implements ISessionWithSocket {
         this._clientID = uuid();
 
         // Connect our kernel and hook up status changes
-        this._kernel = createRawKernel(kernelProcess, this._clientID, (msg: Kernel.IAnyMessageArgs) =>
-            this._anyMessage.emit(msg)
-        );
+        this._kernel = createRawKernel(kernelProcess, this._clientID);
         this._kernel.statusChanged.connect(this.onKernelStatus, this);
         this._kernel.iopubMessage.connect(this.onIOPubMessage, this);
         this._kernel.connectionStatusChanged.connect(this.onKernelConnectionStatus, this);

src/kernels/raw/session/rawSocket.node.ts

@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 import { IDisposable } from '@fluentui/react';
-import type { Kernel, KernelMessage } from '@jupyterlab/services';
+import type { KernelMessage } from '@jupyterlab/services';
 import * as wireProtocol from '@nteract/messaging/lib/wire-protocol';
 import * as uuid from 'uuid/v4';
 import * as WebSocketWS from 'ws';
@@ -49,7 +49,7 @@ export class RawSocket implements IWebSocketLike, IKernelSocket, IDisposable {
         private connection: IKernelConnection,
         private serialize: (msg: KernelMessage.IMessage) => string | ArrayBuffer,
         private deserialize: (data: ArrayBuffer | string) => KernelMessage.IMessage,
-        private readonly onAnyMessage: (msg: Kernel.IAnyMessageArgs) => void
+        private readonly onAnyMessage: (msg: KernelMessage.IMessage) => void
     ) {
         // Setup our ZMQ channels now
         this.channels = this.generateChannels(connection);

src/test/datascience/kernelLauncher.vscode.test.ts

@@ -22,7 +22,6 @@ import { CancellationTokenSource, PortAutoForwardAction } from 'vscode';
 import { createRawKernel } from '../../kernels/raw/session/rawKernel.node';
 import { IKernelConnection, IKernelLauncher } from '../../kernels/raw/types';
 import { IJupyterKernelSpec } from '../../kernels/types';
-import { noop } from '../../platform/common/utils/misc';
 use(chaiAsPromised);
 
 const test_Timeout = 30_000;
@@ -116,7 +115,7 @@ suite('DataScience - Kernel Launcher', () => {
         assert.isOk<IKernelConnection | undefined>(kernel.connection, 'Connection not found');
 
         // Send a request to print out the env vars
-        const rawKernel = createRawKernel(kernel, uuid(), noop);
+        const rawKernel = createRawKernel(kernel, uuid());
 
         const result = await requestExecute(rawKernel, 'import os\nprint(os.getenv("TEST_VAR"))');
         assert.ok(result, 'No result returned');