Speed up version client by partial response reads (#807)

__tests__/download/download-version.test.ts

@@ -1,5 +1,6 @@
 import { beforeEach, describe, expect, it, jest } from "@jest/globals";
 import * as semver from "semver";
+import { VERSIONS_MANIFEST_URL } from "../../src/utils/constants";
 
 const mockInfo = jest.fn();
 const mockWarning = jest.fn();
@@ -36,11 +37,14 @@ const mockGetLatestVersion = jest.fn<any>();
 // biome-ignore lint/suspicious/noExplicitAny: Mock requires flexible typing in tests.
 const mockGetAllVersions = jest.fn<any>();
 // biome-ignore lint/suspicious/noExplicitAny: Mock requires flexible typing in tests.
+const mockGetFirstMatchingVersion = jest.fn<any>();
+// biome-ignore lint/suspicious/noExplicitAny: Mock requires flexible typing in tests.
 const mockGetArtifact = jest.fn<any>();
 
 jest.unstable_mockModule("../../src/download/manifest", () => ({
   getAllVersions: mockGetAllVersions,
   getArtifact: mockGetArtifact,
+  getFirstMatchingVersion: mockGetFirstMatchingVersion,
   getLatestVersion: mockGetLatestVersion,
 }));
 
@@ -65,6 +69,7 @@ describe("download-version", () => {
     mockCacheDir.mockReset();
     mockGetLatestVersion.mockReset();
     mockGetAllVersions.mockReset();
+    mockGetFirstMatchingVersion.mockReset();
     mockGetArtifact.mockReset();
     mockValidateChecksum.mockReset();
 
@@ -85,14 +90,43 @@ describe("download-version", () => {
       expect(mockGetLatestVersion).toHaveBeenCalledWith(undefined);
     });
 
-    it("uses the default manifest to resolve available versions", async () => {
-      mockGetAllVersions.mockResolvedValue(["0.9.26", "0.9.25"]);
+    it("stops at the first matching version in the default manifest", async () => {
+      mockGetFirstMatchingVersion.mockImplementation(
+        (predicate: (version: string) => boolean) =>
+          ["0.9.26", "0.9.25"].find(predicate),
+      );
 
       const version = await resolveVersion("^0.9.0", undefined);
 
       expect(version).toBe("0.9.26");
-      expect(mockGetAllVersions).toHaveBeenCalledTimes(1);
-      expect(mockGetAllVersions).toHaveBeenCalledWith(undefined);
+      expect(mockGetFirstMatchingVersion).toHaveBeenCalledTimes(1);
+      expect(mockGetAllVersions).not.toHaveBeenCalled();
+    });
+
+    it("streams ranges when the default manifest URL is explicit", async () => {
+      mockGetFirstMatchingVersion.mockImplementation(
+        (predicate: (version: string) => boolean) =>
+          ["0.9.26", "0.9.25"].find(predicate),
+      );
+
+      const version = await resolveVersion("^0.9.0", VERSIONS_MANIFEST_URL);
+
+      expect(version).toBe("0.9.26");
+      expect(mockGetFirstMatchingVersion).toHaveBeenCalledTimes(1);
+      expect(mockGetAllVersions).not.toHaveBeenCalled();
+    });
+
+    it("streams PEP 440 ranges from the default manifest", async () => {
+      mockGetFirstMatchingVersion.mockImplementation(
+        (predicate: (version: string) => boolean) =>
+          ["0.9.26", "0.9.25"].find(predicate),
+      );
+
+      const version = await resolveVersion("!=0.9.26", undefined);
+
+      expect(version).toBe("0.9.25");
+      expect(mockGetFirstMatchingVersion).toHaveBeenCalledTimes(1);
+      expect(mockGetAllVersions).not.toHaveBeenCalled();
     });
 
     it("treats == exact pins as explicit versions", async () => {

__tests__/download/manifest.test.ts

@@ -17,6 +17,7 @@ const {
   fetchManifest,
   getAllVersions,
   getArtifact,
+  getFirstMatchingVersion,
   getLatestVersion,
   parseManifest,
 } = await import("../../src/download/manifest");
@@ -33,6 +34,7 @@ function createMockResponse(
   data: string,
 ) {
   return {
+    body: null,
     ok,
     status,
     statusText,
@@ -40,6 +42,31 @@ function createMockResponse(
   };
 }
 
+function createStreamingMockResponse(
+  chunks: string[],
+  cancel: () => void | Promise<void> = () => {},
+  close = false,
+) {
+  const encoder = new TextEncoder();
+  return {
+    body: new ReadableStream<Uint8Array>({
+      cancel,
+      start(controller) {
+        for (const chunk of chunks) {
+          controller.enqueue(encoder.encode(chunk));
+        }
+        if (close) {
+          controller.close();
+        }
+      },
+    }),
+    ok: true,
+    status: 200,
+    statusText: "OK",
+    text: async () => chunks.join(""),
+  };
+}
+
 describe("manifest", () => {
   beforeEach(() => {
     clearManifestCache();
@@ -105,6 +132,75 @@ describe("manifest", () => {
         getLatestVersion("https://example.com/custom.ndjson"),
       ).resolves.toBe("0.9.26");
     });
+
+    it("stops reading the default manifest after the first record", async () => {
+      const [latestVersion] = sampleManifestResponse.split("\n");
+      const cancel = jest.fn();
+      mockFetch.mockResolvedValue(
+        createStreamingMockResponse(
+          [
+            latestVersion.slice(0, 100),
+            `${latestVersion.slice(100)}\n`,
+            "invalid trailing data\n",
+          ],
+          cancel,
+        ),
+      );
+
+      await expect(getLatestVersion()).resolves.toBe("0.9.26");
+      await expect(
+        getArtifact("0.9.26", "aarch64", "apple-darwin"),
+      ).resolves.toBeDefined();
+
+      expect(cancel).toHaveBeenCalledTimes(1);
+      expect(mockFetch).toHaveBeenCalledTimes(1);
+    });
+
+    it("does not fail when canceling the remaining response fails", async () => {
+      const [latestVersion] = sampleManifestResponse.split("\n");
+      mockFetch.mockResolvedValue(
+        createStreamingMockResponse([`${latestVersion}\n`], () =>
+          Promise.reject(new Error("cancel failed")),
+        ),
+      );
+
+      await expect(getLatestVersion()).resolves.toBe("0.9.26");
+    });
+  });
+
+  describe("getFirstMatchingVersion", () => {
+    it("stops at the first matching record", async () => {
+      const cancel = jest.fn();
+      mockFetch.mockResolvedValue(
+        createStreamingMockResponse(
+          [`${sampleManifestResponse}\n`, "invalid trailing data\n"],
+          cancel,
+        ),
+      );
+
+      await expect(
+        getFirstMatchingVersion((version) => version === "0.9.25"),
+      ).resolves.toBe("0.9.25");
+
+      expect(cancel).toHaveBeenCalledTimes(1);
+    });
+
+    it("caches a fully consumed response stream", async () => {
+      mockFetch.mockResolvedValue(
+        createStreamingMockResponse(
+          [`${sampleManifestResponse}\n`],
+          undefined,
+          true,
+        ),
+      );
+
+      await expect(
+        getFirstMatchingVersion((version) => version === "0.0.1"),
+      ).resolves.toBeUndefined();
+      await expect(getLatestVersion()).resolves.toBe("0.9.26");
+
+      expect(mockFetch).toHaveBeenCalledTimes(1);
+    });
   });
 
   describe("getArtifact", () => {
@@ -168,6 +264,60 @@ describe("manifest", () => {
 
       expect(artifact).toBeUndefined();
     });
+
+    it("does not cache records from a failed stream read", async () => {
+      const [latestVersion] = sampleManifestResponse.split("\n");
+      mockFetch
+        .mockResolvedValueOnce(
+          createStreamingMockResponse([
+            `${latestVersion}\n`,
+            "invalid manifest record\n",
+          ]),
+        )
+        .mockResolvedValueOnce(
+          createMockResponse(true, 200, "OK", sampleManifestResponse),
+        );
+
+      await expect(
+        getArtifact("0.0.1", "aarch64", "apple-darwin"),
+      ).rejects.toThrow("Failed to parse manifest data");
+      await expect(getLatestVersion()).resolves.toBe("0.9.26");
+
+      expect(mockFetch).toHaveBeenCalledTimes(2);
+    });
+
+    it("does not cache records when the response stream fails", async () => {
+      const [latestVersion] = sampleManifestResponse.split("\n");
+      const encoder = new TextEncoder();
+      let sentVersion = false;
+      const body = new ReadableStream<Uint8Array>({
+        pull(controller) {
+          if (!sentVersion) {
+            sentVersion = true;
+            controller.enqueue(encoder.encode(`${latestVersion}\n`));
+            return;
+          }
+          controller.error(new Error("response stream failed"));
+        },
+      });
+      mockFetch
+        .mockResolvedValueOnce({
+          body,
+          ok: true,
+          status: 200,
+          statusText: "OK",
+        })
+        .mockResolvedValueOnce(
+          createMockResponse(true, 200, "OK", sampleManifestResponse),
+        );
+
+      await expect(
+        getArtifact("0.0.1", "aarch64", "apple-darwin"),
+      ).rejects.toThrow("response stream failed");
+      await expect(getLatestVersion()).resolves.toBe("0.9.26");
+
+      expect(mockFetch).toHaveBeenCalledTimes(2);
+    });
   });
 
   describe("parseManifest", () => {