From c346e82c22f0635df01186fbc868539f678b8ffa Mon Sep 17 00:00:00 2001
From: Kaizen Conroy <36202692+kaizencc@users.noreply.github.com>
Date: Fri, 31 Jan 2025 06:28:02 -0500
Subject: [PATCH] chore(toolkit): option to strip color and style (#33244)

Closes [#33237](https://github.com/aws/aws-cdk/issues/33237) along with #33243.

```ts
new Toolkit({ ioHost, color: false });
```

Can strip color and style by specifying the property when creating the `Toolkit`.

### Checklist
- [x] My code adheres to the [CONTRIBUTING GUIDE](https://github.com/aws/aws-cdk/blob/main/CONTRIBUTING.md) and [DESIGN GUIDELINES](https://github.com/aws/aws-cdk/blob/main/docs/DESIGN_GUIDELINES.md)

----

*By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*
---
 .../toolkit/lib/api/io/private/logger.ts      | 29 +++++++++++++++++++
 .../@aws-cdk/toolkit/lib/toolkit/toolkit.ts   | 15 +++++++++-
 .../toolkit/test/toolkit/toolkit.test.ts      | 21 ++++++++++++++
 3 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/packages/@aws-cdk/toolkit/lib/api/io/private/logger.ts b/packages/@aws-cdk/toolkit/lib/api/io/private/logger.ts
index 6d7e6976968ed..2df2b444cd0ee 100644
--- a/packages/@aws-cdk/toolkit/lib/api/io/private/logger.ts
+++ b/packages/@aws-cdk/toolkit/lib/api/io/private/logger.ts
@@ -7,6 +7,10 @@ import type { ToolkitAction } from '../../../toolkit';
 import { formatSdkLoggerContent } from '../../aws-cdk';
 import type { IIoHost } from '../io-host';
 
+/**
+ * An IoHost wrapper that adds the given action to an actionless message before
+ * sending the message to the given IoHost
+ */
 export function withAction(ioHost: IIoHost, action: ToolkitAction) {
   return {
     notify: async <T>(msg: Omit<IoMessage<T>, 'action'>) => {
@@ -24,6 +28,31 @@ export function withAction(ioHost: IIoHost, action: ToolkitAction) {
   };
 }
 
+/**
+ * An IoHost wrapper that strips out ANSI colors and styles from the message before
+ * sending the message to the given IoHost
+ */
+export function withoutColor(ioHost: IIoHost): IIoHost {
+  return {
+    notify: async <T>(msg: IoMessage<T>) => {
+      await ioHost.notify({
+        ...msg,
+        message: stripColor(msg.message),
+      });
+    },
+    requestResponse: async <T, U>(msg: IoRequest<T, U>) => {
+      return ioHost.requestResponse({
+        ...msg,
+        message: stripColor(msg.message),
+      });
+    },
+  };
+}
+
+function stripColor(msg: string): string {
+  return msg.replace(/[\u001b\u009b][[()#;?]*(?:[0-9]{1,4}(?:;[0-9]{0,4})*)?[0-9A-ORZcf-nqry=><]/g, '');
+}
+
 /**
  * An IoHost wrapper that strips out emojis from the message before
  * sending the message to the given IoHost
diff --git a/packages/@aws-cdk/toolkit/lib/toolkit/toolkit.ts b/packages/@aws-cdk/toolkit/lib/toolkit/toolkit.ts
index a9e6dfa0f0544..315966b98177d 100644
--- a/packages/@aws-cdk/toolkit/lib/toolkit/toolkit.ts
+++ b/packages/@aws-cdk/toolkit/lib/toolkit/toolkit.ts
@@ -19,7 +19,7 @@ import { CachedCloudAssemblySource, IdentityCloudAssemblySource, StackAssembly,
 import { ALL_STACKS, CloudAssemblySourceBuilder } from '../api/cloud-assembly/private';
 import { ToolkitError } from '../api/errors';
 import { IIoHost, IoMessageCode, IoMessageLevel } from '../api/io';
-import { asSdkLogger, withAction, Timer, confirm, error, highlight, info, success, warn, ActionAwareIoHost, debug, result, withoutEmojis } from '../api/io/private';
+import { asSdkLogger, withAction, Timer, confirm, error, highlight, info, success, warn, ActionAwareIoHost, debug, result, withoutEmojis, withoutColor } from '../api/io/private';
 
 /**
  * The current action being performed by the CLI. 'none' represents the absence of an action.
@@ -54,6 +54,16 @@ export interface ToolkitOptions {
    */
   emojis?: boolean;
 
+  /**
+   * Whether to allow ANSI colors and formatting in IoHost messages.
+   * Setting this value to `falsez enforces that no color or style shows up
+   * in messages sent to the IoHost.
+   * Setting this value to true is a no-op; it is equivalent to the default.
+   *
+   * @default - detects color from the TTY status of the IoHost
+   */
+  color?: boolean;
+
   /**
    * Configuration options for the SDK.
    */
@@ -102,6 +112,9 @@ export class Toolkit extends CloudAssemblySourceBuilder implements AsyncDisposab
     if (props.emojis === false) {
       ioHost = withoutEmojis(ioHost);
     }
+    if (props.color === false) {
+      ioHost = withoutColor(ioHost);
+    }
     this.ioHost = ioHost;
   }
 
diff --git a/packages/@aws-cdk/toolkit/test/toolkit/toolkit.test.ts b/packages/@aws-cdk/toolkit/test/toolkit/toolkit.test.ts
index 437b0260b0ff5..4680f75873a59 100644
--- a/packages/@aws-cdk/toolkit/test/toolkit/toolkit.test.ts
+++ b/packages/@aws-cdk/toolkit/test/toolkit/toolkit.test.ts
@@ -5,6 +5,7 @@
  *  - Source Builders: Tests for the Cloud Assembly Source Builders are in `test/api/cloud-assembly/source-builder.test.ts`
  */
 
+import * as chalk from 'chalk';
 import { Toolkit } from '../../lib';
 import { TestIoHost } from '../_helpers';
 
@@ -27,3 +28,23 @@ test('emojis can be stripped from message', async () => {
     message: 'Smile123',
   }));
 });
+
+test('color can be stripped from message', async () => {
+  const ioHost = new TestIoHost();
+  const toolkit = new Toolkit({ ioHost, color: false });
+
+  await toolkit.ioHost.notify({
+    message: chalk.red('RED') + chalk.bold('BOLD') + chalk.blue('BLUE'),
+    action: 'deploy',
+    level: 'info',
+    code: 'CDK_TOOLKIT_I0000',
+    time: new Date(),
+  });
+
+  expect(ioHost.notifySpy).toHaveBeenCalledWith(expect.objectContaining({
+    action: 'deploy',
+    level: 'info',
+    code: 'CDK_TOOLKIT_I0000',
+    message: 'REDBOLDBLUE',
+  }));
+});