fix(core): batch processes should be tracked by sandboxing [Self-Healing CI Rerun]

packages/nx/src/tasks-runner/forked-process-task-runner.ts

@@ -7,7 +7,6 @@ import { output } from '../utils/output';
 import { stripIndents } from '../utils/strip-indents';
 import { BatchMessageType } from './batch/batch-messages';
 import { DefaultTasksRunnerOptions } from './default-tasks-runner';
-import { getProcessMetricsService } from './process-metrics-service';
 import {
   createPseudoTerminal as createPseudoTerminalWithShutdown,
   PseudoTerminal,
@@ -19,7 +18,10 @@ import {
   NodeChildProcessWithNonDirectOutput,
 } from './running-tasks/node-child-process';
 import { RunningTask } from './running-tasks/running-task';
-import { registerTaskProcessStart } from './task-io-service';
+import {
+  registerBatchProcessStart,
+  registerTaskProcessStart,
+} from './task-io-service';
 import { Batch } from './tasks-schedule';
 import { getCliPath, getPrintableCommandArgsForTask } from './utils';
 
@@ -73,10 +75,10 @@ export class ForkedProcessTaskRunner {
       },
     });
 
-    // Register batch worker process with all tasks
+    // Register batch worker process with all tasks in both IO and metrics services
     if (p.pid) {
       const taskIds = Object.keys(batchTaskGraph.tasks);
-      getProcessMetricsService().registerBatch(batchId, taskIds, p.pid);
+      registerBatchProcessStart(batchId, taskIds, p.pid);
     }
 
     const cp = new BatchProcess(p, executorName);

packages/nx/src/tasks-runner/task-io-service.ts

