From 0b74e71a33acf41446f89693530f3cbcfc5ee0af Mon Sep 17 00:00:00 2001 From: WuYongJie <549727955@qq.com> Date: Thu, 5 Nov 2026 12:04:11 +0800 Subject: [PATCH] Retry background removal result reads without recreating tasks --- lib/ai/rmbg/service.ts | 49 ++++++++++++++++++++++-------- lib/request-error.ts | 9 ++++++ tests/routes/rmbg-delivery.test.ts | 39 +++++++++++++++++++++++- 3 files changed, 84 insertions(+), 13 deletions(-) diff --git a/lib/ai/rmbg/service.ts b/lib/ai/rmbg/service.ts index 078482b03..bfd28cf4d 100644 --- a/lib/ai/rmbg/service.ts +++ b/lib/ai/rmbg/service.ts @@ -148,6 +148,18 @@ const withTimeout = async ( } }; +async function readAosResult(read: (signal: AbortSignal) => Promise, timeoutMs: number): Promise { + for (let attempt = 0; ; attempt++) { + try { + return await withTimeout(read, timeoutMs); + } catch (error) { + const name = error instanceof Error ? error.name : ""; + if (attempt >= 2 || (name !== "TypeError" && name !== "AbortError")) throw error; + await sleep(1000 * (attempt + 1)); + } + } +} + function getAosFriendlyErrorMessage(detail: string): string | null { const resolutionMatch = detail.match(/resolution excess,\s*width=(\d+),\s*height=(\d+)/i); if (resolutionMatch) { @@ -449,17 +461,23 @@ async function processWithAosSegmentation( await sleep(1000); stage = "poll"; - const pollResponse = await withTimeout( - (signal) => - fetch(taskUrl, { + const pollResponse = await readAosResult( + async (signal) => { + const response = await fetch(taskUrl, { method: "GET", headers, signal, - }), + }); + const payload = await response.json().catch((error) => { + if (error instanceof Error && error.name !== "SyntaxError") throw error; + return null; + }); + return { ok: response.ok, status: response.status, statusText: response.statusText, payload }; + }, AOS_POLL_TIMEOUT_MS ); - const pollPayload = await pollResponse.json().catch(() => null); + const pollPayload = pollResponse.payload; if (!pollResponse.ok || !pollPayload?.data) { logger.warn("aos_task_poll_failed", { hasPayload: pollPayload != null, @@ -482,12 +500,14 @@ async function processWithAosSegmentation( taskId, }); stage = "download"; - const imageResponse = await withTimeout( - (signal) => - fetch(imageUrl, { + const imageResponse = await readAosResult( + async (signal) => { + const response = await fetch(imageUrl, { method: "GET", signal, - }), + }); + return { ok: response.ok, status: response.status, statusText: response.statusText, buffer: response.ok ? Buffer.from(await response.arrayBuffer()) : null }; + }, AOS_DOWNLOAD_TIMEOUT_MS ); @@ -502,7 +522,7 @@ async function processWithAosSegmentation( }; } - const resultBuffer = Buffer.from(await imageResponse.arrayBuffer()); + const resultBuffer = imageResponse.buffer!; return { success: true, imageBuffer: resultBuffer, @@ -535,8 +555,13 @@ async function processWithAosSegmentation( : `AOS结果下载超时，任务已完成（任务 ${providerTaskId}）`; return { success: false, error: message }; } - logger.warn("aos_provider_request_error", error); - return { success: false, error: error instanceof Error ? error.message : "未知错误" }; + logger.warn("aos_provider_request_error", error, { requestId: config.requestId, stage, taskId: providerTaskId }); + const detail = error instanceof Error ? error.message : "未知错误"; + return { success: false, error: stage === "create" + ? `AOS创建任务连接失败，结果尚未确认：${detail}` + : stage === "poll" + ? `AOS查询任务连接失败，结果尚未确认（任务 ${providerTaskId}）：${detail}` + : `AOS结果下载连接失败，任务已完成（任务 ${providerTaskId}）：${detail}` }; } } diff --git a/lib/request-error.ts b/lib/request-error.ts index 0201155f9..f3f9304dc 100644 --- a/lib/request-error.ts +++ b/lib/request-error.ts @@ -234,9 +234,18 @@ export function getRequestFailureDetails( if (/AOS创建任务超时/.test(text)) { return detail("RMBG_CREATE_TIMEOUT", "抠图服务创建任务超时（120秒），结果尚未确认。", "请先查看原任务或联系管理员核对；不要连续点击，以免重复处理。", "unknown"); } + if (/AOS创建任务连接失败/.test(text)) { + return detail("RMBG_CREATE_CONNECTION_FAILED", "抠图服务创建任务时连接中断，结果尚未确认。", "请先查看原任务或联系管理员核对；不要连续点击，以免重复处理。", "unknown"); + } + if (/AOS查询任务连接失败/.test(text)) { + return detail("RMBG_POLL_CONNECTION_FAILED", "抠图服务查询任务时连接中断，原任务可能仍在处理。", "请保留当前页面并联系管理员核对原任务，避免重新提交。", "unknown"); + } if (/AOS查询任务超时/.test(text)) { return detail("RMBG_POLL_TIMEOUT", "抠图服务查询任务超时，原任务可能仍在处理。", "请保留当前页面并联系管理员核对原任务，避免重新提交。", "unknown"); } + if (/AOS结果下载连接失败/.test(text)) { + return detail("RMBG_RESULT_DOWNLOAD_CONNECTION_FAILED", "抠图已完成，但结果图片传输中断。", "请联系管理员按请求编号取回原任务结果，避免重复抠图。", "unknown"); + } if (/AOS结果下载超时/.test(text)) { return detail("RMBG_RESULT_DOWNLOAD_TIMEOUT", "抠图已完成，但结果图片下载超时。", "请联系管理员按请求编号取回原任务结果，避免重复抠图。", "unknown"); } diff --git a/tests/routes/rmbg-delivery.test.ts b/tests/routes/rmbg-delivery.test.ts index 532e50cfc..4da52f79d 100644 --- a/tests/routes/rmbg-delivery.test.ts +++ b/tests/routes/rmbg-delivery.test.ts @@ -12,7 +12,7 @@ function rmbgLoader(fetch: typeof globalThis.fetch, timer?: typeof setTimeout) { "@/lib/db/schema": { generatedImage: {} }, "@/lib/storage/service": { uploadFile: async (name: string) => ({ url: `${cdn}/${name}` }) }, fetch, - setTimeout: timer || ((callback: () => void, ms: number) => setTimeout(callback, ms === 1000 ? 0 : ms)), + setTimeout: timer || ((callback: () => void, ms: number) => setTimeout(callback, ms <= 2000 ? 0 : ms)), }, { NEXT_PUBLIC_CANVAS_CDN_ORIGIN: cdn })("lib/ai/rmbg/service.ts"); } @@ -80,3 +80,40 @@ test("poll and result download timeouts retain their original task and distinct expect(download).toMatchObject({ code: "RMBG_RESULT_DOWNLOAD_TIMEOUT", outcome: "unknown" }); expect(download.recommendation).toContain("取回原任务结果"); }); + +test("network and body interruptions retry only the original task reads", async () => { + const image = await sharp({ create: { width: 32, height: 32, channels: 4, background: "#abcdef00" } }).png().toBuffer(); + let creates = 0, polls = 0, downloads = 0; + let networkError: Error; + const service = rmbgLoader(async (input, options) => { + if (options?.method === "POST") { creates++; return Response.json({ data: { task_id: "original-task" } }); } + if (String(input).endsWith("/original-task")) { + if (++polls === 1) throw networkError; + return Response.json({ data: { state: 1, image: "https://output.example/original.png" } }); + } + if (++downloads === 1) return new Response(new ReadableStream({ start(controller) { controller.error(networkError); } })); + return new Response(image, { headers: { "content-type": "image/png" } }); + }); + networkError = new service.RmbgServiceError("fetch failed"); networkError.name = "TypeError"; + const result = await service.executeRmbgJob({ userId: "test", imageBuffer: image, sourceImageUrl: `${cdn}/source.png`, skipGallerySave: true }); + expect(result.original_url).toContain(cdn); + expect({ creates, polls, downloads }).toEqual({ creates: 1, polls: 2, downloads: 2 }); +}); + +test("exhausted result reads preserve the completed task and never re-create it", async () => { + const image = await sharp({ create: { width: 16, height: 16, channels: 3, background: "white" } }).png().toBuffer(); + let creates = 0, downloads = 0; + let networkError: Error; + const service = rmbgLoader(async (input, options) => { + if (options?.method === "POST") { creates++; return Response.json({ data: { task_id: "completed-task" } }); } + if (String(input).endsWith("/completed-task")) return Response.json({ data: { state: 1, image: "https://output.example/original.png" } }); + downloads++; throw networkError; + }); + networkError = new service.RmbgServiceError("fetch failed"); networkError.name = "TypeError"; + let failure: Error | undefined; + try { await service.executeRmbgJob({ userId: "test", imageBuffer: image, sourceImageUrl: `${cdn}/source.png`, skipGallerySave: true }); } + catch (error) { failure = error as Error; } + expect({ creates, downloads }).toEqual({ creates: 1, downloads: 3 }); + expect(failure?.message).toContain("completed-task"); + expect(getRequestFailureDetails(failure)).toMatchObject({ code: "RMBG_RESULT_DOWNLOAD_CONNECTION_FAILED", outcome: "unknown" }); +});