ci: parallelize extension batch groups

.github/workflows/ci.yml

@@ -862,6 +862,7 @@ jobs:
 
       - name: Run extension shard
         env:
+          OPENCLAW_EXTENSION_BATCH_PARALLEL: 2
           OPENCLAW_EXTENSION_BATCH: ${{ matrix.extensions_csv }}
         run: pnpm test:extensions:batch -- "$OPENCLAW_EXTENSION_BATCH"
 

scripts/test-extension-batch.mjs

@@ -1,8 +1,12 @@
 #!/usr/bin/env node
 
+import path from "node:path";
 import { resolveExtensionBatchPlan } from "./lib/extension-test-plan.mjs";
 import { isDirectScriptRun, runVitestBatch } from "./lib/vitest-batch-runner.mjs";
 
+const FS_MODULE_CACHE_PATH_ENV_KEY = "OPENCLAW_VITEST_FS_MODULE_CACHE_PATH";
+const PARALLEL_ENV_KEY = "OPENCLAW_EXTENSION_BATCH_PARALLEL";
+
 function printUsage() {
   console.error("Usage: pnpm test:extensions:batch <extension[,extension...]> [vitest args...]");
   console.error(
@@ -27,6 +31,114 @@ function parseExtensionIds(rawArgs) {
   return { extensionIds, passthroughArgs: args };
 }
 
+function parsePositiveInt(value) {
+  const parsed = Number.parseInt(value ?? "", 10);
+  return Number.isFinite(parsed) && parsed > 0 ? parsed : null;
+}
+
+export function resolveExtensionBatchParallelism(groupCount, env = process.env) {
+  const override = parsePositiveInt(env[PARALLEL_ENV_KEY]);
+  return Math.min(Math.max(1, override ?? 1), Math.max(1, groupCount));
+}
+
+function sanitizeCacheSegment(value) {
+  return (
+    value
+      .replace(/[^a-zA-Z0-9._-]+/gu, "-")
+      .replace(/^-+|-+$/gu, "")
+      .slice(0, 180) || "default"
+  );
+}
+
+function createGroupEnv({ baseEnv, group, groupIndex, useDedicatedCache }) {
+  if (!useDedicatedCache || baseEnv[FS_MODULE_CACHE_PATH_ENV_KEY]?.trim()) {
+    return baseEnv;
+  }
+
+  return {
+    ...baseEnv,
+    [FS_MODULE_CACHE_PATH_ENV_KEY]: path.join(
+      process.cwd(),
+      "node_modules",
+      ".experimental-vitest-cache",
+      "extension-batch",
+      sanitizeCacheSegment(`${groupIndex}-${group.config}`),
+    ),
+  };
+}
+
+function orderPlanGroups(planGroups, parallelism) {
+  if (parallelism <= 1) {
+    return planGroups;
+  }
+  return [...planGroups].toSorted((left, right) => {
+    if (left.estimatedCost !== right.estimatedCost) {
+      return right.estimatedCost - left.estimatedCost;
+    }
+    if (left.testFileCount !== right.testFileCount) {
+      return right.testFileCount - left.testFileCount;
+    }
+    return left.config.localeCompare(right.config);
+  });
+}
+
+async function runPlanGroup(group, params) {
+  console.log(
+    `[test-extension-batch] ${group.config}: ${group.extensionIds.join(", ")} (${group.testFileCount} files)`,
+  );
+  return await params.runGroup({
+    args: params.vitestArgs,
+    config: group.config,
+    env: createGroupEnv({
+      baseEnv: params.env,
+      group,
+      groupIndex: params.groupIndex,
+      useDedicatedCache: params.useDedicatedCache,
+    }),
+    targets: group.roots,
+  });
+}
+
+export async function runExtensionBatchPlan(batchPlan, params = {}) {
+  const env = params.env ?? process.env;
+  const vitestArgs = params.vitestArgs ?? [];
+  const runGroup = params.runGroup ?? runVitestBatch;
+  const parallelism = resolveExtensionBatchParallelism(batchPlan.planGroups.length, env);
+  const orderedGroups = orderPlanGroups(batchPlan.planGroups, parallelism);
+  const useDedicatedCache = parallelism > 1;
+
+  if (parallelism > 1) {
+    console.log(`[test-extension-batch] Running up to ${parallelism} config groups in parallel`);
+  }
+
+  let nextGroupIndex = 0;
+  let exitCode = 0;
+  async function worker() {
+    while (exitCode === 0) {
+      const groupIndex = nextGroupIndex;
+      nextGroupIndex += 1;
+      const group = orderedGroups[groupIndex];
+      if (!group) {
+        return;
+      }
+      const groupExitCode = await runPlanGroup(group, {
+        env,
+        groupIndex,
+        runGroup,
+        useDedicatedCache,
+        vitestArgs,
+      });
+      if (groupExitCode !== 0) {
+        exitCode = groupExitCode;
+        return;
+      }
+    }
+  }
+
+  await Promise.all(Array.from({ length: parallelism }, () => worker()));
+  return exitCode;
+}
+
 async function run() {
   const rawArgs = process.argv.slice(2);
   if (rawArgs.includes("--help") || rawArgs.includes("-h")) {
@@ -51,19 +163,12 @@ async function run() {
     `[test-extension-batch] Running ${batchPlan.testFileCount} test files across ${batchPlan.extensionCount} extensions`,
   );
 
-  for (const group of batchPlan.planGroups) {
-    console.log(
-      `[test-extension-batch] ${group.config}: ${group.extensionIds.join(", ")} (${group.testFileCount} files)`,
-    );
-    const exitCode = await runVitestBatch({
-      args: vitestArgs,
-      config: group.config,
-      env: process.env,
-      targets: group.roots,
-    });
-    if (exitCode !== 0) {
-      process.exit(exitCode);
-    }
+  const exitCode = await runExtensionBatchPlan(batchPlan, {
+    env: process.env,
+    vitestArgs,
+  });
+  if (exitCode !== 0) {
+    process.exit(exitCode);
   }
 }
 

scripts/test-projects.test-support.mjs

@@ -189,6 +189,9 @@ const TOOLING_SOURCE_TEST_TARGETS = new Map([
     "scripts/run-vitest.mjs",
     ["test/scripts/test-projects.test.ts", "test/scripts/vitest-local-scheduling.test.ts"],
   ],
+  ["scripts/test-extension-batch.mjs", ["test/scripts/test-extension.test.ts"]],
+  ["scripts/lib/extension-test-plan.mjs", ["test/scripts/test-extension.test.ts"]],
+  ["scripts/lib/vitest-batch-runner.mjs", ["test/scripts/test-extension.test.ts"]],
   ["scripts/test-projects.mjs", ["test/scripts/test-projects.test.ts"]],
   ["scripts/test-projects.test-support.d.mts", ["test/scripts/test-projects.test.ts"]],
   ["scripts/test-projects.test-support.mjs", ["test/scripts/test-projects.test.ts"]],

test/scripts/test-extension.test.ts

@@ -1,6 +1,6 @@
 import { execFileSync } from "node:child_process";
 import path from "node:path";
-import { describe, expect, it } from "vitest";
+import { describe, expect, it, vi } from "vitest";
 import {
   detectChangedExtensionIds,
   listAvailableExtensionIds,
@@ -12,6 +12,10 @@ import {
   resolveExtensionBatchPlan,
   resolveExtensionTestPlan,
 } from "../../scripts/lib/extension-test-plan.mjs";
+import {
+  resolveExtensionBatchParallelism,
+  runExtensionBatchPlan,
+} from "../../scripts/test-extension-batch.mjs";
 import { bundledPluginFile, bundledPluginRoot } from "../helpers/bundled-plugin-paths.js";
 
 const scriptPath = path.join(process.cwd(), "scripts", "test-extension.mjs");
@@ -436,6 +440,88 @@ describe("scripts/test-extension.mjs", () => {
     expect(msTeamsShardIndex).not.toBe(feishuShardIndex);
   });
 
+  it("runs extension batch config groups concurrently when requested", async () => {
+    const started: string[] = [];
+    const resolvers: Array<() => void> = [];
+    const runGroup = vi.fn(
+      (params: {
+        args: string[];
+        config: string;
+        env: Record<string, string | undefined>;
+        targets: string[];
+      }) => {
+        started.push(params.config);
+        return new Promise<number>((resolve) => {
+          resolvers.push(() => resolve(0));
+        });
+      },
+    );
+    const runPromise = runExtensionBatchPlan(
+      {
+        extensionCount: 3,
+        extensionIds: ["one", "two", "three"],
+        estimatedCost: 60,
+        hasTests: true,
+        planGroups: [
+          {
+            config: "light",
+            estimatedCost: 10,
+            extensionIds: ["one"],
+            roots: ["extensions/one"],
+            testFileCount: 1,
+          },
+          {
+            config: "heavy",
+            estimatedCost: 30,
+            extensionIds: ["two"],
+            roots: ["extensions/two"],
+            testFileCount: 3,
+          },
+          {
+            config: "middle",
+            estimatedCost: 20,
+            extensionIds: ["three"],
+            roots: ["extensions/three"],
+            testFileCount: 2,
+          },
+        ],
+        testFileCount: 6,
+      },
+      {
+        env: { OPENCLAW_EXTENSION_BATCH_PARALLEL: "2" },
+        runGroup,
+        vitestArgs: ["--reporter=dot"],
+      },
+    );
+
+    await Promise.resolve();
+    expect(started).toEqual(["heavy", "middle"]);
+    resolvers.shift()?.();
+    await new Promise<void>((resolve) => setImmediate(resolve));
+    expect(started).toEqual(["heavy", "middle", "light"]);
+    while (resolvers.length > 0) {
+      resolvers.shift()?.();
+    }
+    await expect(runPromise).resolves.toBe(0);
+    expect(runGroup).toHaveBeenCalledTimes(3);
+    expect(runGroup.mock.calls[0]?.[0]).toMatchObject({
+      args: ["--reporter=dot"],
+      config: "heavy",
+      targets: ["extensions/two"],
+    });
+    expect(runGroup.mock.calls[0]?.[0].env.OPENCLAW_VITEST_FS_MODULE_CACHE_PATH).toContain(
+      path.join("node_modules", ".experimental-vitest-cache", "extension-batch", "0-heavy"),
+    );
+  });
+
+  it("keeps extension batch parallelism bounded by group count", () => {
+    expect(resolveExtensionBatchParallelism(3, { OPENCLAW_EXTENSION_BATCH_PARALLEL: "2" })).toBe(2);
+    expect(resolveExtensionBatchParallelism(1, { OPENCLAW_EXTENSION_BATCH_PARALLEL: "4" })).toBe(1);
+    expect(resolveExtensionBatchParallelism(3, { OPENCLAW_EXTENSION_BATCH_PARALLEL: "nope" })).toBe(
+      1,
+    );
+  });
+
   it("treats extensions without tests as a no-op by default", () => {
     const extensionId = findExtensionWithoutTests();
     const stdout = runScript([extensionId]);

test/scripts/test-projects.test.ts

@@ -44,6 +44,13 @@ describe("scripts/test-projects changed-target routing", () => {
     });
   });
 
+  it("keeps extension batch runner edits on extension script tests", () => {
+    expect(resolveChangedTestTargetPlan(["scripts/test-extension-batch.mjs"])).toEqual({
+      mode: "targets",
+      targets: ["test/scripts/test-extension.test.ts"],
+    });
+  });
+
   it("routes changed extension vitest configs to their own shard", () => {
     expect(
       buildVitestRunPlans(["--changed", "origin/main"], process.cwd(), () => [