From eaba5ab7e21f0ea65879b5e8000d34f30c917794 Mon Sep 17 00:00:00 2001
From: Gert Hengeveld <info@ghengeveld.nl>
Date: Wed, 20 Sep 2023 21:59:16 +0200
Subject: [PATCH] Add experimental_abortSignal to Node API

---
 node-src/lib/tasks.ts                                |  1 +
 node-src/main.test.ts                                |  9 +++++++++
 node-src/runBuild.ts                                 | 10 ++++++++--
 node-src/types.ts                                    |  3 +++
 node-src/ui/messages/errors/buildCanceled.stories.ts |  7 +++++++
 node-src/ui/messages/errors/buildCanceled.ts         | 10 ++++++++++
 6 files changed, 38 insertions(+), 2 deletions(-)
 create mode 100644 node-src/ui/messages/errors/buildCanceled.stories.ts
 create mode 100644 node-src/ui/messages/errors/buildCanceled.ts

diff --git a/node-src/lib/tasks.ts b/node-src/lib/tasks.ts
index 62164a19e..8e46ae970 100644
--- a/node-src/lib/tasks.ts
+++ b/node-src/lib/tasks.ts
@@ -26,6 +26,7 @@ export const createTask = ({
 
     // eslint-disable-next-line no-restricted-syntax
     for (const step of steps) {
+      ctx.options.experimental_abortSignal?.throwIfAborted();
       // eslint-disable-next-line no-await-in-loop
       await step(ctx, task);
     }
diff --git a/node-src/main.test.ts b/node-src/main.test.ts
index 11ab2fae7..93d4c444f 100644
--- a/node-src/main.test.ts
+++ b/node-src/main.test.ts
@@ -415,6 +415,15 @@ it('should exit with code 0 when the current branch is skipped', async () => {
   expect(ctx.exitCode).toBe(0);
 });
 
+it('should exit with code 6 and stop the build when abortSignal is aborted', async () => {
+  const abortSignal = AbortSignal.abort(new Error('Build canceled'));
+  const ctx = getContext(['--project-token=asdf1234']);
+  ctx.extraOptions = { experimental_abortSignal: abortSignal };
+  await runBuild(ctx);
+  expect(ctx.exitCode).toBe(6);
+  expect(uploadFiles).not.toHaveBeenCalled();
+});
+
 it('calls out to npm build script passed and uploads files', async () => {
   const ctx = getContext(['--project-token=asdf1234', '--build-script-name=build-storybook']);
   await runBuild(ctx);
diff --git a/node-src/runBuild.ts b/node-src/runBuild.ts
index e55196508..4c0da2074 100644
--- a/node-src/runBuild.ts
+++ b/node-src/runBuild.ts
@@ -1,12 +1,14 @@
 import Listr from 'listr';
 
 import GraphQLClient from './io/GraphQLClient';
+import { getConfiguration } from './lib/getConfiguration';
 import getOptions from './lib/getOptions';
 import NonTTYRenderer from './lib/NonTTYRenderer';
 import { exitCodes, setExitCode } from './lib/setExitCode';
 import { rewriteErrorMessage } from './lib/utils';
 import getTasks from './tasks';
-import { Context, Options } from './types';
+import { Context } from './types';
+import buildCanceled from './ui/messages/errors/buildCanceled';
 import fatalError from './ui/messages/errors/fatalError';
 import fetchError from './ui/messages/errors/fetchError';
 import graphqlError from './ui/messages/errors/graphqlError';
@@ -15,7 +17,6 @@ import runtimeError from './ui/messages/errors/runtimeError';
 import taskError from './ui/messages/errors/taskError';
 import intro from './ui/messages/info/intro';
 import { endActivity } from './ui/components/activity';
-import { getConfiguration } from './lib/getConfiguration';
 
 export async function runBuild(ctx: Context) {
   ctx.log.info('');
@@ -65,6 +66,11 @@ export async function runBuild(ctx: Context) {
         setExitCode(ctx, exitCodes.BUILD_NO_STORIES);
         throw rewriteErrorMessage(err, missingStories(ctx));
       }
+      if (ctx.extraOptions.experimental_abortSignal?.aborted) {
+        ctx.userError = true;
+        setExitCode(ctx, exitCodes.BUILD_WAS_CANCELED);
+        throw rewriteErrorMessage(err, buildCanceled());
+      }
       throw rewriteErrorMessage(err, taskError(ctx, err));
     } finally {
       // Handle potential runtime errors from JSDOM
diff --git a/node-src/types.ts b/node-src/types.ts
index c14d27cf0..7acab5685 100644
--- a/node-src/types.ts
+++ b/node-src/types.ts
@@ -108,6 +108,9 @@ export interface Options {
 
   /** A callback that is called at the completion of each task */
   experimental_onTaskComplete?: (ctx: Context) => void;
+
+  /** An AbortSignal that terminates the build if aborted */
+  experimental_abortSignal?: AbortSignal;
 }
 
 export { Configuration };
diff --git a/node-src/ui/messages/errors/buildCanceled.stories.ts b/node-src/ui/messages/errors/buildCanceled.stories.ts
new file mode 100644
index 000000000..2030b6cbb
--- /dev/null
+++ b/node-src/ui/messages/errors/buildCanceled.stories.ts
@@ -0,0 +1,7 @@
+import buildCanceled from './buildCanceled';
+
+export default {
+  title: 'CLI/Messages/Errors',
+};
+
+export const BuildCanceled = () => buildCanceled();
diff --git a/node-src/ui/messages/errors/buildCanceled.ts b/node-src/ui/messages/errors/buildCanceled.ts
new file mode 100644
index 000000000..d853be0c2
--- /dev/null
+++ b/node-src/ui/messages/errors/buildCanceled.ts
@@ -0,0 +1,10 @@
+import chalk from 'chalk';
+import { dedent } from 'ts-dedent';
+
+import { error } from '../../components/icons';
+
+export default () =>
+  dedent(chalk`
+    ${error} {bold Build canceled}
+    The build was canceled before it completed.
+  `);