From 451ae7a8e26cac4add8179edfd12f8af87ab8de2 Mon Sep 17 00:00:00 2001
From: Don Jayamanne <don.jayamanne@outlook.com>
Date: Thu, 19 Dec 2024 09:01:14 +1100
Subject: [PATCH] Support fallback of jupyter message serialization

---
 src/kernels/jupyter/jupyterUtils.ts | 36 ++++++++++++++++++++++++++---
 1 file changed, 33 insertions(+), 3 deletions(-)

diff --git a/src/kernels/jupyter/jupyterUtils.ts b/src/kernels/jupyter/jupyterUtils.ts
index 3576e4be92b..bfda8ce64d8 100644
--- a/src/kernels/jupyter/jupyterUtils.ts
+++ b/src/kernels/jupyter/jupyterUtils.ts
@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-import type { ServerConnection } from '@jupyterlab/services';
+import type { KernelMessage, ServerConnection } from '@jupyterlab/services';
 import * as path from '../../platform/vscode-path/path';
 import { ConfigurationTarget, Uri, window } from 'vscode';
 import { IJupyterConnection } from '../types';
@@ -137,7 +137,36 @@ export function createJupyterConnectionInfo(
         requestInit = { ...requestInit, agent: requestAgent };
     }
 
-    const { ServerConnection } = require('@jupyterlab/services');
+    const { ServerConnection } = require('@jupyterlab/services') as typeof import('@jupyterlab/services');
+    const { deserialize, serialize } =
+        require('@jupyterlab/services/lib/kernel/serialize') as typeof import('@jupyterlab/services/lib/kernel/serialize');
+    const { supportedKernelWebSocketProtocols } =
+        require('@jupyterlab/services/lib/kernel/messages') as typeof import('@jupyterlab/services/lib/kernel/messages');
+
+    const serializer: import('@jupyterlab/services').ServerConnection.ISettings['serializer'] = {
+        deserialize: (data: ArrayBuffer, protocol?: string) => {
+            try {
+                return deserialize(data, protocol);
+            } catch (ex) {
+                if (protocol) {
+                    return deserialize(data, '');
+                } else {
+                    return deserialize(data, supportedKernelWebSocketProtocols.v1KernelWebsocketJupyterOrg);
+                }
+            }
+        },
+        serialize: (msg: KernelMessage.IMessage, protocol?: string) => {
+            try {
+                return serialize(msg, protocol);
+            } catch (ex) {
+                if (protocol) {
+                    return serialize(msg, '');
+                } else {
+                    return serialize(msg, supportedKernelWebSocketProtocols.v1KernelWebsocketJupyterOrg);
+                }
+            }
+        }
+    };
     // This replaces the WebSocket constructor in jupyter lab services with our own implementation
     // See _createSocket here:
     // https://github.com/jupyterlab/jupyterlab/blob/cfc8ebda95e882b4ed2eefd54863bb8cdb0ab763/packages/services/src/kernel/default.ts
@@ -155,7 +184,8 @@ export function createJupyterConnectionInfo(
               ) as any),
         fetch: serverUri.fetch || requestCreator.getFetchMethod(),
         Request: requestCreator.getRequestCtor(undefined, allowUnauthorized, getAuthHeader),
-        Headers: requestCreator.getHeadersCtor()
+        Headers: requestCreator.getHeadersCtor(),
+        serializer
     };
 
     const connection: IJupyterConnection = {