omnirush 0.3.3 → 0.4.2

package/assets/collect-once.ts

@@ -3,9 +3,15 @@
 // sessions via the omnirush pi extension).
 //
 // Builds the standard collect envelope (schema_version 1, snapshot_type
-// "start") for a project directory and POSTs it zstd-compressed to
-// {origin}/collect with the device access token. Requires the bundled
-// bun (zstd) or Node >= 22.15.
+// "start") for a project directory and POSTs each zstd-compressed part
+// to {origin}/collect with the device access token. Requires the
+// bundled bun (zstd) or Node >= 22.15.
+//
+// Memory (issue #6): files are read and packed one part at a time and
+// each part is discarded right after its upload — nothing workspace-
+// sized is ever held in full. The pre-flight summary walks sizes only
+// (lstat), never content. Deterministic order: the walker yields paths
+// sorted, each path's entries (manifest + slices) grouped.
 //
 // Usage: omnirush collect [--dir PATH] [--yes] [--dry-run]
 
@@ -13,10 +19,10 @@ import { randomBytes } from "node:crypto";
 import path from "node:path";
 import readline from "node:readline";
 
-import { omniDir, resolveOrigin, gatewayUrlForOrigin } from "./extensions/omnirush/auth";
+import { resolveOrigin, gatewayUrlForOrigin } from "./extensions/omnirush/auth";
 import {
-  buildEnvelopeParts,
-  collectFiles,
+  iterateCollectFiles,
+  iterateEnvelopeParts,
   MAX_PART_COMPRESSED_BYTES,
   MAX_PART_COMPRESSED_HARD,
   MAX_SESSION_BYTES,
@@ -50,49 +56,36 @@ async function confirm(question: string): Promise<boolean> {
   return answer.trim().toLowerCase() === "y" || answer.trim().toLowerCase() === "yes";
 }
 
+/** Cheap workspace summary: path count + on-disk bytes (no content). */
+async function summarize(dir: string): Promise<{ files: number; bytes: number }> {
+  const { lstat } = await import("node:fs/promises");
+  const { listWorkspaceFiles } = await import("./extensions/omnirush/collector-lib");
+  const paths = await listWorkspaceFiles(dir);
+  let bytes = 0;
+  for (const relativePath of paths) {
+    try {
+      bytes += (await lstat(path.join(dir, relativePath))).size;
+    } catch {
+      /* raced away — the streaming pass re-checks */
+    }
+  }
+  return { files: paths.length, bytes };
+}
+
 async function main() {
   const { dir, yes, dryRun } = parseArgs(process.argv.slice(2));
   const sessionId = `manual-${randomBytes(6).toString("hex")}`;
   const workspaceId = path.basename(dir) || "workspace";
 
-  const files = await collectFiles(dir, MAX_SESSION_BYTES, workspaceId, {
-    id: sessionId,
-    segment: 1,
-    resumed: false,
-  });
-  const bytes = files.reduce((total, file) => total + Buffer.byteLength(file.content), 0);
-  if (files.length <= 1) {
+  const summary = await summarize(dir);
+  if (summary.files === 0) {
     console.error("omnirush collect: nothing to upload (empty or fully filtered workspace)");
     process.exit(1);
   }
-
-  const state = { id: sessionId, segment: 1, resumed: false };
-  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: ${parts.length} part(s), ${plainBytes} bytes plain, ${compressedBytes} bytes zstd ` +
-      `(target <= ${MAX_PART_COMPRESSED_BYTES}, rail ${MAX_PART_COMPRESSED_HARD})`,
+    `collect: ${summary.files} files (${summary.bytes} bytes on disk, denylist applied) as session ${sessionId}`,
   );
 
-  if (dryRun) {
-    console.log("dry run — no upload performed");
-    return;
-  }
-  if (!yes && !(await confirm(`Upload ${files.length - 1} files from ${dir} to Omnirush? [y/N] `))) {
-    console.log("aborted");
-    process.exit(1);
-  }
-
   const origin = resolveOrigin(process.env);
   const collectUrl = gatewayUrlForOrigin(origin).replace(/\/v1$/, "/collect");
   const accessToken = (
@@ -104,7 +97,34 @@ async function main() {
     console.error("omnirush collect: no access token (run `omnirush login`)");
     process.exit(1);
   }
-  for (const part of parts) {
+
+  const state = { id: sessionId, segment: 1, resumed: false };
+  // Pass 2 streams content: file -> batch -> part -> upload, with only
+  // one part ever resident. `part` is dropped at each loop iteration.
+  const parts = iterateEnvelopeParts(
+    state,
+    "start",
+    iterateCollectFiles(dir, MAX_SESSION_BYTES, workspaceId, {
+      id: sessionId,
+      segment: 1,
+      resumed: false,
+    }),
+    MAX_PART_COMPRESSED_BYTES,
+    (message, attributes) => console.error(`omnirush collect: ${message} (${JSON.stringify(attributes ?? {})})`),
+  );
+
+  let partCount = 0;
+  let plainBytes = 0;
+  let compressedBytes = 0;
+  let uploaded = false;
+  for await (const part of parts) {
+    partCount += 1;
+    plainBytes += part.payload.length;
+    compressedBytes += part.compressed.length;
+    if (dryRun) {
+      console.log(`part ${partCount}: ${part.compressed.length} bytes zstd (discarded — dry run)`);
+      continue;
+    }
     const response = await fetch(collectUrl, {
       method: "POST",
       headers: {
@@ -123,6 +143,7 @@ async function main() {
       console.error(`upload failed: HTTP ${response.status}`);
       process.exit(1);
     }
+    uploaded = true;
     const extra = [...response.headers]
       .filter(([name]) => name.toLowerCase().startsWith("x-omnirush-"))
       .map(([name, value]) => `${name}: ${value}`);
@@ -132,6 +153,23 @@ async function main() {
         `${extra.length ? "\n  " + extra.join("\n  ") : ""}`,
     );
   }
+
+  if (partCount === 0) {
+    console.error("omnirush collect: nothing to upload (empty or fully filtered workspace)");
+    process.exit(1);
+  }
+  console.log(
+    `envelope: ${partCount} part(s), ${plainBytes} bytes plain, ${compressedBytes} bytes zstd ` +
+      `(target <= ${MAX_PART_COMPRESSED_BYTES}, rail ${MAX_PART_COMPRESSED_HARD})`,
+  );
+  if (dryRun) {
+    console.log("dry run — no upload performed");
+    return;
+  }
+  if (!uploaded) {
+    console.error("omnirush collect: nothing uploaded");
+    process.exit(1);
+  }
 }
 
 main().catch((error) => {

package/assets/extensions/omnirush/auth.js

@@ -29,6 +29,12 @@ import fs from "node:fs";
 import os from "node:os";
 import path from "node:path";
 
+import {
+  isRetryableStatus,
+  retryAttempts,
+  withRetries,
+} from "./retry.js";
+
 export const DEFAULT_ORIGIN = "https://omnirush.ai/omnirush";
 export const DEFAULT_POLL_INTERVAL_SEC = 5;
 export const MAX_POLL_INTERVAL_SEC = 30;
@@ -171,19 +177,36 @@ function detailOf(payload, fallback) {
 }
 
 async function postJson(url, body, fetchImpl, timeoutMs = 20_000) {
-  const response = await fetchImpl(url, {
-    method: "POST",
-    headers: { "Content-Type": "application/json" },
-    body: JSON.stringify(body),
-    signal: AbortSignal.timeout(timeoutMs),
-  });
-  let payload = null;
-  try {
-    payload = await response.json();
-  } catch {
-    /* non-JSON body — fall through with payload null */
-  }
-  return { response, payload };
+  // Shared retry layer: 429/5xx/network failures ride out backend
+  // deploys. 4xx are protocol answers (428 pending, 401 rejected) —
+  // never retried here; callers interpret them.
+  const outcome = await withRetries(
+    async () => {
+      const response = await fetchImpl(url, {
+        method: "POST",
+        headers: { "Content-Type": "application/json" },
+        body: JSON.stringify(body),
+        signal: AbortSignal.timeout(timeoutMs),
+      });
+      let payload = null;
+      try {
+        payload = await response.json();
+      } catch {
+        /* non-JSON body — fall through with payload null */
+      }
+      return {
+        response,
+        payload,
+        retryAfterSec: Number(response?.headers?.get?.("retry-after")) || undefined,
+      };
+    },
+    {
+      attempts: retryAttempts(),
+      isRetryable: ({ response }) => isRetryableStatus(response?.status),
+    },
+  );
+  if (outcome.error) throw outcome.error;
+  return outcome.result;
 }
 
 /** Step 1 of the device flow: create a pending authorization. */
@@ -305,10 +328,26 @@ export async function refreshTokens(origin, { refreshToken, fetchImpl } = {}) {
 
 export async function deviceMe(origin, { accessToken, fetchImpl } = {}) {
   const doFetch = fetchImpl ?? globalThis.fetch;
-  const response = await doFetch(`${origin}/device/me`, {
-    headers: { Authorization: `Bearer ${accessToken}` },
-    signal: AbortSignal.timeout(15_000),
-  });
+  // 401 is a protocol answer (stale access token — the caller refreshes
+  // and retries once); 429/5xx/network retry with backoff.
+  const outcome = await withRetries(
+    async () => {
+      const response = await doFetch(`${origin}/device/me`, {
+        headers: { Authorization: `Bearer ${accessToken}` },
+        signal: AbortSignal.timeout(15_000),
+      });
+      return {
+        response,
+        retryAfterSec: Number(response?.headers?.get?.("retry-after")) || undefined,
+      };
+    },
+    {
+      attempts: retryAttempts(),
+      isRetryable: ({ response }) => isRetryableStatus(response?.status),
+    },
+  );
+  if (outcome.error) throw outcome.error;
+  const response = outcome.result.response;
   if (response.status === 401) return null;
   if (!response.ok) throw new Error(`device me failed: HTTP ${response.status}`);
   return response.json();