Keep track of CLI clean up functions on context (#29273)

I noticed a couple bugs while playing around locally:
* If `npx convex dev --local` errored, it wouldn't clean up the local backend process
* If I did `npx convex dev --local` for project A, killed it, did `npx convex dev --local` for project B, then went back to A, it would error because it thought a deployment for project A was still running on the old ports.

The big change here is keeping track of cleanup functions on `ctx`, which also meant making the `oneoffCtx` into a class instead of a singleton.

The cleanup handlers get run on `flushAndExit` (which happens when we error or otherwise force exit the process), and also on SIGINT for `npx convex dev`

GitOrigin-RevId: d6701a88a1eaeeb7d1ff90a541e0f17547bff65a

scripts/bundle-server.mjs

@@ -14,8 +14,9 @@ const out = [];
 // Only bundle "setup.ts" from `udf/_system`.
 const udfDir = process.argv[2];
 const setupPath = path.join(udfDir, "setup.ts");
+const ctx = oneoffContext();
 const setupBundles = (
-  await bundle(oneoffContext, process.argv[2], [setupPath], true, "browser")
+  await bundle(ctx, process.argv[2], [setupPath], true, "browser")
 ).modules;
 if (setupBundles.length !== 1) {
   throw new Error("Got more than one setup bundle?");
@@ -26,19 +27,9 @@ for (const systemDir of systemDirs) {
   if (path.basename(systemDir) !== "_system") {
     throw new Error(`Refusing to bundle non-system directory ${systemDir}`);
   }
-  const entryPoints = await entryPointsByEnvironment(
-    oneoffContext,
-    systemDir,
-    false,
-  );
+  const entryPoints = await entryPointsByEnvironment(ctx, systemDir, false);
   const bundles = (
-    await bundle(
-      oneoffContext,
-      systemDir,
-      entryPoints.isolate,
-      false,
-      "browser",
-    )
+    await bundle(ctx, systemDir, entryPoints.isolate, false, "browser")
   ).modules;
   out.push(...bundles);
 }

src/bundler/context.ts

@@ -45,6 +45,8 @@ export interface Context {
     errForSentry?: any;
     printedMessage: string | null;
   }): Promise<never>;
+  registerCleanup(fn: () => Promise<void>): string;
+  removeCleanup(handle: string): (() => Promise<void>) | null;
 }
 
 async function flushAndExit(exitCode: number, err?: any) {
@@ -64,24 +66,45 @@ export type OneoffCtx = Context & {
   flushAndExit: (exitCode: number, err?: any) => Promise<never>;
 };
 
-export const oneoffContext: OneoffCtx = {
-  fs: nodeFs,
-  deprecationMessagePrinted: false,
-  spinner: undefined,
-  async crash(args: {
+class OneoffContextImpl {
+  private _cleanupFns: Record<string, () => Promise<void>> = {};
+  public fs: Filesystem = nodeFs;
+  public deprecationMessagePrinted: boolean = false;
+  public spinner: Ora | undefined = undefined;
+  crash = async (args: {
     exitCode: number;
     errorType?: ErrorType;
     errForSentry?: any;
     printedMessage: string | null;
-  }) {
+  }) => {
     if (args.printedMessage !== null) {
-      logFailure(oneoffContext, args.printedMessage);
+      logFailure(this, args.printedMessage);
     }
-    return await flushAndExit(args.exitCode, args.errForSentry);
-  },
-  flushAndExit,
-};
+    return await this.flushAndExit(args.exitCode, args.errForSentry);
+  };
+  flushAndExit = async (exitCode: number, err?: any) => {
+    logVerbose(this, "Flushing and exiting");
+    const fns = Object.values(this._cleanupFns);
+    logVerbose(this, `Running ${fns.length} cleanup functions`);
+    for (const fn of fns) {
+      await fn();
+    }
+    logVerbose(this, "All cleanup functions ran");
+    return flushAndExit(exitCode, err);
+  };
+  registerCleanup(fn: () => Promise<void>) {
+    const handle = Math.random().toString(36).slice(2);
+    this._cleanupFns[handle] = fn;
+    return handle;
+  }
+  removeCleanup(handle: string) {
+    const value = this._cleanupFns[handle];
+    delete this._cleanupFns[handle];
+    return value ?? null;
+  }
+}
 
+export const oneoffContext: () => OneoffCtx = () => new OneoffContextImpl();
 // console.error before it started being red by default in Node v20
 function logToStderr(...args: unknown[]) {
   process.stderr.write(`${format(...args)}\n`);

src/bundler/index.test.ts

@@ -35,21 +35,11 @@ test("bundle function is present", () => {
 
 test("bundle finds JavaScript functions", async () => {
   const fixtureDir = dirname + "/test_fixtures/js/project01";
-  const entryPoints = await entryPointsByEnvironment(
-    oneoffContext,
-    fixtureDir,
-    false,
-  );
+  const ctx = oneoffContext();
+  const entryPoints = await entryPointsByEnvironment(ctx, fixtureDir, false);
   const bundles = sorted(
-    (
-      await bundle(
-        oneoffContext,
-        fixtureDir,
-        entryPoints.isolate,
-        false,
-        "browser",
-      )
-    ).modules,
+    (await bundle(ctx, fixtureDir, entryPoints.isolate, false, "browser"))
+      .modules,
     (b) => b.path,
   ).filter((bundle) => !bundle.path.includes("_deps"));
   expect(bundles).toHaveLength(2);
@@ -116,23 +106,23 @@ test("returns true when multiple imports and httpRouter is imported", async () =
 test("bundle warns about https.js|ts at top level", async () => {
   const fixtureDir = dirname + "/test_fixtures/js/project_with_https";
   const logSpy = vi.spyOn(process.stderr, "write");
-  await entryPoints(oneoffContext, fixtureDir, false);
+  await entryPoints(oneoffContext(), fixtureDir, false);
   expect(logSpy).toHaveBeenCalledWith(expect.stringContaining("https"));
 });
 
 test("bundle does not warn about https.js|ts which is not at top level", async () => {
   const fixtureDir =
     dirname + "/test_fixtures/js/project_with_https_not_at_top_level";
   const logSpy = vi.spyOn(process.stderr, "write");
-  await entryPoints(oneoffContext, fixtureDir, false);
+  await entryPoints(oneoffContext(), fixtureDir, false);
   expect(logSpy).toHaveBeenCalledTimes(0);
 });
 
 test("bundle does not warn about https.js|ts which does not import httpRouter", async () => {
   const fixtureDir =
     dirname + "/test_fixtures/js/project_with_https_without_router";
   const logSpy = vi.spyOn(process.stderr, "write");
-  await entryPoints(oneoffContext, fixtureDir, false);
+  await entryPoints(oneoffContext(), fixtureDir, false);
   expect(logSpy).toHaveBeenCalledTimes(0);
 });
 

src/cli/auth.ts

@@ -2,7 +2,7 @@ import { Command, Option } from "@commander-js/extra-typings";
 import { oneoffContext } from "../bundler/context.js";
 
 const list = new Command("list").action(async () => {
-  const ctx = oneoffContext;
+  const ctx = oneoffContext();
   await ctx.crash({
     exitCode: 1,
     errorType: "fatal",
@@ -13,7 +13,7 @@ const list = new Command("list").action(async () => {
 });
 
 const rm = new Command("remove").action(async () => {
-  const ctx = oneoffContext;
+  const ctx = oneoffContext();
   await ctx.crash({
     exitCode: 1,
     errorType: "fatal",
@@ -27,7 +27,7 @@ const add = new Command("add")
   .addOption(new Option("--identity-provider-url <url>").hideHelp())
   .addOption(new Option("--application-id <applicationId>").hideHelp())
   .action(async () => {
-    const ctx = oneoffContext;
+    const ctx = oneoffContext();
     await ctx.crash({
       exitCode: 1,
       errorType: "fatal",

src/cli/codegen.ts

@@ -35,7 +35,7 @@ export const codegen = new Command("codegen")
     ).hideHelp(),
   )
   .action(async (options) => {
-    const ctx = oneoffContext;
+    const ctx = oneoffContext();
 
     await runCodegen(ctx, {
       dryRun: !!options.dryRun,

src/cli/configure.ts

@@ -47,7 +47,6 @@ import { promptOptions, promptString } from "./lib/utils/prompts.js";
 type DeploymentCredentials = {
   url: string;
   adminKey: string;
-  cleanupHandle: (() => Promise<void>) | null;
 };
 
 /**
@@ -79,15 +78,14 @@ export async function deploymentCredentialsOrConfigure(
 ): Promise<
   DeploymentCredentials & {
     deploymentName?: DeploymentName;
-    cleanupHandle: null | (() => Promise<void>);
   }
 > {
   if (cmdOptions.url !== undefined && cmdOptions.adminKey !== undefined) {
     const credentials = await handleManuallySetUrlAndAdminKey(ctx, {
       url: cmdOptions.url,
       adminKey: cmdOptions.adminKey,
     });
-    return { ...credentials, cleanupHandle: null };
+    return { ...credentials };
   }
   const { projectSlug, teamSlug } = await selectProject(
     ctx,
@@ -103,7 +101,6 @@ export async function deploymentCredentialsOrConfigure(
     deploymentName,
     deploymentUrl: url,
     adminKey,
-    cleanupHandle,
   } = await ensureDeploymentProvisioned(ctx, {
     teamSlug,
     projectSlug,
@@ -117,7 +114,7 @@ export async function deploymentCredentialsOrConfigure(
     deploymentType: deploymentOptions.kind,
   });
 
-  return { deploymentName, url, adminKey, cleanupHandle };
+  return { deploymentName, url, adminKey };
 }
 
 async function handleManuallySetUrlAndAdminKey(
@@ -429,7 +426,6 @@ async function ensureDeploymentProvisioned(
         );
       return {
         ...credentials,
-        cleanupHandle: null,
         onActivity: null,
       };
     }

src/cli/convexExport.ts

@@ -46,7 +46,7 @@ export const convexExport = new Command("export")
   .addDeploymentSelectionOptions(actionDescription("Export data from"))
   .showHelpAfterError()
   .action(async (options) => {
-    const ctx = oneoffContext;
+    const ctx = oneoffContext();
 
     const deploymentSelection = deploymentSelectionFromOptions(options);
 

src/cli/convexImport.ts

@@ -86,7 +86,7 @@ export const convexImport = new Command("import")
   .argument("<path>", "Path to the input file")
   .showHelpAfterError()
   .action(async (filePath, options, command) => {
-    const ctx = oneoffContext;
+    const ctx = oneoffContext();
 
     if (command.args.length > 1) {
       return await ctx.crash({

src/cli/dashboard.ts

@@ -21,7 +21,7 @@ export const dashboard = new Command("dashboard")
   .addDeploymentSelectionOptions(actionDescription("Open the dashboard for"))
   .showHelpAfterError()
   .action(async (options) => {
-    const ctx = oneoffContext;
+    const ctx = oneoffContext();
 
     const deploymentSelection = deploymentSelectionFromOptions(options);
     const { deploymentName } = await fetchDeploymentCredentialsProvisionProd(

src/cli/data.ts

@@ -54,7 +54,7 @@ export const data = new Command("data")
   .addDeploymentSelectionOptions(actionDescription("Inspect the database in"))
   .showHelpAfterError()
   .action(async (tableName, options) => {
-    const ctx = oneoffContext;
+    const ctx = oneoffContext();
     const deploymentSelection = deploymentSelectionFromOptions(options);
 
     const {

src/cli/deploy.ts

@@ -114,7 +114,7 @@ export const deploy = new Command("deploy")
   .addOption(new Option("--live-component-sources").hideHelp())
   .showHelpAfterError()
   .action(async (cmdOptions) => {
-    const ctx = oneoffContext;
+    const ctx = oneoffContext();
 
     storeAdminKeyEnvVar(cmdOptions.adminKey);
     const configuredDeployKey = readAdminKeyFromEnvVar() ?? null;
@@ -232,11 +232,10 @@ async function deployToNewPreviewDeployment(
     debugBundlePath: options.debugBundlePath,
     codegen: options.codegen === "enable",
     url: previewUrl,
-    cleanupHandle: null,
     liveComponentSources: false,
   };
   showSpinner(ctx, `Deploying to ${previewUrl}...`);
-  await runPush(oneoffContext, pushOptions);
+  await runPush(ctx, pushOptions);
   logFinishedStep(ctx, `Deployed Convex functions to ${previewUrl}`);
 
   if (options.previewRun !== undefined) {
@@ -329,14 +328,13 @@ async function deployToExistingDeployment(
     codegen: options.codegen === "enable",
     url,
     writePushRequest: options.writePushRequest,
-    cleanupHandle: null,
     liveComponentSources: !!options.liveComponentSources,
   };
   showSpinner(
     ctx,
     `Deploying to ${url}...${options.dryRun ? " [dry run]" : ""}`,
   );
-  await runPush(oneoffContext, pushOptions);
+  await runPush(ctx, pushOptions);
   logFinishedStep(
     ctx,
     `${

src/cli/deployments.ts

@@ -14,7 +14,7 @@ type Deployment = {
 export const deployments = new Command("deployments")
   .description("List deployments associated with a project")
   .action(async () => {
-    const ctx = oneoffContext;
+    const ctx = oneoffContext();
     const { projectConfig: config } = await readProjectConfig(ctx);
 
     const url = `teams/${config.team}/projects/${config.project}/deployments`;

src/cli/dev.ts

@@ -115,7 +115,11 @@ export const dev = new Command("dev")
   .addOption(new Option("--live-component-sources").hideHelp())
   .showHelpAfterError()
   .action(async (cmdOptions) => {
-    const ctx = oneoffContext;
+    const ctx = oneoffContext();
+    process.on("SIGINT", async () => {
+      logVerbose(ctx, "Received SIGINT, cleaning up...");
+      await ctx.flushAndExit(-2);
+    });
 
     if (cmdOptions.debugBundlePath !== undefined && !cmdOptions.once) {
       return await ctx.crash({
@@ -175,19 +179,6 @@ export const dev = new Command("dev")
       ...cmdOptions,
       localOptions,
     });
-    let cleanupHandle = credentials.cleanupHandle;
-    process.on("SIGINT", async () => {
-      logVerbose(ctx, "Received SIGINT, cleaning up...");
-      if (cleanupHandle !== null) {
-        logVerbose(ctx, "About to run cleanup handle.");
-        // Sometimes `SIGINT` gets sent twice, so set `cleanupHandle` to null to prevent double-cleanup
-        const f = cleanupHandle;
-        cleanupHandle = null;
-        await f();
-      }
-      logVerbose(ctx, "Cleaned up. Exiting.");
-      process.exit(-2);
-    });
 
     await usageStateWarning(ctx);
 
@@ -214,6 +205,7 @@ export const dev = new Command("dev")
       ),
     );
     await Promise.race(promises);
+    await ctx.flushAndExit(0);
   });
 
 export async function watchAndPush(
@@ -296,15 +288,9 @@ export async function watchAndPush(
       stopSpinner(ctx);
     }
     if (cmdOptions.once) {
-      if (options.cleanupHandle !== null) {
-        await options.cleanupHandle();
-      }
       return;
     }
     if (pushed && cmdOptions.untilSuccess) {
-      if (options.cleanupHandle !== null) {
-        await options.cleanupHandle();
-      }
       return;
     }
     const fileSystemWatch = getFileSystemWatch(ctx, watch, cmdOptions);

src/cli/docs.ts

@@ -9,7 +9,7 @@ export const docs = new Command("docs")
   .description("Open the docs in the browser")
   .option("--no-open", "Print docs URL instead of opening it in your browser")
   .action(async (options) => {
-    const ctx = oneoffContext;
+    const ctx = oneoffContext();
     // Usually we'd call `getConfiguredDeploymentName` but in this
     // command we don't care at all if the user is in the right directory
     const configuredDeployment = getTargetDeploymentName();