Properly pass through buffers on comm_open events (#7228)

Author: nstrayer

extensions/positron-ipywidgets/renderer/src/manager.ts

@@ -1,5 +1,5 @@
 /*---------------------------------------------------------------------------------------------
- *  Copyright (C) 2024 Posit Software, PBC. All rights reserved.
+ *  Copyright (C) 2024-2025 Posit Software, PBC. All rights reserved.
  *  Licensed under the Elastic License 2.0. See LICENSE.txt for license information.
  *--------------------------------------------------------------------------------------------*/
 
@@ -134,6 +134,9 @@ export class PositronWidgetManager extends ManagerBase implements base.IWidgetMa
 					target_name: message.target_name,
 					data: message.data as JSONObject,
 				},
+				// Sometimes these buffers are not aligned, so we need to convert them to Uint8Arrays to avoid mangled data.
+				// This is handled by the Comm class for subsequent messages.
+				buffers: message.buffers?.map(buffer => new Uint8Array(buffer.buffer)),
 				// This is expected to at least contain the backend widget protocol 'version', which
 				// should match the frontend version.
 				metadata: message.metadata as JSONObject,

extensions/positron-supervisor/src/KallichoreSession.ts

@@ -1526,7 +1526,7 @@ export class KallichoreSession implements JupyterLanguageRuntimeSession {
 	 */
 	handleJupyterMessage(data: any) {
 		// Deserialize the message buffers from base64, if any
-		if (data.buffers) {
+		if (data.buffers?.length > 0) {
 			data.buffers = data.buffers.map((b: string) => {
 				return Buffer.from(b, 'base64');
 			});

extensions/positron-supervisor/src/RuntimeMessageEmitter.ts

@@ -1,5 +1,5 @@
 /*---------------------------------------------------------------------------------------------
- *  Copyright (C) 2023-2024 Posit Software, PBC. All rights reserved.
+ *  Copyright (C) 2023-2025 Posit Software, PBC. All rights reserved.
  *  Licensed under the Elastic License 2.0. See LICENSE.txt for license information.
  *--------------------------------------------------------------------------------------------*/
 
@@ -184,6 +184,7 @@ export class RuntimeMessageEmitter {
 			target_name: data.target_name,
 			data: data.data,
 			metadata: message.metadata,
+			buffers: message.buffers,
 		} as positron.LanguageRuntimeCommOpen);
 	}
 

src/vs/workbench/contrib/positronIPyWidgets/browser/positronIPyWidgetsService.ts

@@ -1,5 +1,5 @@
 /*---------------------------------------------------------------------------------------------
- *  Copyright (C) 2023-2024 Posit Software, PBC. All rights reserved.
+ *  Copyright (C) 2023-2025 Posit Software, PBC. All rights reserved.
  *  Licensed under the Elastic License 2.0. See LICENSE.txt for license information.
  *--------------------------------------------------------------------------------------------*/
 import { Disposable, DisposableStore, toDisposable } from '../../../../base/common/lifecycle.js';
@@ -323,6 +323,7 @@ export class IPyWidgetsInstance extends Disposable {
 				target_name: client.getClientType(),
 				data: message.data,
 				metadata: message.metadata,
+				buffers: message.buffers,
 			});
 		}));
 

src/vs/workbench/contrib/positronIPyWidgets/test/browser/positronIPyWidgetsService.test.ts

@@ -27,6 +27,7 @@ import { PositronTestServiceAccessor, positronWorkbenchInstantiationService } fr
 import { INotebookRendererInfo, INotebookStaticPreloadInfo } from '../../../notebook/common/notebookCommon.js';
 import { NotebookOutputRendererInfo } from '../../../notebook/common/notebookOutputRenderer.js';
 import { ExtensionIdentifier } from '../../../../../platform/extensions/common/extensions.js';
