feat(streams): Add support for remote iterators and generators between vats (#574)

Ferrying the `async function* ()` and `for await (...)` syntax between
vats is complicated slightly by the use of the unserializable
`Symbol.asyncIterator` to implement these features. This PR adds two
functions to `@metamask/streams` which are intended to be used in
conjunction.
- `makeFarGenerator` - which makes a remotable reference from an async
generator
- `makeEventualIterator` - which makes a for-awaitable iterator from a
remote generator presence

Author: grypez

packages/kernel-test/package.json

@@ -47,6 +47,7 @@
     "@agoric/store": "0.9.3-u19.0",
     "@endo/eventual-send": "^1.3.1",
     "@endo/exo": "^1.5.9",
+    "@endo/far": "^1.1.11",
     "@endo/marshal": "^1.6.4",
     "@endo/patterns": "^1.5.0",
     "@endo/promise-kit": "^1.1.10",

packages/kernel-test/src/async-generator.test.ts

@@ -0,0 +1,73 @@
+import { makeSQLKernelDatabase } from '@metamask/kernel-store/sqlite/nodejs';
+import { waitUntilQuiescent } from '@metamask/kernel-utils';
+import { Kernel } from '@metamask/ocap-kernel';
+import { describe, expect, it } from 'vitest';
+
+import {
+  getBundleSpec,
+  makeKernel,
+  makeTestLogger,
+  runTestVats,
+  extractTestLogs,
+} from './utils.ts';
+
+const testSubcluster = {
+  bootstrap: 'consumer',
+  forceReset: true,
+  vats: {
+    consumer: {
+      bundleSpec: getBundleSpec('async-generator-iterator-vat'),
+      parameters: {
+        name: 'alice',
+      },
+    },
+    producer: {
+      bundleSpec: getBundleSpec('async-generator-iterator-vat'),
+      parameters: {
+        name: 'bob',
+      },
+    },
+  },
+};
+
+describe(
+  'Async generator consumption between vats',
+  {
+    timeout: 2000,
+  },
+  () => {
+    let kernel: Kernel;
+
+    it('alice can consume async generator from Bob using for await', async () => {
+      const kernelDatabase = await makeSQLKernelDatabase({
+        dbFilename: ':memory:',
+      });
+      const { logger, entries } = makeTestLogger();
+      kernel = await makeKernel(kernelDatabase, true, logger);
+
+      await runTestVats(kernel, testSubcluster);
+      await waitUntilQuiescent(100);
+
+      const aliceLogs = extractTestLogs(entries, 'alice');
+      const bobLogs = extractTestLogs(entries, 'bob');
+      expect(aliceLogs).toStrictEqual([
+        'alice buildRootObject',
+        'alice is bootstrap',
+        'alice iterating 0',
+        'alice iterating 1',
+        'alice iterating 2',
+        'alice iterating 3',
+        'alice iterating 4',
+      ]);
+
+      expect(bobLogs).toStrictEqual([
+        'bob buildRootObject',
+        'bob generating 0',
+        'bob generating 1',
+        'bob generating 2',
+        'bob generating 3',
+        'bob generating 4',
+      ]);
+    });
+  },
+);

packages/kernel-test/src/vats/async-generator-iterator-vat.js

@@ -0,0 +1,45 @@
+import { E, Far } from '@endo/far';
+// ESLint's import-x/no-unresolved rule with commonjs:false doesn't support subpath exports
+// eslint-disable-next-line import-x/no-unresolved
+import { makeEventualIterator, makeFarGenerator } from '@metamask/streams/vat';
+
+/**
+ * Build function for testing async generators.
+ *
+ * @param {object} vatPowers - The powers of the vat.
+ * @param {object} vatPowers.logger - The logger to use.
+ * @param {object} parameters - The parameters of the vat.
+ * @param {string} parameters.name - The name of the vat.
+ * @returns {object} The root object for the vat.
+ */
+export function buildRootObject({ logger }, { name }) {
+  const tlogger = logger.subLogger({ tags: ['test', name] });
+  const tlog = (...args) => tlogger.log(...args);
+
+  tlog(`${name} buildRootObject`);
+
+  return Far('root', {
+    async bootstrap({ consumer, producer }, _services) {
+      tlog(`${name} is bootstrap`);
+      await E(consumer).iterate(producer);
+    },
+
+    generate: async (stop) =>
+      makeFarGenerator(
+        (async function* () {
+          for (let i = 0; i < stop; i++) {
+            tlog(`${name} generating ${i}`);
+            yield i;
+          }
+          // Note the IIFE.
+        })(),
+      ),
+
+    iterate: async (producer) => {
+      const remoteGenerator = await E(producer).generate(5);
+      for await (const value of makeEventualIterator(remoteGenerator)) {
+        tlog(`${name} iterating ${value}`);
+      }
+    },
+  });
+}

packages/streams/package.json

@@ -39,6 +39,16 @@
         "default": "./dist/browser/index.cjs"
       }
     },
