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/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,39 @@
+import yargs, {Arguments, Argv} from 'yargs';
+import {Log} from './logging.js';
+
+// 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);
+}
+
+/** Error to be thrown when yargs completes without running a command. */
+export class YargsError extends Error {
+  constructor() {
+    super('Error is parse or validation of command');
+  }
+}
+
+/**
+ * Run the yargs process, as configured by the supplied function, calling a set of completion
+ * functions after the command completes.
+ */
+export function runParserWithCompletedFunctions(applyConfiguration: (argv: Argv) => Argv) {
+  applyConfiguration(yargs([])).parse(
+    process.argv.slice(2),
+    async (err: Error | null, _: Arguments, output: string) => {
+      if (output) {
+        err = new YargsError();
+        Log.log(output);
+      }
+      for (const completedFunc of completedFunctions) {
+        await completedFunc(err);
+      }
+    },
+  );
+}