test(core): run Code Mode's executor on plain workerd, and check what core keeps after many runs

.github/workflows/ci.yml

@@ -125,6 +125,8 @@ jobs:
             - args: ["--frozen-lockfile"]
       - run: vp run --filter @grasp-os/web build
       - run: vp run smoke:workerd
+      # Code Mode's executor, with a Worker Loader, on plain workerd.
+      - run: vp run smoke:code-mode
 
   # Builds a release bundle as main will publish it (core, connect, the
   # frontend's static files and D1 migrations, with its manifest) and checks

apps/core/src/workspace.ts

@@ -188,7 +188,7 @@ export const maxChatChars = 4_000_000;
  * Most ended code runs an object remembers, to refuse and log a call from
  * one: a few turns' worth (at most 30 runs each).
  */
-const endedRunsKept = 1000;
+export const endedRunsKept = 1000;
 
 /**
  * Most code run records a chat keeps, newest first: a long chat's run
@@ -1183,6 +1183,27 @@ export class Workspace extends DurableObject<Env> {
     return "revoked";
   }
 
+  /**
+   * What code runs keep in this object now, read from its own state:
+   * leases open and closed, and outputs held. After any number of runs,
+   * none open and no output, and at most {@link endedRunsKept} closed.
+   * What the runtime holds for them (isolates, RPC sessions) isn't here:
+   * workerd offers no signal of it to read.
+   */
+  codeRunRetention(): { open: number; closed: number; outputs: number } {
+    let open = 0;
+    for (const lease of this.#codeRuns.values()) {
+      if (!("closed" in lease)) {
+        open += 1;
+      }
+    }
+    return {
+      open,
+      closed: this.#codeRuns.size - open,
+      outputs: this.#outputs.size,
+    };
+  }
+
   /**
    * Takes a line a code run of the chat logged, for its watchers as the
    * code step `callId`'s output: counted and bounded here, whatever the

apps/core/test/agent.test.ts

@@ -29,6 +29,7 @@ import {
   utf8Bytes,
 } from "../src/code-mode.ts";
 import { workspace } from "../src/durable-objects.ts";
+import { endedRunsKept } from "../src/workspace.ts";
 import {
   chatOf,
   codeResults,
@@ -2511,3 +2512,76 @@ describe("chat agent code output", () => {
     });
   });
 });
+
+/** Orders text alphabetically. */
+const byText = (a: string, b: string): number => a.localeCompare(b, "en");
+
+describe("chat agent code run retention", () => {
+  it("keeps no more closed leases than it may, however many runs end", async () => {
+    const { stub, chat } = await newChat();
+    const chatId = chatIdSchema.parse(chat.id);
+    const runs = endedRunsKept + 50;
+
+    // Leases opened and closed as runs do, past the most it keeps.
+    const kept = await runInDurableObject(stub, (instance) => {
+      const ids: string[] = [];
+      for (let run = 0; run < runs; run += 1) {
+        const id = instance.openCodeRun(chatId, testRunOpening);
+        instance.closeCodeRun(chatId, id, "completed");
+        ids.push(id);
+      }
+      return {
+        retention: instance.codeRunRetention(),
+        // The oldest are forgotten, and answer as ended, unlogged.
+        oldest: instance.callFromCodeRun(chatId, ids[0] ?? ""),
+        newest: instance.callFromCodeRun(chatId, ids.at(-1) ?? ""),
+      };
+    });
+
+    expect(kept).toStrictEqual({
+      retention: { open: 0, closed: endedRunsKept, outputs: 0 },
+      oldest: { call: "ended", first: false },
+      newest: { call: "ended", first: true },
+    });
+  });
+
+  it("keeps no open lease or output after many runs, whether they complete, throw or fail to load", async () => {
+    const kinds = [
+      "export default async (self, env) => { console.log((await env.chat.info()).chatId); return 1; };",
+      "export default async (self, env) => { await env.chat.info(); throw new Error('nope'); };",
+      "export default async () => {",
+    ];
+    const turns = 4;
+    const replies = Array.from({ length: turns }, (_, turn) => [
+      codeStep(kinds[0] ?? "", 5),
+      codeStep(kinds[1 + (turn % 2)] ?? "", 5),
+      says("Done."),
+    ]).flat();
+    const { stub, chat, ask } = await newChat(...replies);
+
+    for (let turn = 0; turn < turns; turn += 1) {
+      // oxlint-disable-next-line no-await-in-loop -- one turn at a time per chat
+      await ask("Run them.");
+    }
+
+    // Read from the object's own state after the runs, not inferred from
+    // their cleanup. What the runtime holds for them can't be read.
+    await expect(
+      runInDurableObject(stub, (instance) => instance.codeRunRetention())
+    ).resolves.toStrictEqual({
+      open: 0,
+      closed: turns * 10,
+      outputs: 0,
+    });
+    const records = await runRecords(stub, chat.id);
+    expect(
+      records.map(({ status }) => String(status)).toSorted(byText)
+    ).toStrictEqual(
+      [
+        ...Array.from({ length: 20 }, () => "completed"),
+        ...Array.from({ length: 10 }, () => "load_failed"),
+        ...Array.from({ length: 10 }, () => "threw"),
+      ].toSorted(byText)
+    );
+  });
+});

apps/core/test/workerd/code-mode-worker.ts

