From a57d7f12ae37e66f6c56297eab1d12a3a4ccf92c Mon Sep 17 00:00:00 2001
From: Remco Haszing <remcohaszing@gmail.com>
Date: Mon, 18 Sep 2023 09:33:21 +0200
Subject: [PATCH] Patch console methods if --stdio is used (#1301)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Patch console methods if --stdio is used

Closes #1299

* Fix stylistic issue

* Implement more console methods

This implements `console.count`, `console.counrReset`, `console.debug`,
and `console.trace`. It also adds support for `console.dir` options.

---------

Co-authored-by: Dirk Bäumer <dirkb@microsoft.com>
---
 server/src/node/main.ts | 75 ++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 74 insertions(+), 1 deletion(-)

diff --git a/server/src/node/main.ts b/server/src/node/main.ts
index 476432962..649b94252 100644
--- a/server/src/node/main.ts
+++ b/server/src/node/main.ts
@@ -4,6 +4,8 @@
  * ------------------------------------------------------------------------------------------ */
 /// <reference path="../../typings/thenable.d.ts" />
 
+import { inspect } from 'node:util';
+
 import * as Is from '../common/utils/is';
 import { Connection, _, _Connection, Features, WatchDog, createConnection as createCommonConnection } from '../common/server';
 
@@ -190,6 +192,7 @@ function _createConnection<PConsole = _, PTracer = _, PTelemetry = _, PClient =
 	options?: ConnectionStrategy | ConnectionOptions,
 	factories?: Features<PConsole, PTracer, PTelemetry, PClient, PWindow, PWorkspace, PLanguages, PNotebooks>,
 ): _Connection<PConsole, PTracer, PTelemetry, PClient, PWindow, PWorkspace, PLanguages, PNotebooks> {
+	let stdio = false;
 	if (!input && !output && process.argv.length > 2) {
 		let port: number | undefined = void 0;
 		let pipeName: string | undefined = void 0;
@@ -201,6 +204,7 @@ function _createConnection<PConsole = _, PTracer = _, PTelemetry = _, PClient =
 				output = new IPCMessageWriter(process);
 				break;
 			} else if (arg === '--stdio') {
+				stdio = true;
 				input = process.stdin;
 				output = process.stdout;
 				break;
@@ -255,7 +259,76 @@ function _createConnection<PConsole = _, PTracer = _, PTelemetry = _, PClient =
 
 	const connectionFactory = (logger: Logger): ProtocolConnection => {
 		const result = createProtocolConnection(input as any, output as any, logger, options);
+		if(stdio) {
+			patchConsole(logger);
+		}
 		return result;
 	};
 	return createCommonConnection(connectionFactory, watchDog, factories);
-}
\ No newline at end of file
+}
+
+function patchConsole(logger: Logger): undefined {
+	function serialize(args: unknown[]) {
+		return args.map(arg => typeof arg === 'string' ? arg : inspect(arg)).join(' ');
+	}
+
+	const counters = new Map<string, number>();
+
+	console.assert = function assert(assertion, ...args) {
+		if (assertion) {
+			return;
+		}
+		if (args.length === 0) {
+			logger.error('Assertion failed');
+		} else {
+			const [message, ...rest] = args;
+			logger.error(`Assertion failed: ${message} ${serialize(rest)}`);
+		}
+	};
+
+	console.count = function count(label = 'default') {
+		const message = String(label);
+		let counter = counters.get(message) ?? 0;
+		counter += 1;
+		counters.set(message, counter);
+		logger.log(`${message}: ${message}`);
+	};
+
+	console.countReset = function countReset(label) {
+		if(label === undefined) {
+			counters.clear();
+		} else {
+			counters.delete(String(label));
+		}
+	};
+
+	console.debug = function debug(...args) {
+		logger.log(serialize(args));
+	};
+
+	console.dir = function dir(arg, options){
+		// @ts-expect-error https://github.com/DefinitelyTyped/DefinitelyTyped/pull/66626
+		logger.log(inspect(arg, options));
+	};
+
+	console.log = function log(...args) {
+		logger.log(serialize(args));
+	};
+
+	console.error = function error(...args){
+		logger.error(serialize(args));
+	};
+
+	console.trace = function trace(...args) {
+		const stack = new Error().stack!.replace(/(.+\n){2}/, '');
+		let message = 'Trace';
+		if (args.length !== 0) {
+			message += `: ${serialize(args)}`;
+		}
+		logger.log(`${message}\n${stack}`);
+	};
+
+	console.warn = function warn(...args) {
+		logger.warn(serialize(args));
+	};
+}