ci(tests): rebalance extension shards by estimated cost

scripts/lib/extension-test-plan.mjs

@@ -21,6 +21,13 @@ import { listAvailableExtensionIds } from "./changed-extensions.mjs";
 
 const repoRoot = path.resolve(import.meta.dirname, "..", "..");
 export const DEFAULT_EXTENSION_TEST_SHARD_COUNT = 6;
+const EXTENSION_TEST_COST_MULTIPLIERS = {
+  "test/vitest/vitest.extension-feishu.config.ts": 1.6,
+  "test/vitest/vitest.extension-msteams.config.ts": 1.6,
+  // This shared config is comparatively cheap per file, so raw file count
+  // overstates its real wall-clock cost during CI shard planning.
+  "test/vitest/vitest.extensions.config.ts": 0.45,
+};
 
 function normalizeRelative(inputPath) {
   return inputPath.split(path.sep).join("/");
@@ -53,6 +60,11 @@ function countTestFiles(rootPath) {
   return total;
 }
 
+function estimatePlanCost(config, testFileCount) {
+  const multiplier = EXTENSION_TEST_COST_MULTIPLIERS[config] ?? 1;
+  return Math.max(1, Math.ceil(testFileCount * multiplier));
+}
+
 function resolveExtensionDirectory(targetArg, cwd = process.cwd()) {
   if (targetArg) {
     const asGiven = path.resolve(cwd, targetArg);
@@ -152,9 +164,11 @@ export function resolveExtensionTestPlan(params = {}) {
     (sum, root) => sum + countTestFiles(path.join(repoRoot, root)),
     0,
   );
+  const estimatedCost = estimatePlanCost(config, testFileCount);
 
   return {
     config,
+    estimatedCost,
     extensionDir: relativeExtensionDir,
     extensionId,
     hasTests: testFileCount > 0,
@@ -171,11 +185,13 @@ function mergeTestPlans(plans) {
       config: plan.config,
       extensionIds: [],
       roots: [],
+      estimatedCost: 0,
       testFileCount: 0,
     };
 
     current.extensionIds.push(plan.extensionId);
     current.roots.push(...plan.roots);
+    current.estimatedCost += plan.estimatedCost;
     current.testFileCount += plan.testFileCount;
     groupsByConfig.set(plan.config, current);
   }
@@ -193,6 +209,7 @@ function mergeTestPlans(plans) {
     extensionIds: plans
       .map((plan) => plan.extensionId)
       .toSorted((left, right) => left.localeCompare(right)),
+    estimatedCost: plans.reduce((sum, plan) => sum + plan.estimatedCost, 0),
     hasTests: plans.length > 0,
     planGroups,
     testFileCount: plans.reduce((sum, plan) => sum + plan.testFileCount, 0),
@@ -215,6 +232,9 @@ function pickLeastLoadedShard(shards) {
       return index;
     }
     const best = shards[bestIndex];
+    if (shard.estimatedCost !== best.estimatedCost) {
+      return shard.estimatedCost < best.estimatedCost ? index : bestIndex;
+    }
     if (shard.testFileCount !== best.testFileCount) {
       return shard.testFileCount < best.testFileCount ? index : bestIndex;
     }
@@ -233,6 +253,9 @@ export function createExtensionTestShards(params = {}) {
     .map((extensionId) => resolveExtensionTestPlan({ cwd, targetArg: extensionId }))
     .filter((plan) => plan.hasTests)
     .toSorted((left, right) => {
+      if (left.estimatedCost !== right.estimatedCost) {
+        return right.estimatedCost - left.estimatedCost;
+      }
       if (left.testFileCount !== right.testFileCount) {
         return right.testFileCount - left.testFileCount;
       }
@@ -241,6 +264,7 @@ export function createExtensionTestShards(params = {}) {
 
   const effectiveShardCount = Math.min(shardCount, Math.max(1, plans.length));
   const shards = Array.from({ length: effectiveShardCount }, () => ({
+    estimatedCost: 0,
     plans: [],
     testFileCount: 0,
   }));
@@ -248,6 +272,7 @@ export function createExtensionTestShards(params = {}) {
   for (const plan of plans) {
     const targetIndex = pickLeastLoadedShard(shards);
     shards[targetIndex].plans.push(plan);
+    shards[targetIndex].estimatedCost += plan.estimatedCost;
     shards[targetIndex].testFileCount += plan.testFileCount;
   }
 

test/scripts/test-extension.test.ts

@@ -282,96 +282,112 @@ describe("scripts/test-extension.mjs", () => {
     expect(batch.planGroups).toEqual([
       {
         config: "test/vitest/vitest.extension-acpx.config.ts",
+        estimatedCost: expect.any(Number),
         extensionIds: ["acpx"],
         roots: [bundledPluginRoot("acpx")],
         testFileCount: expect.any(Number),
       },
       {
         config: "test/vitest/vitest.extension-bluebubbles.config.ts",
+        estimatedCost: expect.any(Number),
         extensionIds: ["bluebubbles"],
         roots: [bundledPluginRoot("bluebubbles")],
         testFileCount: expect.any(Number),
       },
       {
         config: "test/vitest/vitest.extension-channels.config.ts",
+        estimatedCost: expect.any(Number),
         extensionIds: ["line", "slack"],
         roots: [bundledPluginRoot("slack"), bundledPluginRoot("line")],
         testFileCount: expect.any(Number),
       },
       {
         config: "test/vitest/vitest.extension-diffs.config.ts",
+        estimatedCost: expect.any(Number),
         extensionIds: ["diffs"],
         roots: [bundledPluginRoot("diffs")],
         testFileCount: expect.any(Number),
       },
       {
         config: "test/vitest/vitest.extension-feishu.config.ts",
+        estimatedCost: expect.any(Number),
         extensionIds: ["feishu"],
         roots: [bundledPluginRoot("feishu")],
         testFileCount: expect.any(Number),
       },
       {
         config: "test/vitest/vitest.extension-irc.config.ts",
+        estimatedCost: expect.any(Number),
         extensionIds: ["irc"],
         roots: [bundledPluginRoot("irc")],
         testFileCount: expect.any(Number),
       },
       {
         config: "test/vitest/vitest.extension-matrix.config.ts",
+        estimatedCost: expect.any(Number),
         extensionIds: ["matrix"],
         roots: [bundledPluginRoot("matrix")],
         testFileCount: expect.any(Number),
       },
       {
         config: "test/vitest/vitest.extension-mattermost.config.ts",
+        estimatedCost: expect.any(Number),
         extensionIds: ["mattermost"],
         roots: [bundledPluginRoot("mattermost")],
         testFileCount: expect.any(Number),
       },
       {
         config: "test/vitest/vitest.extension-memory.config.ts",
+        estimatedCost: expect.any(Number),
         extensionIds: ["memory-core"],
         roots: [bundledPluginRoot("memory-core")],
         testFileCount: expect.any(Number),
       },
       {
         config: "test/vitest/vitest.extension-msteams.config.ts",
+        estimatedCost: expect.any(Number),
         extensionIds: ["msteams"],
         roots: [bundledPluginRoot("msteams")],
         testFileCount: expect.any(Number),
       },
       {
         config: "test/vitest/vitest.extension-providers.config.ts",
+        estimatedCost: expect.any(Number),
         extensionIds: ["openai"],
         roots: [bundledPluginRoot("openai")],
         testFileCount: expect.any(Number),
       },
       {
         config: "test/vitest/vitest.extension-telegram.config.ts",
+        estimatedCost: expect.any(Number),
         extensionIds: ["telegram"],
         roots: [bundledPluginRoot("telegram")],
         testFileCount: expect.any(Number),
       },
       {
         config: "test/vitest/vitest.extension-voice-call.config.ts",
+        estimatedCost: expect.any(Number),
         extensionIds: ["voice-call"],
         roots: [bundledPluginRoot("voice-call")],
         testFileCount: expect.any(Number),
       },
       {
         config: "test/vitest/vitest.extension-whatsapp.config.ts",
+        estimatedCost: expect.any(Number),
         extensionIds: ["whatsapp"],
         roots: [bundledPluginRoot("whatsapp")],
         testFileCount: expect.any(Number),
       },
       {
         config: "test/vitest/vitest.extension-zalo.config.ts",
+        estimatedCost: expect.any(Number),
         extensionIds: ["zalo", "zalouser"],
         roots: [bundledPluginRoot("zalo"), bundledPluginRoot("zalouser")],
         testFileCount: expect.any(Number),
       },
       {
         config: "test/vitest/vitest.extensions.config.ts",
+        estimatedCost: expect.any(Number),
         extensionIds: ["firecrawl"],
         roots: [bundledPluginRoot("firecrawl")],
         testFileCount: expect.any(Number),
@@ -379,7 +395,7 @@ describe("scripts/test-extension.mjs", () => {
     ]);
   });
 
-  it("balances extension test shards by test file count", () => {
+  it("balances extension test shards by estimated CI cost", () => {
     const shards = createExtensionTestShards({
       cwd: process.cwd(),
       shardCount: DEFAULT_EXTENSION_TEST_SHARD_COUNT,
@@ -402,8 +418,15 @@ describe("scripts/test-extension.mjs", () => {
     );
     expect(assigned).toHaveLength(expected.length);
 
-    const totals = shards.map((shard) => shard.testFileCount);
+    const totals = shards.map((shard) => shard.estimatedCost);
     expect(Math.max(...totals) - Math.min(...totals)).toBeLessThanOrEqual(1);
+
+    const msTeamsShardIndex = shards.findIndex((shard) => shard.extensionIds.includes("msteams"));
+    const feishuShardIndex = shards.findIndex((shard) => shard.extensionIds.includes("feishu"));
+
+    expect(msTeamsShardIndex).toBeGreaterThanOrEqual(0);
+    expect(feishuShardIndex).toBeGreaterThanOrEqual(0);
+    expect(msTeamsShardIndex).not.toBe(feishuShardIndex);
   });
 
   it("treats extensions without tests as a no-op by default", () => {