fix(cli): follow every redirect a host may answer with

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

@@ -1,12 +1,12 @@
 import { EventEmitter } from "node:events";
-import { existsSync, mkdtempSync, rmSync } from "node:fs";
+import { existsSync, mkdtempSync, readFileSync, readdirSync, rmSync } from "node:fs";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
 import { PassThrough } from "node:stream";
 import { get as httpsGet } from "node:https";
 import type { ClientRequest, IncomingMessage } from "node:http";
 import { afterEach, describe, expect, it, vi } from "vitest";
-import { downloadFile } from "./download.js";
+import { downloadFile, REDIRECT_CODES, redirectTarget } from "./download.js";
 
 const { unlinkSyncMock } = vi.hoisted(() => ({
   unlinkSyncMock: vi.fn(),
@@ -22,12 +22,61 @@ vi.mock("node:fs", async (importOriginal) => {
 const mockGet = vi.mocked(httpsGet);
 const tempDirs: string[] = [];
 
+function httpsResponse(
+  statusCode: number,
+  headers: Record<string, string> = {},
+  body?: string | ((url: string) => string),
+): typeof httpsGet {
+  return ((_url: string, callback: (response: IncomingMessage) => void) => {
+    const response = new PassThrough() as PassThrough & {
+      statusCode: number;
+      headers: Record<string, string>;
+    };
+    response.statusCode = statusCode;
+    response.headers = headers;
+    const request = new EventEmitter() as ClientRequest;
+    request.setTimeout = vi.fn();
+    callback(response as unknown as IncomingMessage);
+    if (body !== undefined) {
+      queueMicrotask(() => response.end(typeof body === "function" ? body(_url) : body));
+    }
+    return request;
+  }) as typeof httpsGet;
+}
+
 afterEach(() => {
   vi.clearAllMocks();
   for (const dir of tempDirs.splice(0)) rmSync(dir, { recursive: true, force: true });
 });
 
 describe("downloadFile", () => {
+  it("keeps concurrent partial downloads separate", async () => {
+    mockGet.mockImplementation(httpsResponse(200, {}, (url) => url));
+
+    const dir = mkdtempSync(join(tmpdir(), "hyperframes-download-"));
+    tempDirs.push(dir);
+    const dest = join(dir, "model.onnx");
+    const first = "https://example.test/first";
+    const second = "https://example.test/second";
+
+    await Promise.all([downloadFile(first, dest), downloadFile(second, dest)]);
+    expect([first, second]).toContain(readFileSync(dest, "utf-8"));
+  });
+
+  it("rejects a response that exceeds its byte limit", async () => {
+    mockGet.mockImplementation(httpsResponse(200, {}, "too large"));
+
+    const dir = mkdtempSync(join(tmpdir(), "hyperframes-download-"));
+    tempDirs.push(dir);
+    const dest = join(dir, "model.onnx");
+
+    await expect(
+      downloadFile("https://example.test/model.onnx", dest, { maxBytes: 3 }),
+    ).rejects.toThrow("Download exceeded 3 bytes");
+    expect(existsSync(dest)).toBe(false);
+    expect(readdirSync(dir)).toEqual([]);
+  });
+
   it("rejects an idle response and removes the partial file", async () => {
     const actualFs = await vi.importActual<typeof import("node:fs")>("node:fs");
     let responseClosed = false;
@@ -87,3 +136,59 @@ describe("downloadFile", () => {
     expect(existsSync(dest)).toBe(false);
   });
 });
+
+describe("redirect handling", () => {
+  it("rejects a redirect that the HTTPS client cannot follow", async () => {
+    mockGet
+      .mockImplementationOnce(
+        httpsResponse(307, { location: "http://cdn.example.test/model.onnx" }),
+      )
+      .mockImplementationOnce(() => {
+        throw new TypeError('Protocol "http:" not supported');
+      });
+
+    const dir = mkdtempSync(join(tmpdir(), "hyperframes-download-"));
+    tempDirs.push(dir);
+    await expect(
+      downloadFile("https://example.test/model.onnx", join(dir, "model.onnx")),
+    ).rejects.toThrow('Protocol "http:" not supported');
+  });
+
+  it("follows every redirect a host may answer with", () => {
+    // 307 is the one that broke: HuggingFace answers the tokenizer with it,
+    // and the original set stopped at 302, so the download fell through to the
+    // not-200 branch and hung.
+    for (const code of [301, 302, 303, 307, 308]) {
+      expect(REDIRECT_CODES.has(code)).toBe(true);
+    }
+  });
+
+  it("does not treat a success or an error as a redirect", () => {
+    for (const code of [200, 204, 404, 500]) {
+      expect(REDIRECT_CODES.has(code)).toBe(false);
+    }
+  });
+
+  it("resolves a relative location against the url that sent it", () => {
+    // The actual failure: handing this path back as a request target is not a
+    // valid URL, so nothing was ever fetched.
+    expect(
+      redirectTarget(
+        "/api/resolve-cache/models/x/tokenizer.json",
+        "https://huggingface.co/x/resolve/main/tokenizer.json",
+      ),
+    ).toBe("https://huggingface.co/api/resolve-cache/models/x/tokenizer.json");
+  });
+
+  it("keeps an absolute location as it is", () => {
+    expect(redirectTarget("https://cdn.example.com/blob", "https://huggingface.co/a/b")).toBe(
+      "https://cdn.example.com/blob",
+    );
+  });
+
+  it("resolves a sibling-relative location", () => {
+    expect(redirectTarget("other.json", "https://example.com/a/b/file.json")).toBe(
+      "https://example.com/a/b/other.json",
+    );
+  });
+});

packages/cli/src/utils/download.ts

@@ -1,15 +1,37 @@
+import { randomUUID } from "node:crypto";
 import { createWriteStream, renameSync, unlinkSync } from "node:fs";
 import { get as httpsGet } from "node:https";
-import type { IncomingMessage } from "node:http";
+import type { ClientRequest, IncomingMessage } from "node:http";
+import { Transform } from "node:stream";
 import { pipeline } from "node:stream/promises";
 
 const DEFAULT_DOWNLOAD_TIMEOUT_MS = 30_000;
 
 export interface DownloadOptions {
   /** Abort after this many milliseconds without network activity. */
   timeoutMs?: number;
+  /** Reject before writing more than this many response bytes. */
+  maxBytes?: number;
 }
 
+/** Every redirect a host may reasonably answer with, not just the two we saw first. */
+export const REDIRECT_CODES = new Set([301, 302, 303, 307, 308]);
+
+/**
+ * Where a redirect actually points.
+ *
+ * Split out because this is the part that was wrong, and it is the part that
+ * can be checked without a socket: hosts answer with relative locations
+ * (`/api/resolve-cache/...`) far more often than the original code assumed, and
+ * handing that string back as a request target fails.
+ */
+export function redirectTarget(location: string, from: string): string {
+  return new URL(location, from).toString();
+}
+
+/** Enough hops for a CDN handoff, few enough that a redirect loop still ends. */
+const MAX_REDIRECTS = 10;
+
 function removePartialFile(path: string): void {
   try {
     unlinkSync(path);
@@ -18,51 +40,97 @@ function removePartialFile(path: string): void {
   }
 }
 
+function enforceByteLimit(maxBytes: number): Transform {
+  let received = 0;
+  return new Transform({
+    transform(chunk: Buffer, _encoding, callback) {
+      received += chunk.byteLength;
+      callback(
+        received > maxBytes ? new Error(`Download exceeded ${maxBytes} bytes`) : null,
+        chunk,
+      );
+    },
+  });
+}
+
 /**
  * Download a file from a URL, following redirects.
  * Uses atomic write (download to .tmp, rename on success) to prevent
  * corrupt partial files from persisting in the cache on interruption.
+ *
+ * Location headers are resolved against the URL that sent them, because they
+ * are frequently relative: a host answering 307 with `/api/resolve-cache/...`
+ * is normal, and passing that string back as a request target is not.
+ *
+ * The timeout is per request, so it re-arms on each hop of a redirect chain
+ * rather than budgeting the whole chain. A stalled socket is what it exists to
+ * catch, and without it the CLI waits forever.
  */
 export function downloadFile(
   url: string,
   dest: string,
   options: DownloadOptions = {},
 ): Promise<void> {
-  const tmp = `${dest}.tmp`;
+  const tmp = `${dest}.${process.pid}.${randomUUID()}.tmp`;
   const timeoutMs = options.timeoutMs ?? DEFAULT_DOWNLOAD_TIMEOUT_MS;
+  const maxBytes = options.maxBytes;
   return new Promise((resolve, reject) => {
-    const follow = (u: string) => {
+    const follow = (u: string, hops = 0) => {
       let activeResponse: IncomingMessage | undefined;
       let responsePipelineStarted = false;
       let requestError: Error | undefined;
-      const request = httpsGet(u, (res) => {
-        activeResponse = res;
-        if (res.statusCode === 301 || res.statusCode === 302) {
-          const location = res.headers.location;
-          if (location) {
+      let request: ClientRequest;
+      try {
+        request = httpsGet(u, (res) => {
+          activeResponse = res;
+          if (res.statusCode && REDIRECT_CODES.has(res.statusCode)) {
+            const location = res.headers.location;
+            if (location) {
+              if (hops >= MAX_REDIRECTS) {
+                res.resume();
+                removePartialFile(tmp);
+                reject(
+                  new Error(`Download failed: more than ${MAX_REDIRECTS} redirects from ${url}`),
+                );
+                return;
+              }
+              res.resume();
+              try {
+                follow(redirectTarget(location, u), hops + 1);
+              } catch (error) {
+                removePartialFile(tmp);
+                reject(error);
+              }
+              return;
+            }
+          }
+          if (res.statusCode !== 200) {
             res.resume();
-            follow(location);
+            removePartialFile(tmp);
+            reject(new Error(`Download failed: HTTP ${res.statusCode}`));
             return;
           }
-        }
-        if (res.statusCode !== 200) {
-          res.resume();
-          removePartialFile(tmp);
-          reject(new Error(`Download failed: HTTP ${res.statusCode}`));
-          return;
-        }
-        const file = createWriteStream(tmp);
-        responsePipelineStarted = true;
-        pipeline(res, file)
-          .then(() => {
-            renameSync(tmp, dest);
-            resolve();
-          })
-          .catch((err) => {
-            removePartialFile(tmp);
-            reject(requestError ?? err);
-          });
-      });
+          const file = createWriteStream(tmp);
+          responsePipelineStarted = true;
+          const transfer =
+            maxBytes === undefined
+              ? pipeline(res, file)
+              : pipeline(res, enforceByteLimit(maxBytes), file);
+          transfer
+            .then(() => {
+              renameSync(tmp, dest);
+              resolve();
+            })
+            .catch((err) => {
+              removePartialFile(tmp);
+              reject(requestError ?? err);
+            });
+        });
+      } catch (error) {
+        removePartialFile(tmp);
+        reject(error);
+        return;
+      }
       request.setTimeout(timeoutMs, () => {
         request.destroy(new Error(`Download timed out after ${timeoutMs}ms`));
       });