+    "./vat": {
+      "import": {
+        "types": "./dist/vat/index.d.mts",
+        "default": "./dist/vat/index.mjs"
+      },
+      "require": {
+        "types": "./dist/vat/index.d.cts",
+        "default": "./dist/vat/index.cjs"
+      }
+    },
     "./package.json": "./package.json"
   },
   "main": "./dist/index.cjs",
@@ -68,6 +78,9 @@
     "test:watch": "vitest --config vitest.config.ts"
   },
   "dependencies": {
+    "@endo/exo": "^1.5.9",
+    "@endo/far": "^1.1.11",
+    "@endo/patterns": "^1.5.0",
     "@endo/promise-kit": "^1.1.10",
     "@endo/stream": "^1.2.10",
     "@metamask/kernel-errors": "workspace:^",

packages/streams/src/index.test.ts

@@ -8,6 +8,8 @@ describe('index', () => {
       'NodeWorkerDuplexStream',
       'NodeWorkerReader',
       'NodeWorkerWriter',
+      'makeEventualIterator',
+      'makeFarGenerator',
       'split',
     ]);
   });

packages/streams/src/index.ts

@@ -6,3 +6,5 @@ export {
   NodeWorkerDuplexStream,
 } from './node/NodeWorkerStream.ts';
 export { split } from './split.ts';
+export { makeEventualIterator } from './vat/eventual-iterator.ts';
+export { makeFarGenerator } from './vat/far-generator.ts';

packages/streams/src/vat/eventual-iterator.ts

@@ -0,0 +1,21 @@
+// This type is used in the docs.
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+import type { makeFarGenerator } from './far-generator.ts';
+import { makeRefIterator } from './ref-reader.ts';
+
+/**
+ * Make an iterator from a remote generator. Intended to be used in conjunction
+ * with {@link makeFarGenerator}. This is the consuming end of the pair.
+ *
+ * Enables async iterator syntax like below.
+ * ```ts
+ * const eventualIterator = makeEventualIterator(remoteGeneratorRef);
+ * for await (const value of eventualIterator) {
+ *   console.log(`A faraway vat yielded: ${value}`);
+ * }
+ * ```
+ *
+ * @param iteratorRef - The remotable presence to make an iterator from.
+ * @returns An iterator that wraps the remotable presence, enabling async iterator syntax.
+ */
+export const makeEventualIterator = makeRefIterator;

packages/streams/src/vat/far-generator.test.ts

