Compress notebook output streams before rendering (#160667)

* Compress notebook output streams before rendering

* OOps

* Combine the buffers manually

* Address code review

* oops

* Fixes

* We can have multiple stream mimes in an output

* oops

extensions/ipynb/src/serializers.ts

@@ -7,6 +7,7 @@ import type * as nbformat from '@jupyterlab/nbformat';
 import { NotebookCell, NotebookCellData, NotebookCellKind, NotebookCellOutput } from 'vscode';
 import { CellMetadata, CellOutputMetadata } from './common';
 import { textMimeTypes } from './deserializers';
+import { compressOutputItemStreams } from './streamCompressor';
 
 const textDecoder = new TextDecoder();
 
@@ -270,21 +271,17 @@ type JupyterOutput =
 
 function convertStreamOutput(output: NotebookCellOutput): JupyterOutput {
 	const outputs: string[] = [];
-	output.items
-		.filter((opit) => opit.mime === CellOutputMimeTypes.stderr || opit.mime === CellOutputMimeTypes.stdout)
-		.map((opit) => textDecoder.decode(opit.data))
-		.forEach(value => {
-			// Ensure each line is a seprate entry in an array (ending with \n).
-			const lines = value.split('\n');
-			// If the last item in `outputs` is not empty and the first item in `lines` is not empty, then concate them.
-			// As they are part of the same line.
-			if (outputs.length && lines.length && lines[0].length > 0) {
-				outputs[outputs.length - 1] = `${outputs[outputs.length - 1]}${lines.shift()!}`;
-			}
-			for (const line of lines) {
-				outputs.push(line);
-			}
-		});
+	const compressedStream = output.items.length ? new TextDecoder().decode(compressOutputItemStreams(output.items[0].mime, output.items)) : '';
+	// Ensure each line is a separate entry in an array (ending with \n).
+	const lines = compressedStream.split('\n');
+	// If the last item in `outputs` is not empty and the first item in `lines` is not empty, then concate them.
+	// As they are part of the same line.
+	if (outputs.length && lines.length && lines[0].length > 0) {
+		outputs[outputs.length - 1] = `${outputs[outputs.length - 1]}${lines.shift()!}`;
+	}
+	for (const line of lines) {
+		outputs.push(line);
+	}
 
 	for (let index = 0; index < (outputs.length - 1); index++) {
 		outputs[index] = `${outputs[index]}\n`;

extensions/ipynb/src/streamCompressor.ts

@@ -0,0 +1,63 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import type { NotebookCellOutputItem } from 'vscode';
+
+
+/**
+ * Given a stream of individual stdout outputs, this function will return the compressed lines, escaping some of the common terminal escape codes.
+ * E.g. some terminal escape codes would result in the previous line getting cleared, such if we had 3 lines and
+ * last line contained such a code, then the result string would be just the first two lines.
+ */
+export function compressOutputItemStreams(mimeType: string, outputs: NotebookCellOutputItem[]) {
+	// return outputs.find(op => op.mime === mimeType)!.data.buffer;
+
+	const buffers: Uint8Array[] = [];
+	let startAppending = false;
+	// Pick the first set of outputs with the same mime type.
+	for (const output of outputs) {
+		if (output.mime === mimeType) {
+			if ((buffers.length === 0 || startAppending)) {
+				buffers.push(output.data);
+				startAppending = true;
+			}
+		} else if (startAppending) {
+			startAppending = false;
+		}
+	}
+	compressStreamBuffer(buffers);
+	const totalBytes = buffers.reduce((p, c) => p + c.byteLength, 0);
+	const combinedBuffer = new Uint8Array(totalBytes);
+	let offset = 0;
+	for (const buffer of buffers) {
+		combinedBuffer.set(buffer, offset);
+		offset = offset + buffer.byteLength;
+	}
+	return combinedBuffer;
+}
+const MOVE_CURSOR_1_LINE_COMMAND = `${String.fromCharCode(27)}[A`;
+const MOVE_CURSOR_1_LINE_COMMAND_BYTES = MOVE_CURSOR_1_LINE_COMMAND.split('').map(c => c.charCodeAt(0));
+const LINE_FEED = 10;
+function compressStreamBuffer(streams: Uint8Array[]) {
+	streams.forEach((stream, index) => {
+		if (index === 0 || stream.length < MOVE_CURSOR_1_LINE_COMMAND.length) {
+			return;
+		}
+
+		const previousStream = streams[index - 1];
+
+		// Remove the previous line if required.
+		const command = stream.subarray(0, MOVE_CURSOR_1_LINE_COMMAND.length);
+		if (command[0] === MOVE_CURSOR_1_LINE_COMMAND_BYTES[0] && command[1] === MOVE_CURSOR_1_LINE_COMMAND_BYTES[1] && command[2] === MOVE_CURSOR_1_LINE_COMMAND_BYTES[2]) {
+			const lastIndexOfLineFeed = previousStream.lastIndexOf(LINE_FEED);
+			if (lastIndexOfLineFeed === -1) {
+				return;
+			}
+			streams[index - 1] = previousStream.subarray(0, lastIndexOfLineFeed);
+			streams[index] = stream.subarray(MOVE_CURSOR_1_LINE_COMMAND.length);
+		}
+	});
+	return streams;
+}

src/vs/workbench/api/common/extHostTypes.ts

@@ -17,7 +17,7 @@ import { IExtensionDescription } from 'vs/platform/extensions/common/extensions'
 import { FileSystemProviderErrorCode, markAsFileSystemProviderError } from 'vs/platform/files/common/files';
 import { RemoteAuthorityResolverErrorCode } from 'vs/platform/remote/common/remoteAuthorityResolver';
 import { IRelativePatternDto } from 'vs/workbench/api/common/extHost.protocol';
-import { CellEditType, ICellPartialMetadataEdit, IDocumentMetadataEdit } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { CellEditType, ICellPartialMetadataEdit, IDocumentMetadataEdit, isTextStreamMime } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { checkProposedApiEnabled } from 'vs/workbench/services/extensions/common/extensions';
 import type * as vscode from 'vscode';
 
@@ -3517,7 +3517,8 @@ export class NotebookCellOutput {
 		for (let i = 0; i < items.length; i++) {
 			const item = items[i];
 			const normalMime = normalizeMimeType(item.mime);
-			if (!seen.has(normalMime)) {
+			// We can have multiple text stream mime types in the same output.
+			if (!seen.has(normalMime) || isTextStreamMime(normalMime)) {
 				seen.add(normalMime);
 				continue;
 			}

src/vs/workbench/contrib/notebook/browser/view/renderers/backLayerWebView.ts

@@ -37,7 +37,7 @@ import { NOTEBOOK_WEBVIEW_BOUNDARY } from 'vs/workbench/contrib/notebook/browser
 import { preloadsScriptStr } from 'vs/workbench/contrib/notebook/browser/view/renderers/webviewPreloads';
 import { transformWebviewThemeVars } from 'vs/workbench/contrib/notebook/browser/view/renderers/webviewThemeMapping';
 import { MarkupCellViewModel } from 'vs/workbench/contrib/notebook/browser/viewModel/markupCellViewModel';
-import { CellUri, INotebookRendererInfo, NotebookSetting, RendererMessagingSpec } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { CellUri, INotebookRendererInfo, isTextStreamMime, NotebookSetting, RendererMessagingSpec } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { INotebookKernel } from 'vs/workbench/contrib/notebook/common/notebookKernelService';
 import { IScopedRendererMessaging } from 'vs/workbench/contrib/notebook/common/notebookRendererMessagingService';
 import { INotebookService } from 'vs/workbench/contrib/notebook/common/notebookService';
@@ -46,6 +46,7 @@ import { WebviewWindowDragMonitor } from 'vs/workbench/contrib/webview/browser/w
 import { IEditorGroupsService } from 'vs/workbench/services/editor/common/editorGroupsService';
 import { IWorkbenchEnvironmentService } from 'vs/workbench/services/environment/common/environmentService';
 import { FromWebviewMessage, IAckOutputHeight, IClickedDataUrlMessage, ICodeBlockHighlightRequest, IContentWidgetTopRequest, IControllerPreload, ICreationContent, ICreationRequestMessage, IFindMatch, IMarkupCellInitialization, RendererMetadata, ToWebviewMessage } from './webviewMessages';
+import { compressOutputItemStreams } from 'vs/workbench/contrib/notebook/browser/view/renderers/stdOutErrorPreProcessor';
 
 export interface ICachedInset<K extends ICommonCellInfo> {
 	outputId: string;
@@ -1278,12 +1279,14 @@ var requirejs = (function() {
 		let updatedContent: ICreationContent | undefined = undefined;
 		if (content.type === RenderOutputType.Extension) {
 			const output = content.source.model;
-			const first = output.outputs.find(op => op.mime === content.mimeType)!;
+			const firstBuffer = isTextStreamMime(content.mimeType) ?
+				compressOutputItemStreams(content.mimeType, output.outputs) :
+				output.outputs.find(op => op.mime === content.mimeType)!.data.buffer;
 			updatedContent = {
 				type: RenderOutputType.Extension,
 				outputId: outputCache.outputId,
-				mimeType: first.mime,
-				valueBytes: first.data.buffer,
+				mimeType: content.mimeType,
+				valueBytes: firstBuffer,
 				metadata: output.metadata,
 			};
 		}

src/vs/workbench/contrib/notebook/browser/view/renderers/stdOutErrorPreProcessor.ts

@@ -0,0 +1,56 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { VSBuffer } from 'vs/base/common/buffer';
+import type { IOutputItemDto } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+
+
+/**
+ * Given a stream of individual stdout outputs, this function will return the compressed lines, escaping some of the common terminal escape codes.
+ * E.g. some terminal escape codes would result in the previous line getting cleared, such if we had 3 lines and
+ * last line contained such a code, then the result string would be just the first two lines.
+ */
+export function compressOutputItemStreams(mimeType: string, outputs: IOutputItemDto[]) {
+	const buffers: Uint8Array[] = [];
+	let startAppending = false;
+
+	// Pick the first set of outputs with the same mime type.
+	for (const output of outputs) {
+		if (output.mime === mimeType) {
+			if ((buffers.length === 0 || startAppending)) {
+				buffers.push(output.data.buffer);
+				startAppending = true;
+			}
+		} else if (startAppending) {
+			startAppending = false;
+		}
+	}
+	compressStreamBuffer(buffers);
+	return VSBuffer.concat(buffers.map(buffer => VSBuffer.wrap(buffer))).buffer;
+}
+const MOVE_CURSOR_1_LINE_COMMAND = `${String.fromCharCode(27)}[A`;
+const MOVE_CURSOR_1_LINE_COMMAND_BYTES = MOVE_CURSOR_1_LINE_COMMAND.split('').map(c => c.charCodeAt(0));
+const LINE_FEED = 10;
+function compressStreamBuffer(streams: Uint8Array[]) {
+	streams.forEach((stream, index) => {
+		if (index === 0 || stream.length < MOVE_CURSOR_1_LINE_COMMAND.length) {
+			return;
+		}
+
+		const previousStream = streams[index - 1];
+
+		// Remove the previous line if required.
+		const command = stream.subarray(0, MOVE_CURSOR_1_LINE_COMMAND.length);
+		if (command[0] === MOVE_CURSOR_1_LINE_COMMAND_BYTES[0] && command[1] === MOVE_CURSOR_1_LINE_COMMAND_BYTES[1] && command[2] === MOVE_CURSOR_1_LINE_COMMAND_BYTES[2]) {
+			const lastIndexOfLineFeed = previousStream.lastIndexOf(LINE_FEED);
+			if (lastIndexOfLineFeed === -1) {
+				return;
+			}
+			streams[index - 1] = previousStream.subarray(0, lastIndexOfLineFeed);
+			streams[index] = stream.subarray(MOVE_CURSOR_1_LINE_COMMAND.length);
+		}
+	});
+	return streams;
+}

src/vs/workbench/contrib/notebook/common/notebookCommon.ts

@@ -947,3 +947,11 @@ export interface NotebookExtensionDescription {
 	readonly id: ExtensionIdentifier;
 	readonly location: UriComponents | undefined;
 }
+
+/**
+ * Whether the provided mime type is a text streamn like `stdout`, `stderr`.
+ */
+export function isTextStreamMime(mimeType: string) {
+	return ['application/vnd.code.notebook.stdout', 'application/x.notebook.stdout', 'application/x.notebook.stream', 'application/vnd.code.notebook.stderr', 'application/x.notebook.stderr'].includes(mimeType);
+}
+