refactor(cli): make corked logging part of CliIoHost

packages/aws-cdk/lib/cli/cdk-toolkit.ts

@@ -47,9 +47,10 @@ import {
 import { printSecurityDiff, printStackDiff, RequireApproval } from '../diff';
 import { ResourceImporter, removeNonImportResources } from '../import';
 import { listStacks } from '../list-stacks';
-import { data, debug, error, highlight, info, success, warning, withCorkedLogging } from '../logging';
+import { data, debug, error, highlight, info, success, warning } from '../logging';
 import { ResourceMigrator } from '../migrator';
 import { deserializeStructure, obscureTemplate, serializeStructure } from '../serialize';
+import { CliIoHost } from '../toolkit/cli-io-host';
 import { ToolkitError } from '../toolkit/error';
 import { numberFromBool, partition } from '../util';
 import { formatErrorMessage } from '../util/error';
@@ -79,6 +80,11 @@ export interface CdkToolkitProps {
    */
   deployments: Deployments;
 
+  /**
+   * The CliIoHost that's used for I/O operations
+   */
+  ioHost?: CliIoHost;
+
   /**
    * Whether to be verbose
    *
@@ -136,7 +142,11 @@ export enum AssetBuildTime {
  * deploys applies them to `cloudFormation`.
  */
 export class CdkToolkit {
-  constructor(private readonly props: CdkToolkitProps) {}
+  private ioHost: CliIoHost;
+
+  constructor(private readonly props: CdkToolkitProps) {
+    this.ioHost = props.ioHost ?? CliIoHost.instance();
+  }
 
   public async metadata(stackName: string, json: boolean) {
     const stacks = await this.selectSingleStackByName(stackName);
@@ -371,6 +381,7 @@ export class CdkToolkit {
         const currentTemplate = await this.props.deployments.readCurrentTemplate(stack);
         if (printSecurityDiff(currentTemplate, stack, requireApproval)) {
           await askUserConfirmation(
+            this.ioHost,
             concurrency,
             '"--require-approval" is enabled and stack includes security-sensitive updates',
             'Do you wish to deploy these changes',
@@ -451,6 +462,7 @@ export class CdkToolkit {
                 warning(`${motivation}. Rolling back first (--force).`);
               } else {
                 await askUserConfirmation(
+                  this.ioHost,
                   concurrency,
                   motivation,
                   `${motivation}. Roll back first and then proceed with deployment`,
@@ -476,6 +488,7 @@ export class CdkToolkit {
                 warning(`${motivation}. Proceeding with regular deployment (--force).`);
               } else {
                 await askUserConfirmation(
+                  this.ioHost,
                   concurrency,
                   motivation,
                   `${motivation}. Perform a regular deployment`,
@@ -1818,11 +1831,12 @@ function buildParameterMap(
  * cannot be interactively obtained from a human at the keyboard.
  */
 async function askUserConfirmation(
+  ioHost: CliIoHost,
   concurrency: number,
   motivation: string,
   question: string,
 ) {
-  await withCorkedLogging(async () => {
+  await ioHost.withCorkedLogging(async () => {
     // only talk to user if STDIN is a terminal (otherwise, fail)
     if (!TESTING && !process.stdin.isTTY) {
       throw new ToolkitError(`${motivation}, but terminal (TTY) is not attached so we are unable to get a confirmation from the user`);

packages/aws-cdk/lib/cli/cli.ts

@@ -195,6 +195,7 @@ export async function exec(args: string[], synthesizer?: Synthesizer): Promise<n
     };
 
     const cli = new CdkToolkit({
+      ioHost,
       cloudExecutable,
       deployments: cloudFormation,
       verbose: argv.trace || argv.verbose > 0,

packages/aws-cdk/lib/logging.ts

@@ -1,84 +1,11 @@
 import * as util from 'util';
 import * as chalk from 'chalk';
-import { IoMessageLevel, IoMessage, CliIoHost, IoMessageSpecificCode, IoMessageCode, IoMessageCodeCategory, IoCodeLevel, levelPriority } from './toolkit/cli-io-host';
-
-// Corking mechanism
-let CORK_COUNTER = 0;
-const logBuffer: IoMessage<any>[] = [];
-
-/**
- * Executes a block of code with corked logging. All log messages during execution
- * are buffered and only written when all nested cork blocks complete (when CORK_COUNTER reaches 0).
- * @param block - Async function to execute with corked logging
- * @returns Promise that resolves with the block's return value
- */
-export async function withCorkedLogging<T>(block: () => Promise<T>): Promise<T> {
-  CORK_COUNTER++;
-  try {
-    return await block();
-  } finally {
-    CORK_COUNTER--;
-    if (CORK_COUNTER === 0) {
-      // Process each buffered message through notify
-      for (const ioMessage of logBuffer) {
-        void CliIoHost.instance().notify(ioMessage);
-      }
-      logBuffer.splice(0);
-    }
-  }
-}
-
-interface LogMessage {
-  /**
-   * The log level to use
-   */
-  readonly level: IoMessageLevel;
-  /**
-   * The message to log
-   */
-  readonly message: string;
-  /**
-   * Whether to force stdout
-   * @default false
-   */
-  readonly forceStdout?: boolean;
-  /**
-   * Message code of the format [CATEGORY]_[NUMBER_CODE]
-   * @pattern [A-Z]+_[0-2][0-9]{3}
-   * @default TOOLKIT_[0/1/2]000
-   */
-  readonly code: IoMessageCode;
-}
-
-/**
- * Internal core logging function that writes messages through the CLI IO host.
- * @param msg Configuration options for the log message. See  {@link LogMessage}
- */
-function log(msg: LogMessage) {
-  const ioMessage: IoMessage<undefined> = {
-    level: msg.level,
-    message: msg.message,
-    forceStdout: msg.forceStdout,
-    time: new Date(),
-    action: CliIoHost.instance().currentAction,
-    code: msg.code,
-  };
-
-  if (CORK_COUNTER > 0) {
-    if (levelPriority[msg.level] > levelPriority[CliIoHost.instance().logLevel]) {
-      return;
-    }
-    logBuffer.push(ioMessage);
-    return;
-  }
-
-  void CliIoHost.instance().notify(ioMessage);
-}
+import { IoMessageLevel, IoMessage, CliIoHost, IoMessageSpecificCode, IoMessageCode, IoMessageCodeCategory, IoCodeLevel } from './toolkit/cli-io-host';
 
 /**
  * Internal helper that processes log inputs into a consistent format.
  * Handles string interpolation, format strings, and object parameter styles.
- * Applies optional styling and prepares the final message for logging.
+ * Applies optional styling and sends the message to the IoHost.
  */
 function formatMessageAndLog(
   level: IoMessageLevel,
@@ -98,12 +25,17 @@ function formatMessageAndLog(
   // Apply style if provided
   const finalMessage = style ? style(formattedMessage) : formattedMessage;
 
-  log({
+  const ioHost = CliIoHost.instance();
+  const ioMessage: IoMessage<undefined> = {
+    time: new Date(),
+    action: ioHost.currentAction,
     level,
     message: finalMessage,
     code,
     forceStdout,
-  });
+  };
+
+  void ioHost.notify(ioMessage);
 }
 
 function getDefaultCode(level: IoMessageLevel, category: IoMessageCodeCategory = 'TOOLKIT'): IoMessageCode {

packages/aws-cdk/lib/toolkit/cli-io-host.ts

@@ -175,12 +175,17 @@ export class CliIoHost implements IIoHost {
    */
   private static _instance: CliIoHost | undefined;
 
+  // internal state for getters/setter
   private _currentAction: ToolkitAction;
   private _isCI: boolean;
   private _isTTY: boolean;
   private _logLevel: IoMessageLevel;
   private _internalIoHost?: IIoHost;
 
+  // Corked Logging
+  private corkedCounter = 0;
+  private readonly corkedLoggingBuffer: IoMessage<any>[] = [];
+
   private constructor(props: CliIoHostProps = {}) {
     this._currentAction = props.currentAction ?? 'none' as ToolkitAction;
     this._isTTY = props.isTTY ?? process.stdout.isTTY ?? false;
@@ -259,6 +264,31 @@ export class CliIoHost implements IIoHost {
     this._logLevel = level;
   }
 
+  /**
+   * Executes a block of code with corked logging. All log messages during execution
+   * are buffered and only written when all nested cork blocks complete (when CORK_COUNTER reaches 0).
+   * The corking is bound to the specific instance of the CliIoHost.
+   *
+   * @param block - Async function to execute with corked logging
+   * @returns Promise that resolves with the block's return value
+   */
+  public async withCorkedLogging<T>(block: () => Promise<T>): Promise<T> {
+    this.corkedCounter++;
+    try {
+      return await block();
+    } finally {
+      this.corkedCounter--;
+      if (this.corkedCounter === 0) {
+        // Process each buffered message through notify
+        for (const ioMessage of this.corkedLoggingBuffer) {
+          await this.notify(ioMessage);
+        }
+        // remove all buffered messages in-place
+        this.corkedLoggingBuffer.splice(0);
+      }
+    }
+  }
+
   /**
    * Notifies the host of a message.
    * The caller waits until the notification completes.
@@ -272,24 +302,20 @@ export class CliIoHost implements IIoHost {
       return;
     }
 
+    if (this.corkedCounter > 0) {
+      this.corkedLoggingBuffer.push(msg);
+      return;
+    }
+
     const output = this.formatMessage(msg);
-    const stream = this.stream(msg.level, msg.forceStdout ?? false);
-
-    return new Promise((resolve, reject) => {
-      stream.write(output, (err) => {
-        if (err) {
-          reject(err);
-        } else {
-          resolve();
-        }
-      });
-    });
+    const stream = this.selectStream(msg.level, msg.forceStdout ?? false);
+    stream.write(output);
   }
 
   /**
    * Determines which output stream to use based on log level and configuration.
    */
-  private stream(level: IoMessageLevel, forceStdout: boolean) {
+  private selectStream(level: IoMessageLevel, forceStdout: boolean) {
     // For legacy purposes all log streams are written to stderr by default, unless
     // specified otherwise, by passing `forceStdout`, which is used by the `data()` logging function, or
     // if the CDK is running in a CI environment. This is because some CI environments will immediately

packages/aws-cdk/test/api/logs/logging.test.ts

@@ -1,4 +1,4 @@
-import { data, success, highlight, error, warning, info, debug, trace, withCorkedLogging } from '../../../lib/logging';
+import { data, success, highlight, error, warning, info, debug, trace } from '../../../lib/logging';
 import { CliIoHost } from '../../../lib/toolkit/cli-io-host';
 
 describe('logging', () => {
@@ -249,65 +249,6 @@ describe('logging', () => {
     });
   });
 
-  describe('corked logging', () => {
-    test('buffers messages when corked', async () => {
-      await withCorkedLogging(async () => {
-        info('message 1');
-        info({ message: 'message 2' });
-        expect(mockStderr).not.toHaveBeenCalled();
-      });
-
-      expect(mockStderr).toHaveBeenCalledWith('message 1\n');
-      expect(mockStderr).toHaveBeenCalledWith('message 2\n');
-    });
-
-    test('handles nested corking correctly', async () => {
-      await withCorkedLogging(async () => {
-        info('outer 1');
-        await withCorkedLogging(async () => {
-          info({ message: 'inner' });
-        });
-        info({ message: 'outer 2' });
-        expect(mockStderr).not.toHaveBeenCalled();
-      });
-
-      expect(mockStderr).toHaveBeenCalledTimes(3);
-      expect(mockStderr).toHaveBeenCalledWith('outer 1\n');
-      expect(mockStderr).toHaveBeenCalledWith('inner\n');
-      expect(mockStderr).toHaveBeenCalledWith('outer 2\n');
-    });
-
-    test('handles errors in corked block while preserving buffer', async () => {
-      await expect(withCorkedLogging(async () => {
-        info('message 1');
-        throw new Error('test error');
-      })).rejects.toThrow('test error');
-
-      // The buffered message should still be printed even if the block throws
-      expect(mockStderr).toHaveBeenCalledWith('message 1\n');
-    });
-
-    test('maintains correct order with mixed log levels in corked block', async () => {
-      // Set threshold to debug to allow debug messages
-      ioHost.logLevel = 'debug';
-
-      await withCorkedLogging(async () => {
-        error('error message');
-        warning('warning message');
-        success('success message');
-        debug('debug message');
-      });
-
-      const calls = mockStderr.mock.calls.map(call => call[0]);
-      expect(calls).toEqual([
-        'error message\n',
-        'warning message\n',
-        'success message\n',
-        expect.stringMatching(/^\[\d{2}:\d{2}:\d{2}\] debug message\n$/),
-      ]);
-    });
-  });
-
   describe('CI mode behavior', () => {
     test('correctly switches between stdout and stderr based on CI mode', () => {
       ioHost.isCI = true;