From 396935d54ddbd16bd09de3dda538d06a0d9d6857 Mon Sep 17 00:00:00 2001 From: kapelame <168134658+kapelame@users.noreply.github.com> Date: Mon, 14 Sep 2026 11:30:04 -0400 Subject: [PATCH] fix(cloud): wait for the download writer to close before cleanup --- packages/cli/src/cloud/download.test.ts | 60 ++++++++++++++++++++++++- packages/cli/src/cloud/download.ts | 7 ++- 2 files changed, 65 insertions(+), 2 deletions(-) diff --git a/packages/cli/src/cloud/download.test.ts b/packages/cli/src/cloud/download.test.ts index 4c5b4908e8..e526e87b8e 100644 --- a/packages/cli/src/cloud/download.test.ts +++ b/packages/cli/src/cloud/download.test.ts @@ -1,5 +1,6 @@ import { chmodSync, + createWriteStream, lstatSync, mkdtempSync, readFileSync, @@ -11,9 +12,14 @@ import { } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; -import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { downloadToFile } from "./download.js"; +vi.mock("node:fs", async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, createWriteStream: vi.fn(actual.createWriteStream) }; +}); + function makeBytesFetch(bytes: Uint8Array, headers: Record = {}): typeof fetch { return (async () => new Response(new Blob([bytes as unknown as BlobPart]), { @@ -187,6 +193,58 @@ describe("cloud/download", () => { ); }); + it("waits for a pending file open to close before cleaning up a failed download", async () => { + const actual = await vi.importActual("node:fs"); + const dest = join(dir, "out.mp4"); + writeFileSync(dest, "previous render"); + let releaseOpen = () => {}; + let writerClosed = Promise.resolve(); + const openStarted = new Promise((resolveOpen) => { + vi.mocked(createWriteStream).mockImplementationOnce((path, options) => { + const writer = actual.createWriteStream(path, { + ...(typeof options === "object" ? options : { encoding: options }), + fs: { + open: (...args: Parameters) => { + releaseOpen = () => actual.open(...args); + resolveOpen(); + }, + write: actual.write, + close: actual.close, + }, + }); + writerClosed = new Promise((resolveClose) => writer.once("close", resolveClose)); + return writer; + }); + }); + const body = new ReadableStream({ + start(controller) { + controller.error(new Error("network interrupted")); + }, + }); + let settled = false; + let failure: unknown; + const downloading = downloadToFile("https://example/x", dest, { + fetchImpl: (async () => new Response(body)) as typeof fetch, + }).then( + () => { + settled = true; + }, + (error: unknown) => { + failure = error; + settled = true; + }, + ); + await openStarted; + await new Promise((resolve) => setImmediate(resolve)); + const settledBeforeOpen = settled; + releaseOpen(); + await Promise.all([downloading, writerClosed]); + expect(settledBeforeOpen).toBe(false); + expect(failure).toEqual(new Error("network interrupted")); + expect(readFileSync(dest, "utf8")).toBe("previous render"); + expect(readdirSync(dir)).toEqual(["out.mp4"]); + }); + it.skipIf(process.platform === "win32")( "preserves existing and dangling output symlinks", async () => { diff --git a/packages/cli/src/cloud/download.ts b/packages/cli/src/cloud/download.ts index 6bf3a5e646..70b0ac5f98 100644 --- a/packages/cli/src/cloud/download.ts +++ b/packages/cli/src/cloud/download.ts @@ -74,7 +74,10 @@ export async function downloadToFile( const stage = mkdtempSync(join(dirname(destination), ".hf-download-")); const stagedFile = join(stage, "download"); let bytes = 0; + let writerClosed: Promise | undefined; try { + const writer = createWriteStream(stagedFile, { flags: "wx" }); + writerClosed = new Promise((resolveClose) => writer.once("close", resolveClose)); await pipeline( async function* () { for await (const chunk of res.body as unknown as AsyncIterable) { @@ -89,7 +92,7 @@ export async function downloadToFile( ); } }, - createWriteStream(stagedFile, { flags: "wx" }), + writer, { signal: options.signal }, ); options.signal?.throwIfAborted(); @@ -100,6 +103,8 @@ export async function downloadToFile( if (options.signal?.aborted && reason instanceof Error) throw reason; throw error; } finally { + // An early source error can reject pipeline before the writer's async open finishes. + await writerClosed; rmSync(stage, { recursive: true, force: true }); } return { path: destPath, bytes };