feat(ng-dev): create system for registering functions to be called on command completion

ng-dev/auth/shared/ng-dev-token.ts

@@ -7,6 +7,7 @@ import {randomBytes, createCipheriv, createDecipheriv, createHash} from 'crypto'
 import {RawData, WebSocket} from 'ws';
 import {AuthenticatedGitClient} from '../../utils/git/authenticated-git-client.js';
 import {assertValidGithubConfig, getConfig} from '../../utils/config.js';
+import {registerCompletedFunction} from '../../utils/yargs.js';
 
 /** Algorithm to use for encryption. */
 const algorithm = 'aes-256-ctr';
@@ -159,6 +160,9 @@ export function configureAuthorizedGitClientWithTemporaryToken() {
         },
       });
 
+      // Close the socket whenever the command which established it is complete.
+      registerCompletedFunction(() => socket.close());
+
       // When the token is provided via the websocket message, use the token to set up
       // the AuthenticatedGitClient. The token is valid as long as the socket remains open,
       // with the server emposing a limit of 1 hour.

ng-dev/cli.ts

@@ -7,7 +7,7 @@
  * Use of this source code is governed by an MIT-style license that can be
  * found in the LICENSE file at https://angular.io/license
  */
-import yargs from 'yargs';
+import {runParserWithCompletedFunctions} from './utils/yargs.js';
 
 import {buildCaretakerParser} from './caretaker/cli.js';
 import {buildCiParser} from './ci/cli.js';
@@ -21,23 +21,25 @@ import {buildReleaseParser} from './release/cli.js';
 import {tsCircularDependenciesBuilder} from './ts-circular-dependencies/index.js';
 import {captureLogOutputForCommand} from './utils/logging.js';
 import {buildAuthParser} from './auth/cli.js';
+import {Argv} from 'yargs';
 
-yargs(process.argv.slice(2))
-  .scriptName('ng-dev')
-  .middleware(captureLogOutputForCommand)
-  .demandCommand()
-  .recommendCommands()
-  .command('auth <command>', false, buildAuthParser)
-  .command('commit-message <command>', '', buildCommitMessageParser)
-  .command('format <command>', '', buildFormatParser)
-  .command('pr <command>', '', buildPrParser)
-  .command('pullapprove <command>', '', buildPullapproveParser)
-  .command('release <command>', '', buildReleaseParser)
-  .command('ts-circular-deps <command>', '', tsCircularDependenciesBuilder)
-  .command('caretaker <command>', '', buildCaretakerParser)
-  .command('misc <command>', '', buildMiscParser)
-  .command('ngbot <command>', false, buildNgbotParser)
-  .command('ci <command>', false, buildCiParser)
-  .wrap(120)
-  .strict()
-  .parse();
+runParserWithCompletedFunctions((yargs: Argv) => {
+  return yargs
+    .scriptName('ng-dev')
+    .middleware(captureLogOutputForCommand, true)
+    .demandCommand()
+    .recommendCommands()
+    .command('auth <command>', false, buildAuthParser)
+    .command('commit-message <command>', '', buildCommitMessageParser)
+    .command('format <command>', '', buildFormatParser)
+    .command('pr <command>', '', buildPrParser)
+    .command('pullapprove <command>', '', buildPullapproveParser)
+    .command('release <command>', '', buildReleaseParser)
+    .command('ts-circular-deps <command>', '', tsCircularDependenciesBuilder)
+    .command('caretaker <command>', '', buildCaretakerParser)
+    .command('misc <command>', '', buildMiscParser)
+    .command('ngbot <command>', false, buildNgbotParser)
+    .command('ci <command>', false, buildCiParser)
+    .wrap(120)
+    .strict();
+});

ng-dev/format/cli.ts

@@ -6,6 +6,7 @@
  * found in the LICENSE file at https://angular.io/license
  */
 import {Argv} from 'yargs';
+import {registerCompletedFunction} from '../utils/yargs.js';
 import {AllFilesModule} from './all.js';
 import {ChangedModule} from './changed.js';
 import {FilesModule} from './files.js';

ng-dev/utils/logging.ts

@@ -133,8 +133,10 @@ const LOG_LEVEL_COLUMNS = 7;
  * response is executed.
  */
 export async function captureLogOutputForCommand(argv: Arguments) {
+  // TODO(josephperrott): remove this guard against running multiple times after
+  //   https://github.com/yargs/yargs/issues/2223 is fixed
   if (FILE_LOGGING_ENABLED) {
-    throw Error('`captureLogOutputForCommand` cannot be called multiple times');
+    return;
   }
 
   const repoDir = determineRepoBaseDirFromCwd();

ng-dev/utils/ng-dev-service.ts

@@ -28,6 +28,9 @@ const firebaseConfig = {
   appId: '1:823469418460:web:009b51c93132b218761119',
 };
 
+/** Whether or not the middleware has already been run. */
+let ngDevServiceMiddlewareHasRun = false;
+
 /**
  * Sets up middleware to ensure that configuration and setup is completed for commands which
  *  require the ng-dev service
@@ -54,6 +57,13 @@ export async function useNgDevService<T>(
       })
       .middleware(
         async (args: Arguments<T & {githubToken: string | null; githubEscapeHatch: boolean}>) => {
+          // TODO(josephperrott): remove this guard against running multiple times after
+          //   https://github.com/yargs/yargs/issues/2223 is fixed
+          if (ngDevServiceMiddlewareHasRun) {
+            return;
+          }
+          ngDevServiceMiddlewareHasRun = true;
+
           initializeApp(firebaseConfig);
           await restoreNgTokenFromDiskIfValid();
 
@@ -82,6 +92,7 @@ export async function useNgDevService<T>(
           Log.log('  yarn ng-dev auth login');
           throw new Error('The user is not logged in');
         },
+        true,
       )
   );
 }

ng-dev/utils/yargs.ts

@@ -0,0 +1,34 @@
+import yargs, {Arguments, Argv} from 'yargs';
+
+// A function to be called when the command completes.
+type CompletedFn = (err: Error | null) => Promise<void> | void;
+
+/** List of functions to be called upon command completion. */
+const completedFunctions: CompletedFn[] = [];
+
+/** Register a function to be called when the command completes. */
+export function registerCompletedFunction(fn: CompletedFn) {
+  completedFunctions.push(fn);
+}
+
+/**
+ * Run the yargs process, as configured by the supplied function, calling a set of completion
+ * functions after the command completes.
+ */
+export async function runParserWithCompletedFunctions(applyConfiguration: (argv: Argv) => Argv) {
+  let err: Error | null = null;
+  try {
+    await applyConfiguration(yargs(process.argv.slice(2)))
+      .exitProcess(false)
+      .parse();
+  } catch (e) {
+    err = e as Error;
+    if ([undefined, 0].includes(process.exitCode)) {
+      process.exitCode = 1;
+    }
+  } finally {
+    for (const completedFunc of completedFunctions) {
+      await completedFunc(err);
+    }
+  }
+}