fix(cli): stopping remove-background on a Mac no longer aborts the process (#6848)

* fix(cli): stopping remove-background on a Mac no longer aborts the process

* fix(cli): a quit after background removal released its session exits at once

* Revert "fix(cli): a quit after background removal released its session exits at once"

This reverts commit 3b9d6ffc24.

packages/cli/src/background-removal/inference.stop.test.ts

@@ -0,0 +1,96 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+const cleanups = vi.hoisted(() => new Set<() => Promise<void>>());
+const native = vi.hoisted(() => ({
+  frames: [] as (() => void)[],
+  released: 0,
+  releasedDuringFrame: 0,
+  inFrame: false,
+}));
+
+vi.mock("../browser/launch.js", () => ({
+  registerBrowserSessionCleanup: (cleanup: () => Promise<void>) => {
+    cleanups.add(cleanup);
+    return () => cleanups.delete(cleanup);
+  },
+}));
+vi.mock("./manager.js", () => ({
+  ensureModel: async () => "model.onnx",
+  selectProviders: () => ({ providers: ["coreml", "cpu"], label: "CoreML" }),
+}));
+vi.mock("./coremlTemp.js", () => ({
+  trackCoreMlTemp: () => ({ record: () => undefined, forget: () => undefined }),
+}));
+// A frame inside the native session ends when the test says so; releasing the session meanwhile is what aborts
+// the process on a Mac ("mutex lock failed").
+vi.mock("../utils/optionalPackages.js", () => {
+  const chain: Record<string, unknown> = {};
+  for (const step of ["resize", "raw", "toColourspace"]) chain[step] = () => chain;
+  chain.toBuffer = async () => Buffer.alloc(320 * 320 * 3);
+  return {
+    loadSharp: async () => () => chain,
+    loadOptionalPackage: async () => ({
+      Tensor: class {},
+      InferenceSession: {
+        create: async () => ({
+          inputNames: ["in"],
+          outputNames: ["out"],
+          run: () =>
+            new Promise((resolve) => {
+              native.inFrame = true;
+              native.frames.push(() => {
+                native.inFrame = false;
+                resolve({ out: { data: new Float32Array(320 * 320) } });
+              });
+            }),
+          release: async () => {
+            native.released++;
+            if (native.inFrame) native.releasedDuringFrame++;
+          },
+        }),
+      },
+    }),
+  };
+});
+
+import { createSession } from "./inference.js";
+
+const frame = (session: Awaited<ReturnType<typeof createSession>>) =>
+  session.process(Buffer.alloc(48), 4, 4);
+const settle = () => new Promise((resolve) => setTimeout(resolve, 5));
+
+beforeEach(() => {
+  cleanups.clear();
+  Object.assign(native, { frames: [], released: 0, releasedDuringFrame: 0, inFrame: false });
+});
+
+describe("a stop signal while a frame is inside the session", () => {
+  it("releases the session only after that frame has ended", async () => {
+    const session = await createSession();
+    const running = frame(session);
+    await settle();
+    const [stop] = [...cleanups];
+
+    const stopped = stop!();
+    await settle();
+    expect(native.released).toBe(0);
+
+    native.frames.shift()!();
+    await running.catch(() => undefined);
+    await stopped;
+
+    expect(native.released).toBe(1);
+    expect(native.releasedDuringFrame).toBe(0);
+  });
+
+  it("starts no frame after the stop, and closing afterwards does not release a second time", async () => {
+    const session = await createSession();
+    const [stop] = [...cleanups];
+    await stop!();
+
+    await expect(frame(session)).rejects.toThrow(/stopped/i);
+    await session.close();
+
+    expect(native.released).toBe(1);
+  });
+});

packages/cli/src/background-removal/inference.ts

@@ -9,6 +9,7 @@ import type { InferenceSession, Tensor } from "onnxruntime-node";
 import type sharpType from "sharp";
 import { loadOptionalPackage, loadSharp } from "../utils/optionalPackages.js";
 import { registerBrowserSessionCleanup } from "../browser/launch.js";
+import { drainToExit } from "../utils/commandResult.js";
 import { trackCoreMlTemp } from "./coremlTemp.js";
 import { ensureModel, selectProviders, type Device, type ModelId } from "./manager.js";
 
@@ -75,7 +76,23 @@ export async function createSession(options: CreateSessionOptions = {}): Promise
 
   const coreMlTemp = choice.label === "CoreML" ? trackCoreMlTemp() : undefined;
   // A signal stops the run through the CLI's shutdown, which releases the session so CoreML removes its temp files.
-  const unregister = coreMlTemp && registerBrowserSessionCleanup(async () => session?.release());
+  // Releasing it while a frame is inside it aborts the process, so the stop waits for that frame and refuses the next.
+  if (coreMlTemp) drainToExit();
+  let stopped = false;
+  let inFlight: Promise<unknown> | undefined;
+  let released = false;
+  const release = async () => {
+    if (released || !session) return;
+    released = true;
+    await session.release();
+  };
+  const unregister =
+    coreMlTemp &&
+    registerBrowserSessionCleanup(async () => {
+      stopped = true;
+      await inFlight;
+      await release();
+    });
   let session: InferenceSession;
   let providerUsed = choice.label;
   try {
@@ -125,38 +142,45 @@ export async function createSession(options: CreateSessionOptions = {}): Promise
     }
   };
 
