From ecc8fe5dc27eb9b8bc59910fdfbf008c676a3940 Mon Sep 17 00:00:00 2001
From: Peter Steinberger <steipete@gmail.com>
Date: Tue, 24 Mar 2026 01:43:46 +0000
Subject: [PATCH] ci: rebalance sharded test lanes

---
 .github/workflows/ci-bun.yml | 32 +++++++++++++++++++--
 .github/workflows/ci.yml     |  9 ++++--
 scripts/test-parallel.mjs    | 54 ++++++++++++++++++++++++++++++++++++
 3 files changed, 91 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/ci-bun.yml b/.github/workflows/ci-bun.yml
index eba544c3cc5..0501779fc6c 100644
--- a/.github/workflows/ci-bun.yml
+++ b/.github/workflows/ci-bun.yml
@@ -12,7 +12,32 @@ env:
   FORCE_JAVASCRIPT_ACTIONS_TO_NODE24: "true"
 
 jobs:
+  build-bun-artifacts:
+    runs-on: blacksmith-16vcpu-ubuntu-2404
+    timeout-minutes: 20
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v6
+        with:
+          submodules: false
+
+      - name: Setup Node environment
+        uses: ./.github/actions/setup-node-env
+        with:
+          install-bun: "false"
+          use-sticky-disk: "false"
+
+      - name: Build A2UI bundle
+        run: pnpm canvas:a2ui:bundle
+
+      - name: Upload A2UI bundle artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: canvas-a2ui-bundle
+          path: src/canvas-host/a2ui/
+
   bun-checks:
+    needs: [build-bun-artifacts]
     runs-on: blacksmith-16vcpu-ubuntu-2404
     timeout-minutes: 20
     strategy:
@@ -37,8 +62,11 @@ jobs:
           install-bun: "true"
           use-sticky-disk: "false"
 
-      - name: Build A2UI bundle
-        run: pnpm canvas:a2ui:bundle
+      - name: Download A2UI bundle artifact
+        uses: actions/download-artifact@v8
+        with:
+          name: canvas-a2ui-bundle
+          path: src/canvas-host/a2ui/
 
       - name: Run Bun test shard
         run: ${{ matrix.command }}
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index c9234692e70..be1db16f9d7 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -257,7 +257,7 @@ jobs:
           - runtime: node
             task: channels
             shard_index: 1
-            shard_count: 2
+            shard_count: 3
             command: pnpm test:channels
           - runtime: node
             task: contracts
@@ -265,7 +265,12 @@ jobs:
           - runtime: node
             task: channels
             shard_index: 2
-            shard_count: 2
+            shard_count: 3
+            command: pnpm test:channels
+          - runtime: node
+            task: channels
+            shard_index: 3
+            shard_count: 3
             command: pnpm test:channels
           - runtime: node
             task: protocol
diff --git a/scripts/test-parallel.mjs b/scripts/test-parallel.mjs
index 0d8bf5896bb..ed3cc6ca534 100644
--- a/scripts/test-parallel.mjs
+++ b/scripts/test-parallel.mjs
@@ -807,6 +807,53 @@ const targetedEntries = (() => {
     return [createTargetedEntry(owner, false, uniqueFilters)];
   }).flat();
 })();
+const estimateTopLevelEntryDurationMs = (entry) => {
+  const filters = getExplicitEntryFilters(entry.args);
+  if (filters.length === 0) {
+    return unitTimingManifest.defaultDurationMs;
+  }
+  return filters.reduce((totalMs, file) => {
+    if (isUnitConfigTestFile(file)) {
+      return totalMs + estimateUnitDurationMs(file);
+    }
+    if (channelTestPrefixes.some((prefix) => file.startsWith(prefix))) {
+      return totalMs + 3_000;
+    }
+    if (file.startsWith("extensions/")) {
+      return totalMs + 2_000;
+    }
+    return totalMs + 1_000;
+  }, 0);
+};
+const topLevelSingleShardAssignments = (() => {
+  if (shardIndexOverride === null || shardCount <= 1) {
+    return new Map();
+  }
+
+  // Single-file and other non-shardable explicit lanes would otherwise run on
+  // every shard. Assign them to one top-level shard instead.
+  const entriesNeedingAssignment = runs.filter((entry) => {
+    const explicitFilterCount = countExplicitEntryFilters(entry.args);
+    if (explicitFilterCount === null) {
+      return false;
+    }
+    const effectiveShardCount = Math.min(shardCount, Math.max(1, explicitFilterCount - 1));
+    return effectiveShardCount <= 1;
+  });
+
+  const assignmentMap = new Map();
+  const buckets = packFilesByDuration(
+    entriesNeedingAssignment,
+    shardCount,
+    estimateTopLevelEntryDurationMs,
+  );
+  for (const [bucketIndex, bucket] of buckets.entries()) {
+    for (const entry of bucket) {
+      assignmentMap.set(entry, bucketIndex + 1);
+    }
+  }
+  return assignmentMap;
+})();
 // Node 25 local runs still show cross-process worker shutdown contention even
 // after moving the known heavy files into singleton lanes.
 const topLevelParallelEnabled =
@@ -1258,6 +1305,13 @@ const runOnce = (entry, extraArgs = []) =>
 
 const run = async (entry, extraArgs = []) => {
   const explicitFilterCount = countExplicitEntryFilters(entry.args);
+  const topLevelAssignedShard = topLevelSingleShardAssignments.get(entry);
+  if (topLevelAssignedShard !== undefined) {
+    if (shardIndexOverride !== null && shardIndexOverride !== topLevelAssignedShard) {
+      return 0;
+    }
+    return runOnce(entry, extraArgs);
+  }
   // Vitest requires the shard count to stay strictly below the number of
   // resolved test files, so explicit-filter lanes need a `< fileCount` cap.
   const effectiveShardCount =