computer: Stream ranged tool reads

Author: aron-cf

packages/computer/src/stub.test.ts

@@ -205,6 +205,16 @@ describe("WorkspaceStub", () => {
     });
   });
 
+  it("fs.readFile forwards ranged stream options", async () => {
+    await withStub(async (ws) => {
+      const stub = ws.stub();
+      await stub.fs.writeFile("/bin", new Uint8Array([1, 2, 3, 4, 5]));
+      const stream = await stub.fs.readFile("/bin", { byteOffset: 1, byteLength: 3 });
+      const bytes = new Uint8Array(await new Response(stream).arrayBuffer());
+      expect(Array.from(bytes)).toEqual([2, 3, 4]);
+    });
+  });
+
   it("fs.readRange forwards bounded byte reads", async () => {
     await withStub(async (ws) => {
       const stub = ws.stub();

packages/computer/src/stub.ts

@@ -104,6 +104,11 @@ export class WorkspaceFilesystemStub extends RpcTarget {
 
   readFile(path: string): Promise<ReadableStream<Uint8Array>>;
   readFile(path: string, encoding: "utf8"): Promise<string>;
+  readFile(
+    path: string,
+    options: ReadFileOptions & { encoding?: undefined },
+  ): Promise<ReadableStream<Uint8Array>>;
+  readFile(path: string, options: ReadFileOptions & { encoding: "utf8" }): Promise<string>;
   readFile(path: string, options: ReadFileOptions): Promise<string | ReadableStream<Uint8Array>>;
   readFile(
     path: string,

packages/computer/src/tools/ai.test.ts

@@ -179,82 +179,124 @@ function memoryStore(options: {
 }
 
 describe("WorkspaceFileStore", () => {
-  it("uses readRange without streaming skipped bytes from Workspace.fs", async () => {
-    const calls: Array<{ offset: number; length: number }> = [];
+  it("opens one ranged stream instead of issuing repeated range calls", async () => {
+    const calls: Array<{ byteOffset?: number; byteLength?: number }> = [];
     const content = bytes("abcdefghij");
     const workspace = {
       fs: {
         async stat() {
-          return {
-            size: content.length,
-            mtime: 1,
-            mode: 0o100644,
-            isFile: true,
-            isDirectory: false,
-          };
+          throw new Error("stat must not be called by readChunks");
         },
-        async readRange(_path: string, offset: number, length: number) {
-          calls.push({ offset, length });
-          return content.slice(offset, offset + length);
+        async readRange() {
+          throw new Error("readRange must not be called by readChunks");
         },
-        async readFile(): Promise<ReadableStream<Uint8Array>> {
-          throw new Error("readFile must not be called for ranged reads");
+        async readFile(
+          _path: string,
+          options: { byteOffset?: number; byteLength?: number } = {},
+        ): Promise<ReadableStream<Uint8Array>> {
+          calls.push(options);
+          const start = options.byteOffset ?? 0;
+          const end = options.byteLength === undefined ? undefined : start + options.byteLength;
+          return new ReadableStream({
+            start(controller) {
+              controller.enqueue(content.slice(start, end));
+              controller.close();
+            },
+          });
         },
         async writeFile() {},
         async mkdir() {},
         async rm() {},
-        async readdir() {
-          return [];
-        },
       },
     };
     const store = new WorkspaceFileStore(workspace);
 
     await expect(
       drainChunks(store.readChunks("/workspace/range.txt", 2, 5)).then(decode),
     ).resolves.toBe("cdefg");
-    expect(calls).toEqual([{ offset: 2, length: 5 }]);
+    expect(calls).toEqual([{ byteOffset: 2, byteLength: 5 }]);
+  });
+
+  it("still validates the path for a zero-length read", async () => {
+    const store = new WorkspaceFileStore(makeWorkspace());
+
+    await expect(drainChunks(store.readChunks("/missing", 0, 0))).rejects.toMatchObject({
+      code: "ENOENT",
+    });
   });
 
-  it("splits unbounded reads into fixed-size ranges", async () => {
-    const calls: Array<{ offset: number; length: number }> = [];
-    const content = new Uint8Array(150_000).fill(7);
+  it("rejects directories instead of treating them as empty files", async () => {
+    const workspace = makeWorkspace();
+    await workspace.fs.mkdir("/directory");
+    const store = new WorkspaceFileStore(workspace);
+
+    await expect(drainChunks(store.readChunks("/directory"))).rejects.toMatchObject({
+      code: "EISDIR",
+    });
+  });
+
+  it("keeps a real multi-chunk workspace read on one snapshot", async () => {
+    const workspace = makeWorkspace();
+    await workspace.fs.mkdir("/workspace", { recursive: true });
+    const original = new Uint8Array(600_000);
+    original.fill(0x41, 0, 500_000);
+    original.fill(0x42, 500_000);
+    await workspace.fs.writeFile("/workspace/large.bin", original);
+    const store = new WorkspaceFileStore(workspace);
+
+    const chunks = store.readChunks("/workspace/large.bin")[Symbol.asyncIterator]();
+    const first = await chunks.next();
+    expect(first.done).toBe(false);
+    await workspace.fs.writeFile(
+      "/workspace/large.bin",
+      new Uint8Array(original.length).fill(0x43),
+    );
+
+    const parts = [first.value];
+    while (true) {
+      const next = await chunks.next();
+      if (next.done) break;
+      parts.push(next.value);
+    }
+    const result = await drainChunks(
+      (async function* () {
+        yield* parts;
+      })(),
+    );
+    expect(result.byteLength).toBe(original.byteLength);
+    expect(result.every((value, index) => value === original[index])).toBe(true);
+  });
+
+  it("cancels a ranged stream when its consumer stops early", async () => {
+    let cancelled = false;
     const workspace = {
       fs: {
         async stat() {
-          return {
-            size: content.length,
-            mtime: 1,
-            mode: 0o100644,
-            isFile: true,
-            isDirectory: false,
-          };
+          throw new Error("stat must not be called by readChunks");
         },
-        async readRange(_path: string, offset: number, length: number) {
-          calls.push({ offset, length });
-          return content.slice(offset, offset + length);
+        async readRange() {
+          throw new Error("readRange must not be called by readChunks");
         },
         async readFile(): Promise<ReadableStream<Uint8Array>> {
-          throw new Error("readFile must not be called for ranged reads");
+          return new ReadableStream({
+            start(controller) {
+              controller.enqueue(bytes("first"));
+              controller.enqueue(bytes("second"));
+            },
+            cancel() {
+              cancelled = true;
+            },
+          });
         },
         async writeFile() {},
         async mkdir() {},
         async rm() {},
-        async readdir() {
-          return [];
-        },
       },
     };
     const store = new WorkspaceFileStore(workspace);
 
-    await expect(drainChunks(store.readChunks("/workspace/large.bin"))).resolves.toHaveLength(
-      content.length,
-    );
-    expect(calls).toEqual([
-      { offset: 0, length: 65_536 },
-      { offset: 65_536, length: 65_536 },
-      { offset: 131_072, length: 18_928 },
-    ]);
+    for await (const _chunk of store.readChunks("/workspace/range.txt")) break;
+    expect(cancelled).toBe(true);
   });
 });
 

packages/computer/src/tools/fs/store.ts

@@ -6,14 +6,13 @@
  * class. This adapter is the bridge from that contract to the public
  * `workspace.fs` surface.
  *
- * Bounded reads go through `fs.readRange`; whole-file reads used by edit
- * and multimodal output still drain `fs.readFile(path)`.
+ * Chunked and ranged reads use one `fs.readFile` stream so remote workspaces
+ * keep one snapshot and one RPC invocation. Whole-file reads used by edit and
+ * multimodal output drain the same stream interface.
  */
 
 import type { FileStat, FileStore } from "./types.js";
 
-const RANGE_CHUNK_BYTES = 64 * 1024;
-
 /**
  * Structural subset of `@cloudflare/computer.Workspace` the tools
  * depend on.
@@ -27,8 +26,10 @@ export interface WorkspaceLike {
       isFile: boolean;
       isDirectory: boolean;
     }>;
-    readFile(path: string): Promise<ReadableStream<Uint8Array>>;
-    readRange(path: string, offset: number, length: number): Promise<Uint8Array>;
+    readFile(
+      path: string,
+      options?: { byteOffset?: number; byteLength?: number },
+    ): Promise<ReadableStream<Uint8Array>>;
     writeFile(path: string, content: Uint8Array, options?: { mode?: number }): Promise<void>;
     mkdir(path: string, options?: { recursive?: boolean }): Promise<void>;
     rm(path: string, options?: { recursive?: boolean; force?: boolean }): Promise<void>;
@@ -84,20 +85,24 @@ export class WorkspaceFileStore implements FileStore {
     if (byteLength !== undefined && (!Number.isSafeInteger(byteLength) || byteLength < 0)) {
       throw new Error("readChunks: byteLength must be a non-negative safe integer");
     }
-    if (byteLength === 0) return;
-
-    const stat = await this.ws.fs.stat(path);
-    let remaining = Math.max(0, stat.size - byteOffset);
-    if (byteLength !== undefined) remaining = Math.min(remaining, byteLength);
-    let offset = byteOffset;
-    while (remaining > 0) {
-      const requested = Math.min(remaining, RANGE_CHUNK_BYTES);
-      const chunk = await this.ws.fs.readRange(path, offset, requested);
-      if (chunk.byteLength === 0) return;
-      yield chunk;
-      offset += chunk.byteLength;
-      remaining -= chunk.byteLength;
-      if (chunk.byteLength < requested) return;
+    const stream = await this.ws.fs.readFile(path, { byteOffset, byteLength });
+    const reader = stream.getReader();
+    let completed = false;
+    try {
+      while (true) {
+        const { value, done } = await reader.read();
+        if (done) {
+          completed = true;
+          return;
+        }
+        if (value !== undefined && value.byteLength > 0) yield value;
+      }
+    } finally {
+      try {
+        if (!completed) await reader.cancel();
+      } finally {
+        reader.releaseLock();
+      }
     }
   }
 }