omnirush 0.2.0 → 0.3.0

package/assets/collect-once.ts

@@ -15,11 +15,11 @@ import readline from "node:readline";
 
 import { omniDir, resolveOrigin, gatewayUrlForOrigin } from "./extensions/omnirush/auth";
 import {
-  buildEnvelopePayload,
+  buildEnvelopeParts,
   collectFiles,
-  compressZstd,
-  MAX_COMPRESSED_BYTES,
-  MAX_COLLECTOR_SESSION_BYTES,
+  MAX_PART_COMPRESSED_BYTES,
+  MAX_PART_COMPRESSED_HARD,
+  MAX_SESSION_BYTES,
 } from "./extensions/omnirush/collector-lib";
 
 function parseArgs(argv: string[]): { dir: string; yes: boolean; dryRun: boolean } {
@@ -55,7 +55,7 @@ async function main() {
   const sessionId = `manual-${randomBytes(6).toString("hex")}`;
   const workspaceId = path.basename(dir) || "workspace";
 
-  const files = await collectFiles(dir, MAX_COLLECTOR_SESSION_BYTES, workspaceId, {
+  const files = await collectFiles(dir, MAX_SESSION_BYTES, workspaceId, {
     id: sessionId,
     segment: 1,
     resumed: false,
@@ -65,15 +65,23 @@ async function main() {
     console.error("omnirush collect: nothing to upload (empty or fully filtered workspace)");
     process.exit(1);
   }
-  console.log(
-    `collect: ${files.length - 1} files (${bytes} bytes, denylist/caps applied) as session ${sessionId}`,
-  );
 
   const state = { id: sessionId, segment: 1, resumed: false };
-  const payload = buildEnvelopePayload(state, "start", files);
-  const compressed = await compressZstd(payload);
+  const parts = await buildEnvelopeParts(
+    state,
+    "start",
+    files,
+    MAX_PART_COMPRESSED_BYTES,
+    (message, attributes) => console.error(`omnirush collect: ${message} (${JSON.stringify(attributes ?? {})})`),
+  );
+  const plainBytes = parts.reduce((total, part) => total + part.payload.length, 0);
+  const compressedBytes = parts.reduce((total, part) => total + part.compressed.length, 0);
+  console.log(
+    `collect: ${files.length - 1} files (${bytes} bytes, denylist applied) as session ${sessionId}`,
+  );
   console.log(
-    `envelope: ${payload.length} bytes plain, ${compressed.length} bytes zstd (cap ${MAX_COMPRESSED_BYTES})`,
+    `envelope: ${parts.length} part(s), ${plainBytes} bytes plain, ${compressedBytes} bytes zstd ` +
+      `(target <= ${MAX_PART_COMPRESSED_BYTES}, rail ${MAX_PART_COMPRESSED_HARD})`,
   );
 
   if (dryRun) {
@@ -96,28 +104,34 @@ async function main() {
     console.error("omnirush collect: no access token (run `omnirush login`)");
     process.exit(1);
   }
-  const response = await fetch(collectUrl, {
-    method: "POST",
-    headers: {
-      Authorization: `Bearer ${accessToken}`,
-      "Content-Type": "application/zstd",
-      "X-Omnirush-Session-ID": sessionId,
-    },
-    body: compressed.buffer.slice(
-      compressed.byteOffset,
-      compressed.byteOffset + compressed.byteLength,
-    ) as ArrayBuffer,
-    signal: AbortSignal.timeout(60_000),
-  });
-  if (!response.ok) {
-    await response.body?.cancel().catch(() => undefined);
-    console.error(`upload failed: HTTP ${response.status}`);
-    process.exit(1);
+  for (const part of parts) {
+    const response = await fetch(collectUrl, {
+      method: "POST",
+      headers: {
+        Authorization: `Bearer ${accessToken}`,
+        "Content-Type": "application/zstd",
+        "X-Omnirush-Session-ID": sessionId,
+      },
+      body: part.compressed.buffer.slice(
+        part.compressed.byteOffset,
+        part.compressed.byteOffset + part.compressed.byteLength,
+      ) as ArrayBuffer,
+      signal: AbortSignal.timeout(120_000),
+    });
+    if (!response.ok) {
+      await response.body?.cancel().catch(() => undefined);
+      console.error(`upload failed: HTTP ${response.status}`);
+      process.exit(1);
+    }
+    const extra = [...response.headers]
+      .filter(([name]) => name.toLowerCase().startsWith("x-omnirush-"))
+      .map(([name, value]) => `${name}: ${value}`);
+    const sequence = JSON.parse(part.payload.toString("utf8")).sequence;
+    console.log(
+      `uploaded part sequence=${sequence} (${part.compressed.length} bytes, HTTP ${response.status})` +
+        `${extra.length ? "\n  " + extra.join("\n  ") : ""}`,
+    );
   }
-  const extra = [...response.headers]
-    .filter(([name]) => name.toLowerCase().startsWith("x-omnirush-"))
-    .map(([name, value]) => `${name}: ${value}`);
-  console.log(`uploaded (HTTP ${response.status})${extra.length ? "\n  " + extra.join("\n  ") : ""}`);
 }
 
 main().catch((error) => {