@@ -0,0 +1,90 @@
+import { E } from '@endo/far';
+import { makePromiseKit } from '@endo/promise-kit';
+import type { Reader, Writer } from '@endo/stream';
+import { makePipe } from '@endo/stream';
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import type { Mocked } from 'vitest';
+
+import { makeFarGenerator } from './far-generator.ts';
+
+vi.mock('@endo/far', () => ({
+  E: vi.fn((obj) => obj),
+}));
+
+vi.mock('@endo/stream', () => ({
+  makePipe: vi.fn(),
+}));
+
+vi.mock('./reader-ref.ts', () => ({
+  makeIteratorRef: vi.fn((reader) => reader),
+}));
+
+describe('far-generator', () => {
+  const mockWriter: Mocked<Writer<unknown>> = {
+    next: vi.fn(),
+    return: vi.fn(),
+    throw: vi.fn(),
+    [Symbol.asyncIterator]: vi.fn(),
+  };
+  const mockReader: Mocked<Reader<unknown>> = {
+    next: vi.fn(),
+    return: vi.fn(),
+    throw: vi.fn(),
+    [Symbol.asyncIterator]: vi.fn(() => mockReader),
+  };
+  beforeEach(() => {
+    vi.clearAllMocks();
+    vi.mocked(makePipe).mockReturnValue([mockWriter, mockReader]);
+  });
+
+  describe('makeFarGenerator', () => {
+    it('should wrap a generator', async () => {
+      const result = makeFarGenerator(
+        (async function* () {
+          // Empty generator
+        })(),
+      );
+
+      // Verify that makeIteratorRef was called with the reader
+      const { makeIteratorRef } = await import('./reader-ref.ts');
+      expect(makeIteratorRef).toHaveBeenCalledTimes(1);
+      expect(result).toBeDefined();
+    });
+
+    it('should pipe values from the generator to the writer', async () => {
+      const generated = makePromiseKit<string>();
+      const result = makeFarGenerator(
+        (async function* () {
+          yield 'test';
+          generated.resolve('yielded');
+        })(),
+      );
+
+      await E(result).next();
+
+      expect(await generated.promise).toBe('yielded');
+
+      expect(mockWriter.next).toHaveBeenCalledOnce();
+      expect(mockWriter.next).toHaveBeenCalledWith('test');
+    });
+
+    it('calls writer.throw if the generator throws', async () => {
+      const generated = makePromiseKit<string>();
+      const error = new Error('test');
+      const result = makeFarGenerator(
+        // eslint-disable-next-line require-yield
+        (async function* () {
+          generated.resolve('threw');
+          throw error;
+        })(),
+      );
+
+      await E(result).next();
+
+      expect(await generated.promise).toBe('threw');
+
+      expect(mockWriter.throw).toHaveBeenCalledOnce();
+      expect(mockWriter.throw).toHaveBeenCalledWith(error);
+    });
+  });
+});

packages/streams/src/vat/far-generator.ts

@@ -0,0 +1,32 @@
+import { E } from '@endo/far';
+import type { FarRef } from '@endo/far';
+import { makePipe } from '@endo/stream';
+import type { Writer, Reader } from '@endo/stream';
+
+// This type is used in the docs.
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+import type { makeEventualIterator } from './eventual-iterator.ts';
+import { makeIteratorRef } from './reader-ref.ts';
+
+/**
+ * Make a remotable generator. Intended to be used in conjunction with
+ * {@link makeEventualIterator}. This is the producing end of the pair.
+ *
+ * @param generator - The generator to make remotable.
+ * @returns A remotable reference to the generator.
+ */
+export const makeFarGenerator = <Item>(
+  generator: AsyncGenerator<Item>,
+): FarRef<AsyncIterator<Item>> => {
+  const [writer, reader] = makePipe<Item>() as unknown as [
+    Writer<Item>,
+    Reader<Item>,
+  ];
+  (async () => {
+    for await (const value of generator) {
+      await E(writer).next(value);
+    }
+    await E(writer).return(undefined);
+  })().catch(async (error) => await E(writer).throw(error));
+  return makeIteratorRef<Item>(reader);
+};

packages/streams/src/vat/index.test.ts

@@ -0,0 +1,12 @@
+import { describe, it, expect } from 'vitest';
+
+import * as indexModule from './index.ts';
+
+describe('index', () => {
+  it('has the expected exports', () => {
+    expect(Object.keys(indexModule).sort()).toStrictEqual([
+      'makeEventualIterator',
+      'makeFarGenerator',
+    ]);
+  });
+});