+import { VSBuffer } from '../../../../../base/common/buffer.js';
 
 class TestNotebookService implements Partial<INotebookService> {
 	getRenderers(): INotebookRendererInfo[] {
@@ -246,6 +247,40 @@ suite('Positron - IPyWidgetsInstance constructor', () => {
 		assert.deepStrictEqual(messaging.messagesToWebview, [{ type: 'initialize_result' }]);
 	});
 
+	test('forwards comm_open with buffers when client is created via session.createClient', async () => {
+		// Create instance and clear initial messages
+		await createIPyWidgetsInstance();
+		messaging.messagesToWebview.length = 0; // Clear initialize_result message
+
+		// Make some sample client data
+		const clientId = 'new-widget-client-id';
+		const clientType = RuntimeClientType.IPyWidget;
+		const messageData = { initial_state: 'some_value' };
+		const messageMetadata = { source: 'test' };
+		const messageBuffers = [
+			VSBuffer.wrap(new Uint8Array([1, 2, 3])),
+			VSBuffer.wrap(new Uint8Array([4, 5]))
+		];
+
+		// Act: Call session.createClient to simulate client creation and trigger the event
+		await session.createClient(
+			clientType,
+			messageData,
+			messageMetadata,
+			clientId,
+			messageBuffers
+		);
+		await timeout(0);
+
+		// Sanity check: only one message should be sent.
+		assert.strictEqual(messaging.messagesToWebview.length, 1, 'Expected exactly one message (comm_open) to be sent to the webview');
+		const sentMessage = messaging.messagesToWebview[0];
+
+		// The message should be comm_open and contain the correct buffers.
+		assert.strictEqual(sentMessage.type, 'comm_open', 'Expected message type to be comm_open');
+		assert.deepStrictEqual(sentMessage.buffers, messageBuffers, 'Expected buffers to be passed correctly in the comm_open message');
+	});
+
 	test('initialized session, one ipywidget client', async () => {
 		const client = await session.createClient(RuntimeClientType.IPyWidget, {}, {}, 'test-client-id');
 
@@ -365,6 +400,7 @@ suite('Positron - IPyWidgetsInstance', () => {
 			target_name: client.getClientType(),
 			data: {},
 			metadata: {},
+			buffers: [],
 		} as ToWebviewMessage]);
 
 		// Close the client.

src/vs/workbench/services/languageRuntime/common/positronIPyWidgetsWebviewMessages.ts

@@ -1,8 +1,10 @@
 /*---------------------------------------------------------------------------------------------
- *  Copyright (C) 2024 Posit Software, PBC. All rights reserved.
+ *  Copyright (C) 2024-2025 Posit Software, PBC. All rights reserved.
  *  Licensed under the Elastic License 2.0. See LICENSE.txt for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import { VSBuffer } from '../../../../base/common/buffer.js';
+
 /**
  * Request the runtime to close a comm.
  *
@@ -99,6 +101,7 @@ export interface ICommOpenToWebview {
 	comm_id: string;
 	target_name: string;
 	data: unknown;
+	buffers?: Array<VSBuffer>;
 	metadata: unknown;
 }
 

src/vs/workbench/services/runtimeSession/test/common/testLanguageRuntimeSession.ts

@@ -13,6 +13,7 @@ import { IRuntimeClientEvent } from '../../../languageRuntime/common/languageRun
 import { TestRuntimeClientInstance } from '../../../languageRuntime/test/common/testRuntimeClientInstance.js';
 import { CancellationError } from '../../../../../base/common/errors.js';
 import { TestUiClientInstance } from '../../../languageRuntime/test/common/testUiClientInstance.js';
+import { VSBuffer } from '../../../../../base/common/buffer.js';
 
 export class TestLanguageRuntimeSession extends Disposable implements ILanguageRuntimeSession {
 	private readonly _onDidChangeRuntimeState = this._register(new Emitter<RuntimeState>());
@@ -140,7 +141,7 @@ export class TestLanguageRuntimeSession extends Disposable implements ILanguageR
 	}
 
 	async createClient(
-		type: RuntimeClientType, params: any, metadata?: any, id?: string
+		type: RuntimeClientType, params: any, metadata?: any, id?: string, buffers?: VSBuffer[]
 	): Promise<TestRuntimeClientInstance> {
 		const client = type === RuntimeClientType.Ui ?
 			new TestUiClientInstance(id ?? generateUuid()) :
@@ -163,7 +164,7 @@ export class TestLanguageRuntimeSession extends Disposable implements ILanguageR
 					parent_id: '',
 					type: LanguageRuntimeMessageType.CommOpen,
 					when: new Date().toISOString(),
-					buffers: [],
+					buffers: buffers ?? [],
 				}
 			}
 		);