diff --git a/packages/ai/src/background-removal.ts b/packages/ai/src/background-removal.ts index 78c95648..bb69f62b 100644 --- a/packages/ai/src/background-removal.ts +++ b/packages/ai/src/background-removal.ts @@ -2,6 +2,7 @@ import { randomUUID } from "node:crypto"; import { readFile, unlink, writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; import { join } from "node:path"; +import { isSafeMessageError, SafeError } from "@snapotter/shared"; import sharp from "sharp"; import { type ProgressCallback, parseStdoutJson, runPythonWithProgress } from "./bridge.js"; @@ -28,6 +29,20 @@ export function isMemoryAllocError(err: unknown): boolean { ); } +/** + * Wrap a background-removal failure in a SafeError so its message survives the + * API's Sentry scrubber, which otherwise reduces a plain Error to "Error: + * Error". The specific sidecar reason is kept as the message (callers and the + * existing tests rely on it, matching the ai-bridge behavior); an empty reason + * falls back to a constant. Errors we already author (the bridge's SafeError + * timeout/OOM) pass through unchanged so their kind is not masked. + */ +function toBgRemovalError(reason: unknown): Error { + if (isSafeMessageError(reason)) return reason; + const message = reason instanceof Error ? reason.message : String(reason ?? ""); + return new SafeError(message || "Background removal failed", { kind: "bug" }); +} + export async function removeBackground( inputBuffer: Buffer, outputDir: string, @@ -96,7 +111,7 @@ async function runAndParse( const isOom = isMemoryAllocError(err); const canFallback = isOom && options.model !== OOM_FALLBACK_MODEL; - if (!canFallback) throw err; + if (!canFallback) throw toBgRemovalError(err); onProgress?.(5, `Retrying with lighter model (${OOM_FALLBACK_MODEL})`); const fallbackOpts = { ...options, model: OOM_FALLBACK_MODEL }; @@ -107,7 +122,7 @@ async function runAndParse( ); const result = parseStdoutJson(stdout); if (!result.success) { - throw new Error(result.error || "Background removal failed"); + throw toBgRemovalError(result.error || "Background removal failed"); } return readFile(outputPath); } diff --git a/tests/unit/ai/background-removal-error.test.ts b/tests/unit/ai/background-removal-error.test.ts new file mode 100644 index 00000000..d7a613fe --- /dev/null +++ b/tests/unit/ai/background-removal-error.test.ts @@ -0,0 +1,78 @@ +import { tmpdir } from "node:os"; +import sharp from "sharp"; +import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; + +// Mock the sidecar bridge by its resolved path so background-removal's +// `./bridge.js` import receives these stubs. +const runPythonWithProgress = vi.fn(); +const parseStdoutJson = vi.fn(); +vi.mock("../../../packages/ai/src/bridge.js", () => ({ + runPythonWithProgress: (...args: unknown[]) => runPythonWithProgress(...args), + parseStdoutJson: (...args: unknown[]) => parseStdoutJson(...args), +})); + +import { isSafeMessageError, SafeError } from "@snapotter/shared"; +import { removeBackground } from "../../../packages/ai/src/background-removal.js"; + +let png: Buffer; +beforeAll(async () => { + png = await sharp({ + create: { width: 4, height: 4, channels: 4, background: { r: 0, g: 0, b: 0, alpha: 0 } }, + }) + .png() + .toBuffer(); +}); + +beforeEach(() => { + runPythonWithProgress.mockReset(); + parseStdoutJson.mockReset(); +}); + +describe("removeBackground error surfacing", () => { + it("throws a SafeError titled 'Background removal failed' with the sidecar reason in the cause", async () => { + runPythonWithProgress.mockResolvedValue({ stdout: "{}" }); + parseStdoutJson.mockReturnValue({ success: false, error: "rembg model load failed" }); + + let caught: unknown; + try { + await removeBackground(png, tmpdir(), { model: "u2net" }); + } catch (e) { + caught = e; + } + + expect(isSafeMessageError(caught)).toBe(true); + // The specific sidecar reason is preserved as the message (and survives scrubbing). + expect((caught as SafeError).message).toBe("rembg model load failed"); + expect((caught as SafeError).kind).toBe("bug"); + expect(runPythonWithProgress).toHaveBeenCalledTimes(1); + }); + + it("still retries with the lighter model on OOM, and wraps the fallback failure too", async () => { + parseStdoutJson + .mockReturnValueOnce({ success: false, error: "CUDA out of memory" }) + .mockReturnValueOnce({ success: false, error: "still failing" }); + runPythonWithProgress.mockResolvedValue({ stdout: "{}" }); + + let caught: unknown; + try { + await removeBackground(png, tmpdir(), { model: "isnet-general-use" }); + } catch (e) { + caught = e; + } + + // OOM detection still works: the fallback attempt fired (two sidecar calls). + expect(runPythonWithProgress).toHaveBeenCalledTimes(2); + expect(isSafeMessageError(caught)).toBe(true); + expect((caught as SafeError).message).toBe("still failing"); + }); + + it("passes a bridge SafeError (e.g. timeout) through unchanged, not re-wrapped as a bug", async () => { + const timeout = new SafeError("Python script timed out", { + kind: "operational", + code: "timeout", + }); + runPythonWithProgress.mockRejectedValue(timeout); + + await expect(removeBackground(png, tmpdir(), { model: "u2net" })).rejects.toBe(timeout); + }); +});