fix(core): ensure workers shutdown after phase cancelled (#34799)

AgentEnder · web-flow · commit 631734028fd8 · 2026-03-17T13:53:56.000-04:00
## Current Behavior
New logic in the daemon can cancel an active graph creation, resulting
in worker shutdown not behaving as intended

## Expected Behavior
When the daemon aborts graph construction, the plugins still shut down

## Related Issue(s)
&lt;!-- Please link the issue being fixed so it gets closed when this is
merged. --&gt;

Fixes #
diff --git a/packages/nx/src/daemon/server/project-graph-incremental-recomputation.ts b/packages/nx/src/daemon/server/project-graph-incremental-recomputation.ts
@@ -357,6 +357,7 @@ async function processFilesAndCreateAndSerializeProjectGraph(
 
     // Early exit if a newer recomputation has started - chain to the newer one
     if (isStale()) {
+      notifyPluginsGraphAborted(plugins);
       return cachedSerializedProjectGraphPromise;
     }
 
@@ -381,6 +382,7 @@ async function processFilesAndCreateAndSerializeProjectGraph(
 
     // Early exit if a newer recomputation has started - chain to the newer one
     if (isStale()) {
+      notifyPluginsGraphAborted(plugins);
       return cachedSerializedProjectGraphPromise;
     }
 
@@ -549,6 +551,15 @@ async function resetInternalStateIfNxDepsMissing() {
   }
 }
 
+function notifyPluginsGraphAborted(plugins: LoadedNxPlugin[]) {
+  // At both abort sites, only createNodes has been called.
+  // createDependencies and createMetadata are called later in
+  // createAndSerializeProjectGraph, which hasn't run yet.
+  for (const plugin of plugins) {
+    plugin.abortGraphPhase?.('createNodes');
+  }
+}
+
 function notifyProjectGraphRecomputationListeners(
   projectGraph: ProjectGraph,
   sourceMaps: ConfigurationSourceMaps,
diff --git a/packages/nx/src/project-graph/plugins/isolation/isolated-plugin.ts b/packages/nx/src/project-graph/plugins/isolation/isolated-plugin.ts
@@ -434,6 +434,12 @@ export class IsolatedPlugin implements LoadedNxPlugin {
     this.shutdown();
   }
 
+  abortGraphPhase(lastCompletedHook: Hook): void {
+    if (this.lifecycle?.abortPhase('graph', lastCompletedHook)) {
+      this.shutdownIfInactive();
+    }
+  }
+
   shutdown(): void {
     if (!this._alive) return;
     this._alive = false;
diff --git a/packages/nx/src/project-graph/plugins/isolation/plugin-lifecycle-manager.spec.ts b/packages/nx/src/project-graph/plugins/isolation/plugin-lifecycle-manager.spec.ts
@@ -349,6 +349,199 @@ describe('PluginLifecycleManager', () => {
     });
   });
 
+  describe('abortPhase', () => {
+    it('should decrement session count when last completed hook is not the last registered hook', () => {
+      const lifecycle = new PluginLifecycleManager([
+        'createNodes',
+        'createDependencies',
+        'createMetadata',
+      ]);
+
+      // Simulate entering graph phase (createNodes is first hook)
+      lifecycle.enterHook('createNodes');
+      lifecycle.exitHook('createNodes'); // not last hook, no decrement
+      expect(lifecycle.getPhaseRefCount('graph')).toBe(1);
+
+      // Abort after createNodes — createDependencies and createMetadata remain
+      const shouldShutdown = lifecycle.abortPhase('graph', 'createNodes');
+      expect(lifecycle.getPhaseRefCount('graph')).toBe(0);
+      expect(shouldShutdown).toBe(true); // graph is the only registered phase
+    });
+
+    it('should be a no-op when last completed hook IS the last registered hook', () => {
+      const lifecycle = new PluginLifecycleManager(['createNodes']);
+
+      // Single hook: createNodes is both first and last.
+      // exitHook already closed the session.
+      lifecycle.enterHook('createNodes');
+      lifecycle.exitHook('createNodes');
+      expect(lifecycle.getPhaseRefCount('graph')).toBe(0);
+
+      // abortPhase with createNodes as last completed — it's the last registered hook
+      const shouldShutdown = lifecycle.abortPhase('graph', 'createNodes');
+      expect(shouldShutdown).toBe(false);
+      expect(lifecycle.getPhaseRefCount('graph')).toBe(0);
+    });
+
+    it('should be a no-op when aborting after completing all hooks in a multi-hook phase', () => {
+      const lifecycle = new PluginLifecycleManager([
+        'createNodes',
+        'createDependencies',
+      ]);
+
+      // Complete the full phase
+      lifecycle.enterHook('createNodes');
+      lifecycle.exitHook('createNodes');
+      lifecycle.enterHook('createDependencies');
+      lifecycle.exitHook('createDependencies'); // last hook, decrements
+      expect(lifecycle.getPhaseRefCount('graph')).toBe(0);
+
+      // Abort after createDependencies — it's the last registered hook, session closed
+      const shouldShutdown = lifecycle.abortPhase(
+        'graph',
+        'createDependencies'
+      );
+      expect(shouldShutdown).toBe(false);
+      expect(lifecycle.getPhaseRefCount('graph')).toBe(0);
+    });
+
+    it('should not steal another callers session in single-hook phase', () => {
+      const lifecycle = new PluginLifecycleManager(['createNodes']);
+
+      // Caller A completes createNodes (session fully closed)
+      lifecycle.enterHook('createNodes');
+      lifecycle.exitHook('createNodes');
+
+      // Caller B enters createNodes (new session)
+      lifecycle.enterHook('createNodes');
+      expect(lifecycle.getPhaseRefCount('graph')).toBe(1);
+
+      // Caller A aborts — createNodes is the last registered hook, no-op
+      const shouldShutdown = lifecycle.abortPhase('graph', 'createNodes');
+      expect(shouldShutdown).toBe(false);
+      expect(lifecycle.getPhaseRefCount('graph')).toBe(1);
+
+      // Caller B completes normally
+      expect(lifecycle.exitHook('createNodes')).toBe(true);
+      expect(lifecycle.getPhaseRefCount('graph')).toBe(0);
+    });
+
+    it('should be a no-op when hook is not registered for the plugin', () => {
+      // Plugin only has createDependencies, not createNodes
+      const lifecycle = new PluginLifecycleManager(['createDependencies']);
+
+      // Abort with createNodes — not registered, so no session was opened
+      const shouldShutdown = lifecycle.abortPhase('graph', 'createNodes');
+      expect(shouldShutdown).toBe(false);
+    });
+
+    it('should return false when phase has no active sessions', () => {
+      const lifecycle = new PluginLifecycleManager([
+        'createNodes',
+        'createDependencies',
+      ]);
+
+      const shouldShutdown = lifecycle.abortPhase('graph', 'createNodes');
+      expect(shouldShutdown).toBe(false);
+      expect(lifecycle.getPhaseRefCount('graph')).toBe(0);
+    });
+
+    it('should return false for unregistered phase', () => {
+      const lifecycle = new PluginLifecycleManager(['createNodes']);
+
+      const shouldShutdown = lifecycle.abortPhase(
+        'pre-task',
+        'preTasksExecution'
+      );
+      expect(shouldShutdown).toBe(false);
+    });
+
+    it('should return false when later phases have hooks', () => {
+      const lifecycle = new PluginLifecycleManager([
+        'createNodes',
+        'createDependencies',
+        'postTasksExecution',
+      ]);
+
+      lifecycle.enterHook('createNodes');
+      lifecycle.exitHook('createNodes');
+      const shouldShutdown = lifecycle.abortPhase('graph', 'createNodes');
+      expect(lifecycle.getPhaseRefCount('graph')).toBe(0);
+      expect(shouldShutdown).toBe(false); // post-task phase still registered
+    });
+
+    it('should decrement by 1 with concurrent callers', () => {
+      const lifecycle = new PluginLifecycleManager([
+        'createNodes',
+        'createDependencies',
+        'createMetadata',
+      ]);
+
+      // Two concurrent callers enter and exit createNodes (first hook, not last)
+      lifecycle.enterHook('createNodes');
+      lifecycle.exitHook('createNodes');
+      lifecycle.enterHook('createNodes');
+      lifecycle.exitHook('createNodes');
+      expect(lifecycle.getPhaseRefCount('graph')).toBe(2);
+
+      // First caller aborts after createNodes
+      expect(lifecycle.abortPhase('graph', 'createNodes')).toBe(false);
+      expect(lifecycle.getPhaseRefCount('graph')).toBe(1);
+
+      // Second caller aborts after createNodes
+      expect(lifecycle.abortPhase('graph', 'createNodes')).toBe(true);
+      expect(lifecycle.getPhaseRefCount('graph')).toBe(0);
+    });
+
+    it('should handle abort after second of three hooks', () => {
+      const lifecycle = new PluginLifecycleManager([
+        'createNodes',
+        'createDependencies',
+        'createMetadata',
+      ]);
+
+      // Caller enters createNodes and createDependencies but not createMetadata
+      lifecycle.enterHook('createNodes');
+      lifecycle.exitHook('createNodes');
+      lifecycle.enterHook('createDependencies');
+      lifecycle.exitHook('createDependencies');
+      expect(lifecycle.getPhaseRefCount('graph')).toBe(1); // session still open
+
+      // Abort after createDependencies — createMetadata remains
+      const shouldShutdown = lifecycle.abortPhase(
+        'graph',
+        'createDependencies'
+      );
+      expect(lifecycle.getPhaseRefCount('graph')).toBe(0);
+      expect(shouldShutdown).toBe(true);
+    });
+
+    it('should allow normal hook flow after abort', () => {
+      const lifecycle = new PluginLifecycleManager([
+        'createNodes',
+        'createDependencies',
+        'createMetadata',
+      ]);
+
+      // First recomputation enters but gets aborted after createNodes
+      lifecycle.enterHook('createNodes');
+      lifecycle.exitHook('createNodes');
+      lifecycle.abortPhase('graph', 'createNodes');
+      expect(lifecycle.getPhaseRefCount('graph')).toBe(0);
+
+      // Second recomputation runs to completion
+      lifecycle.enterHook('createNodes');
+      expect(lifecycle.getPhaseRefCount('graph')).toBe(1);
+      lifecycle.exitHook('createNodes');
+      lifecycle.enterHook('createDependencies');
+      lifecycle.exitHook('createDependencies');
+      lifecycle.enterHook('createMetadata');
+      const shouldShutdown = lifecycle.exitHook('createMetadata');
+      expect(shouldShutdown).toBe(true);
+      expect(lifecycle.getPhaseRefCount('graph')).toBe(0);
+    });
+  });
+
   describe('wrapHook', () => {
     it('should call enterHook before and exitHook after the wrapped function', async () => {
       const lifecycle = new PluginLifecycleManager(['createNodes']);
diff --git a/packages/nx/src/project-graph/plugins/isolation/plugin-lifecycle-manager.ts b/packages/nx/src/project-graph/plugins/isolation/plugin-lifecycle-manager.ts
@@ -174,6 +174,52 @@ export class PluginLifecycleManager {
     return phaseIndex === this.registeredPhaseOrder.length - 1;
   }
 
+  /**
+   * Aborts a single session within a phase by decrementing its count,
+   * but only if the caller's session is still open.
+   *
+   * Called when graph construction is abandoned mid-flight (e.g., a newer
+   * recomputation supersedes the current one). Without this, hooks that
+   * were entered but whose phase never completed would leave the session
+   * count elevated, preventing the worker from ever shutting down.
+   *
+   * @param phase The phase to abort
+   * @param lastCompletedHook The last hook the caller executed before aborting.
+   *   Used to determine whether the caller's session is still open: if there
+   *   are registered hooks after this one, the session is open and needs
+   *   cleanup. If this IS the last registered hook, exitHook already closed
+   *   the session — decrementing again would steal another caller's count.
+   * @returns true if the worker should shut down, false otherwise
+   */
+  abortPhase(phase: Phase, lastCompletedHook: Hook): boolean {
+    const phaseHooks = this.registeredPhases[phase];
+    if (!phaseHooks) {
+      return false;
+    }
+
+    const hookIndex = phaseHooks.indexOf(lastCompletedHook);
+
+    // Hook not registered → this caller never entered the phase → no session to abort.
+    // Hook is the last registered hook → exitHook already closed the session.
+    if (hookIndex === -1 || hookIndex === phaseHooks.length - 1) {
+      return false;
+    }
+
+    const count = this.phaseSessionCount[phase] ?? 0;
+    if (count === 0) {
+      return false;
+    }
+
+    const newCount = count - 1;
+    this.phaseSessionCount[phase] = newCount;
+
+    if (newCount > 0) {
+      return false;
+    }
+
+    return this.isLastRegisteredPhase(phase);
+  }
+
   /**
    * Returns the current session count for a phase. Useful for testing.
    */
diff --git a/packages/nx/src/project-graph/plugins/loaded-nx-plugin.ts b/packages/nx/src/project-graph/plugins/loaded-nx-plugin.ts
@@ -55,6 +55,16 @@ export class LoadedNxPlugin {
   readonly include?: string[];
   readonly exclude?: string[];
 
+  /**
+   * Notifies the plugin that graph construction was aborted mid-flight.
+   * Overridden by IsolatedPlugin to reset lifecycle phase tracking so
+   * the worker can still shut down properly.
+   *
+   * @param lastCompletedHook The name of the last graph-phase hook that
+   *   was called before the abort (e.g. 'createNodes').
+   */
+  abortGraphPhase?(lastCompletedHook: string): void;
+
   constructor(
     plugin: NxPluginV2,
     pluginDefinition: PluginConfiguration,