@@ -0,0 +1,139 @@
+import { WorkerEntrypoint, env, exports } from "cloudflare:workers";
+
+import {
+  codeLimits,
+  codeRuntime,
+  cutMarker,
+  describeRun,
+  runCode,
+  utf8Bytes,
+} from "../../src/code-mode.ts";
+import type { CodeRun } from "../../src/code-mode.ts";
+
+// Code Mode's executor on plain workerd (scripts/workerd-code-mode.ts):
+// the same `runCode` core runs, with a Worker Loader and a capability of
+// this Worker's own as the code's env, and no Wrangler or Miniflare. Each
+// case runs one module and says how it ended; the script checks them.
+
+/** Calls the capability took, across every case. */
+let calls = 0;
+
+/** A capability the code is given, as core gives its APIs: a loopback stub. */
+export class Caps extends WorkerEntrypoint {
+  // oxlint-disable-next-line class-methods-use-this -- an RPC method, called on the entrypoint
+  async double(value: number): Promise<number> {
+    calls += 1;
+    await scheduler.wait(1);
+    return value * 2;
+  }
+}
+
+interface Case {
+  code: string;
+  timeoutMs?: number;
+  /** Cancels the run this long after it starts. */
+  cancelAfterMs?: number;
+}
+
+const cases: Record<string, Case> = {
+  complete: {
+    code: "const twice = (env, n) => env.caps.double(n); export default async (self, env, ctx) => ({ self, ctx: Object.keys(ctx), sum: (await Promise.all([twice(env, 1), twice(env, 2)])).reduce((a, b) => a + b, 0) });",
+  },
+  network: {
+    code: "export default async () => (await fetch('https://example.com')).status;",
+  },
+  syntax: { code: "export default async () => {" },
+  notAFunction: { code: "export default 42;" },
+  throwsAsItLoads: {
+    code: "console.log('loading'); throw new Error('broken module'); export default async () => 1;",
+  },
+  throws: { code: "export default async () => { throw new Error('nope'); };" },
+  oversized: {
+    code: `export default async () => 1; // ${"x".repeat(codeLimits.sourceChars)}`,
+  },
+  neverSettles: {
+    code: "export default async () => { await new Promise(() => {}); };",
+  },
+  waitsPastItsTime: {
+    code: "export default async () => { await scheduler.wait(60_000); };",
+    timeoutMs: 1000,
+  },
+  cancelled: {
+    code: "export default async () => { await scheduler.wait(60_000); };",
+    cancelAfterMs: 300,
+  },
+  manyLines: {
+    code: "export default async () => { for (let i = 0; i < 5000; i++) console.log(i); return 'done'; };",
+  },
+  flood: {
+    code: "export default async () => { for (let i = 0; i < 5000; i++) console.log('é'.repeat(200)); throw new Error('the end'); };",
+  },
+};
+
+/** Whether `value` is a stub a loaded Worker can be given in its env. */
+const isFetcher = (value: unknown): value is Fetcher =>
+  (typeof value === "object" || typeof value === "function") &&
+  value !== null &&
+  "fetch" in value;
+
+/** A loopback stub of {@link Caps}, as core makes its APIs' stubs. */
+const capability = (): Fetcher => {
+  const make: unknown = Reflect.get(exports, "Caps");
+  if (typeof make !== "function") {
+    throw new TypeError("This Worker exports no Caps");
+  }
+  const stub: unknown = Reflect.apply(make, exports, [{}]);
+  if (!isFetcher(stub)) {
+    throw new TypeError("Caps gave no stub");
+  }
+  return stub;
+};
+
+/** One case: how it ended, what the model would read, what streamed. */
+const runCase = async ({ code, timeoutMs, cancelAfterMs }: Case) => {
+  const cancel = new AbortController();
+  if (cancelAfterMs !== undefined) {
+    setTimeout(() => {
+      cancel.abort();
+    }, cancelAfterMs);
+  }
+  const streamed: string[] = [];
+  const started = Date.now();
+  const run: CodeRun = await runCode(
+    env.LOADER,
+    code,
+    { caps: capability() },
+    {
+      signal: cancel.signal,
+      ...(timeoutMs === undefined ? {} : { timeoutMs }),
+      onOutput: (line) => {
+        streamed.push(line);
+      },
+    }
+  );
+  const text = describeRun(run);
+  return {
+    status: run.status,
+    text: text.slice(0, 400),
+    ending: text.slice(-200),
+    ms: Date.now() - started,
+    bytes: utf8Bytes(text),
+    cut: text.includes(cutMarker),
+    streamed: streamed.length,
+  };
+};
+
+export default {
+  async fetch(): Promise<Response> {
+    const results: Record<string, unknown> = {};
+    for (const [name, each] of Object.entries(cases)) {
+      // oxlint-disable-next-line no-await-in-loop -- one run at a time, as a chat runs them
+      results[name] = await runCase(each);
+    }
+    return Response.json({
+      runtime: codeRuntime({ RUNTIME_PROFILE: "workerd" }),
+      calls,
+      results,
+    });
+  },
+};

package.json

@@ -9,6 +9,7 @@
     "dev": "node scripts/dev.ts",
     "e2e": "playwright test",
     "smoke:workerd": "node scripts/workerd-smoke.ts",
+    "smoke:code-mode": "node scripts/workerd-code-mode.ts",
     "check:docs": "node scripts/check-docs.ts",
     "check:migrations": "node scripts/check-migrations.ts",
     "check:wrangler-configs": "node scripts/check-wrangler-configs.ts",