@@ -50,6 +50,15 @@ class TaskIOService {
   private taskInputCallbacks: TaskInputCallback[] = [];
   private taskOutputsCallbacks: TaskOutputsCallback[] = [];
 
+  // Batch state: reverse lookup from taskId to its batchId
+  private taskIdToBatchId: Map<string, string> = new Map();
+  // Buffered inputs/outputs for batch tasks, keyed by batchId -> (taskId -> data)
+  private pendingBatchInputs: Map<
+    string,
+    Map<string, TaskInputInfo['inputs']>
+  > = new Map();
+  private pendingBatchOutputs: Map<string, Map<string, string[]>> = new Map();
+
   /**
    * Subscribe to task PID updates.
    * Receives notifications when processes are added/removed from tasks.
@@ -82,9 +91,24 @@ class TaskIOService {
     this.taskOutputsCallbacks.push(callback);
   }
 
+  /**
+   * Register a batch so that subsequent notifyTaskInputs/notifyTaskOutputs
+   * calls for tasks in this batch are buffered instead of emitted immediately.
+   * The orchestrator calls this at the start of batch execution.
+   */
+  registerBatch(batchId: string, taskIds: string[]): void {
+    for (const taskId of taskIds) {
+      this.taskIdToBatchId.set(taskId, batchId);
+    }
+    this.pendingBatchInputs.set(batchId, new Map());
+    this.pendingBatchOutputs.set(batchId, new Map());
+  }
+
   /**
    * Notify subscribers that hash inputs are available for a task.
    * Called from the hasher when inputs are computed.
+   * If the task is in a registered batch, inputs are buffered until
+   * finalizeBatchInputs is called.
    */
   notifyTaskInputs(
     taskId: string,
@@ -96,6 +120,15 @@ class TaskIOService {
       external: string[];
     }
   ): void {
+    const batchId = this.taskIdToBatchId.get(taskId);
+    if (batchId) {
+      const batchInputs = this.pendingBatchInputs.get(batchId);
+      if (batchInputs) {
+        batchInputs.set(taskId, inputs);
+        return;
+      }
+    }
+
     const taskInputInfo: TaskInputInfo = {
       taskId,
       inputs,
@@ -113,8 +146,19 @@ class TaskIOService {
   /**
    * Notify subscribers that task outputs have been collected.
    * Called from the cache when outputs are stored.
+   * If the task is in a registered batch, outputs are buffered until
+   * finalizeBatchOutputs is called.
    */
   notifyTaskOutputs(taskId: string, outputs: string[]): void {
+    const batchId = this.taskIdToBatchId.get(taskId);
+    if (batchId) {
+      const batchOutputs = this.pendingBatchOutputs.get(batchId);
+      if (batchOutputs) {
+        batchOutputs.set(taskId, outputs);
+        return;
+      }
+    }
+
     const update: TaskOutputsUpdate = {
       taskId,
       outputs,
@@ -142,6 +186,128 @@ class TaskIOService {
       }
     }
   }
+
+  /**
+   * Broadcast a PID update to all tasks in a batch.
+   * Since batch tasks share a single worker process, each task
+   * receives the same PID notification.
+   */
+  notifyBatchPidUpdate(taskIds: string[], pid: number): void {
+    for (const taskId of taskIds) {
+      this.notifyPidUpdate({ taskId, pid });
+    }
+  }
+
+  /**
+   * Merge all buffered inputs for a batch and emit the combined set
+   * to each task in the batch. Called by the orchestrator after all
+   * hashing for the batch is complete.
+   */
+  finalizeBatchInputs(batchId: string): void {
+    const batchInputs = this.pendingBatchInputs.get(batchId);
+    if (!batchInputs || batchInputs.size === 0) {
+      return;
+    }
+
+    const combined = mergeBatchInputs(batchInputs);
+
+    // Emit combined inputs for each task in the batch
+    for (const taskId of batchInputs.keys()) {
+      const taskInputInfo: TaskInputInfo = { taskId, inputs: combined };
+      for (const cb of this.taskInputCallbacks) {
+        try {
+          cb(taskInputInfo);
+        } catch {
+          // Silent failure - don't let one callback break others
+        }
+      }
+    }
+
+    this.pendingBatchInputs.delete(batchId);
+  }
+
+  /**
+   * Merge all buffered outputs for a batch and emit the combined set
+   * to each task in the batch. Called by the orchestrator after all
+   * caching for the batch is complete.
+   */
+  finalizeBatchOutputs(batchId: string): void {
+    const batchOutputs = this.pendingBatchOutputs.get(batchId);
+    if (!batchOutputs || batchOutputs.size === 0) {
+      return;
+    }
+
+    const combined = mergeBatchOutputs(batchOutputs);
+
+    // Emit combined outputs for each task in the batch
+    for (const taskId of batchOutputs.keys()) {
+      const update: TaskOutputsUpdate = { taskId, outputs: combined };
+      for (const cb of this.taskOutputsCallbacks) {
+        try {
+          cb(update);
+        } catch {
+          // Silent failure - don't let one callback break others
+        }
+      }
+    }
+
+    this.pendingBatchOutputs.delete(batchId);
+  }
+
+  /**
+   * Clean up all batch state for a given batch.
+   * Called by the orchestrator when the batch is fully complete.
+   */
+  clearBatch(batchId: string): void {
+    this.pendingBatchInputs.delete(batchId);
+    this.pendingBatchOutputs.delete(batchId);
+    for (const [taskId, id] of this.taskIdToBatchId) {
+      if (id === batchId) {
+        this.taskIdToBatchId.delete(taskId);
+      }
+    }
+  }
+}
+
+/**
+ * Merge inputs from all tasks in a batch into a single combined set.
+ * Uses Set for deduplication across tasks.
+ */
+function mergeBatchInputs(
+  batchInputs: Map<string, TaskInputInfo['inputs']>
+): TaskInputInfo['inputs'] {
+  const files = new Set<string>();
+  const runtime = new Set<string>();
+  const environment = new Set<string>();
+  const depOutputs = new Set<string>();
+  const external = new Set<string>();
+
+  for (const inputs of batchInputs.values()) {
+    for (const f of inputs.files) files.add(f);
+    for (const r of inputs.runtime) runtime.add(r);
+    for (const e of inputs.environment) environment.add(e);
+    for (const d of inputs.depOutputs) depOutputs.add(d);
+    for (const x of inputs.external) external.add(x);
+  }
+
+  return {
+    files: [...files],
+    runtime: [...runtime],
+    environment: [...environment],
+    depOutputs: [...depOutputs],
+    external: [...external],
+  };
+}
+
+/**
+ * Merge outputs from all tasks in a batch into a single deduplicated list.
+ */
+function mergeBatchOutputs(batchOutputs: Map<string, string[]>): string[] {
+  const combined = new Set<string>();
+  for (const outputs of batchOutputs.values()) {
+    for (const o of outputs) combined.add(o);
+  }
+  return [...combined];
 }
 
 // Singleton
@@ -167,3 +333,17 @@ export function registerTaskProcessStart(taskId: string, pid: number): void {
   getTaskIOService().notifyPidUpdate({ taskId, pid });
   getProcessMetricsService().registerTaskProcess(taskId, pid);
 }
+
+/**
+ * Register a batch process start with both IO and metrics services.
+ * Broadcasts the PID to all tasks in the batch since they share a single
+ * worker process and cannot be differentiated.
+ */
+export function registerBatchProcessStart(
+  batchId: string,
+  taskIds: string[],
+  pid: number
+): void {
+  getTaskIOService().notifyBatchPidUpdate(taskIds, pid);
+  getProcessMetricsService().registerBatch(batchId, taskIds, pid);
+}

packages/nx/src/tasks-runner/task-orchestrator.ts

@@ -42,6 +42,7 @@ import {
   getEnvVariablesForTask,
   getTaskSpecificEnv,
 } from './task-env';
+import { getTaskIOService } from './task-io-service';
 import { TaskStatus } from './tasks-runner';
 import { Batch, TasksSchedule } from './tasks-schedule';
 import {
@@ -470,6 +471,11 @@ export class TaskOrchestrator {
       taskIds: Object.keys(batch.taskGraph.tasks),
     });
 
+    // Register batch with TaskIOService so inputs/outputs are buffered
+    const ioService = getTaskIOService();
+    const allBatchTaskIds = Object.keys(batch.taskGraph.tasks);
+    ioService.registerBatch(batch.id, allBatchTaskIds);
+
     const { cachedResults, needsRehashAfterExecution } =
       await this.applyBatchCachedResults(batch, doNotSkipCache, groupId);
 
@@ -514,10 +520,16 @@ export class TaskOrchestrator {
       }
     }
 
+    // Emit combined inputs for all tasks in the batch now that all hashing is done
+    ioService.finalizeBatchInputs(batch.id);
+
     if (batchResults.length > 0) {
       await this.postRunSteps(batchResults, doNotSkipCache, { groupId });
     }
 
+    // Emit combined outputs for all tasks in the batch now that caching is done
+    ioService.finalizeBatchOutputs(batch.id);
+
     // Update batch status based on all task results
     const hasFailures = taskEntries.some(([taskId]) => {
       const status = this.completedTasks[taskId];
@@ -549,6 +561,9 @@ export class TaskOrchestrator {
         groupId
       );
     }
+    // Clean up batch IO state
+    ioService.clearBatch(batch.id);
+
     // Batch is done, mark it as completed
     const applyFromCacheOrRunBatchEnd = performance.mark(
       'TaskOrchestrator-apply-from-cache-or-run-batch:end'