+  const convert = async (rgb: Buffer, width: number, height: number, withBackground: boolean) => {
+    const tensor = await preprocess(sharp, ort, rgb, width, height, inputData);
+    const outputs = await run(tensor);
+    const output = outputs[outputName];
+    if (!output) throw new Error(`Model did not return output '${outputName}'`);
+    const expectedBytes = width * height * 4;
+    if (!rgbaBuf || rgbaBuf.length !== expectedBytes) {
+      rgbaBuf = Buffer.allocUnsafe(expectedBytes);
+    }
+    if (withBackground) {
+      if (!rgbaBgBuf || rgbaBgBuf.length !== expectedBytes) {
+        rgbaBgBuf = Buffer.allocUnsafe(expectedBytes);
+      }
+    }
+    return await postprocess(
+      sharp,
+      output,
+      rgb,
+      width,
+      height,
+      maskBuf,
+      rgbaBuf,
+      withBackground ? rgbaBgBuf : null,
+    );
+  };
+
   return {
     get provider() {
       return providerUsed;
     },
     async process(rgb, width, height, withBackground = false) {
-      const tensor = await preprocess(sharp, ort, rgb, width, height, inputData);
-      const outputs = await run(tensor);
-      const output = outputs[outputName];
-      if (!output) throw new Error(`Model did not return output '${outputName}'`);
-      const expectedBytes = width * height * 4;
-      if (!rgbaBuf || rgbaBuf.length !== expectedBytes) {
-        rgbaBuf = Buffer.allocUnsafe(expectedBytes);
-      }
-      if (withBackground) {
-        if (!rgbaBgBuf || rgbaBgBuf.length !== expectedBytes) {
-          rgbaBgBuf = Buffer.allocUnsafe(expectedBytes);
-        }
-      }
-      return await postprocess(
-        sharp,
-        output,
-        rgb,
-        width,
-        height,
-        maskBuf,
-        rgbaBuf,
-        withBackground ? rgbaBgBuf : null,
-      );
+      if (stopped) throw new Error("Background removal was stopped.");
+      const frame = convert(rgb, width, height, withBackground);
+      inFlight = frame.catch(() => undefined);
+      return await frame;
     },
     async close() {
       unregister?.();
-      await session.release();
+      await release();
       coreMlTemp?.forget();
     },
   };

packages/cli/src/cli.ts

@@ -97,6 +97,7 @@ import {
   CliUsageError,
   USAGE_EXIT_CODE,
   consumeCommandResult,
+  endProcess,
   outputWritten,
   registerRootExitCodeSanitizer,
   registerRootExitRequester,
@@ -373,7 +374,13 @@ registerRootExitRequester((exitCode) => {
     presented: true,
   }).finally(async () => {
     await outputWritten();
-    process.exit(exitCode);
+    endProcess(exitCode, {
+      exit: (code) => process.exit(code),
+      setExitCode: (code) => {
+        process.exitCode = code;
+      },
+      later: setTimeout,
+    });
   });
 });
 

packages/cli/src/utils/commandResult.test.ts

@@ -1,9 +1,11 @@
-import { afterEach, describe, expect, it } from "vitest";
+import { afterEach, describe, expect, it, vi } from "vitest";
 import {
   CliResultSignal,
   CliRuntimeError,
   CliUsageError,
   consumeCommandResult,
+  drainToExit,
+  endProcess,
   failCommand,
   failUsage,
   finishCommand,
@@ -33,3 +35,39 @@ describe("command result contract", () => {
     expect(consumeCommandResult()).toMatchObject({ exitCode: 1, kind: "runtime_error" });
   });
 });
+
+describe("ending the process", () => {
+  const host = () => {
+    const pending: (() => void)[] = [];
+    const calls = { exit: vi.fn(), setExitCode: vi.fn() };
+    const later = ((run: () => void) => {
+      pending.push(run);
+      return { unref: () => undefined };
+    }) as unknown as typeof setTimeout;
+    return { pending, calls, ...calls, later };
+  };
+
+  it("exits at once when nothing native needs tearing down", async () => {
+    vi.resetModules();
+    const fresh = await import("./commandResult.js");
+    const ending = host();
+
+    fresh.endProcess(143, ending);
+
+    expect(ending.exit).toHaveBeenCalledWith(143);
+    expect(ending.pending).toHaveLength(0);
+  });
+
+  it("after CoreML ran, leaves the exit to Node's own teardown and forces it only after a grace", () => {
+    const ending = host();
+    drainToExit();
+
+    endProcess(143, ending);
+
+    expect(ending.exit).not.toHaveBeenCalled();
+    expect(ending.setExitCode).toHaveBeenCalledWith(143);
+    expect(ending.pending).toHaveLength(1);
+    ending.pending[0]!();
+    expect(ending.exit).toHaveBeenCalledWith(143);
+  });
+});

packages/cli/src/utils/commandResult.ts

@@ -120,6 +120,25 @@ export async function outputWritten(): Promise<void> {
   await Promise.all([written(process.stdout), written(process.stderr)]);
 }
 
+let drainsToExit = false;
+const DRAIN_GRACE_MS = 8_000;
+
+export function drainToExit(): void {
+  drainsToExit = true;
+}
+
+export interface ProcessEnd {
+  exit(exitCode: number): void;
+  setExitCode(exitCode: number): void;
+  later: typeof setTimeout;
+}
+
+export function endProcess(exitCode: number, host: ProcessEnd): void {
+  if (!drainsToExit) return host.exit(exitCode);
+  host.setExitCode(exitCode);
+  host.later(() => host.exit(exitCode), DRAIN_GRACE_MS).unref();
+}
+
 /** Ask cli.ts to finalize telemetry/output and then terminate the process. */
 export function requestCliExit(exitCode = 0): void {
   if (!rootExitRequester) {