@in-the-loop-labs/pair-review 3.2.2 → 3.3.0

package/src/git/worktree-pool-usage.js

@@ -0,0 +1,216 @@
+// Copyright 2026 Tim Perkins (tjwp) | SPDX-License-Identifier: Apache-2.0
+'use strict';
+
+const logger = require('../utils/logger');
+
+const GRACE_PERIOD_MS = 30_000; // 30 seconds after last WS disconnect
+
+/**
+ * In-memory tracker that determines whether a pool worktree is "in use".
+ *
+ * A pool worktree is considered in-use when any of:
+ * - At least one WebSocket session is subscribed to its review topic
+ * - An AI analysis is running against it
+ * - The grace period after the last session disconnect hasn't expired
+ *
+ * When a worktree becomes idle (all of the above are false), the `onIdle`
+ * callback fires so the pool manager can mark it available.
+ */
+class WorktreePoolUsageTracker {
+  constructor() {
+    /** @type {Map<string, Set<string>>} worktreeId -> Set of active session keys */
+    this._sessions = new Map();
+    /** @type {Map<string, Set<string>>} worktreeId -> Set of active analysis IDs */
+    this._analyses = new Map();
+    /** @type {Map<string, NodeJS.Timeout>} worktreeId -> grace period timer */
+    this._graceTimers = new Map();
+    /** @type {Function|null} Callback when a worktree becomes idle: (worktreeId) => void */
+    this.onIdle = null;
+  }
+
+  /**
+   * Register an active WebSocket session for a worktree.
+   * Clears any pending grace-period timer.
+   * @param {string} worktreeId - Pool worktree ID
+   * @param {string} sessionKey - Unique key for this WS connection
+   */
+  addSession(worktreeId, sessionKey) {
+    if (!this._sessions.has(worktreeId)) {
+      this._sessions.set(worktreeId, new Set());
+    }
+    this._sessions.get(worktreeId).add(sessionKey);
+
+    // Cancel any pending grace timer
+    const timer = this._graceTimers.get(worktreeId);
+    if (timer) {
+      clearTimeout(timer);
+      this._graceTimers.delete(worktreeId);
+      logger.debug(`Grace period cancelled for pool worktree ${worktreeId} — new session connected`);
+    }
+  }
+
+  /**
+   * Remove a WebSocket session. Starts grace period if no sessions remain
+   * and no analyses are running.
+   * @param {string} worktreeId
+   * @param {string} sessionKey
+   */
+  removeSession(worktreeId, sessionKey) {
+    const sessions = this._sessions.get(worktreeId);
+    if (!sessions) return;
+
+    sessions.delete(sessionKey);
+    if (sessions.size === 0) {
+      this._sessions.delete(worktreeId);
+    }
+
+    this._checkIdle(worktreeId);
+  }
+
+  /**
+   * Register an active analysis for a worktree.
+   * Clears any pending grace-period timer.
+   * @param {string} worktreeId - Pool worktree ID
+   * @param {string} analysisId - Unique analysis identifier
+   */
+  addAnalysis(worktreeId, analysisId) {
+    if (!this._analyses.has(worktreeId)) {
+      this._analyses.set(worktreeId, new Set());
+    }
+    this._analyses.get(worktreeId).add(analysisId);
+
+    // Cancel any pending grace timer
+    const timer = this._graceTimers.get(worktreeId);
+    if (timer) {
+      clearTimeout(timer);
+      this._graceTimers.delete(worktreeId);
+      logger.debug(`Grace period cancelled for pool worktree ${worktreeId} — new analysis started`);
+    }
+  }
+
+  /**
+   * Remove an active analysis hold.
+   * @param {string} worktreeId
+   * @param {string} analysisId
+   */
+  removeAnalysis(worktreeId, analysisId) {
+    const analyses = this._analyses.get(worktreeId);
+    if (!analyses) return;
+
+    analyses.delete(analysisId);
+    if (analyses.size === 0) {
+      this._analyses.delete(worktreeId);
+    }
+
+    this._checkIdle(worktreeId);
+  }
+
+  /**
+   * Remove an analysis hold by analysisId only (without knowing worktreeId).
+   * Searches all worktrees for the analysis.
+   * @param {string} analysisId
+   */
+  removeAnalysisById(analysisId) {
+    for (const [worktreeId, analyses] of this._analyses) {
+      if (analyses.has(analysisId)) {
+        this.removeAnalysis(worktreeId, analysisId);
+        return;
+      }
+    }
+  }
+
+  /**
+   * Check if a worktree is currently in use.
+   * @param {string} worktreeId
+   * @returns {boolean}
+   */
+  isInUse(worktreeId) {
+    const hasSessions = (this._sessions.get(worktreeId)?.size || 0) > 0;
+    const hasAnalyses = (this._analyses.get(worktreeId)?.size || 0) > 0;
+    const hasGraceTimer = this._graceTimers.has(worktreeId);
+    return hasSessions || hasAnalyses || hasGraceTimer;
+  }
+
+  /**
+   * Return the set of active analysis IDs for a worktree (may be empty).
+   * @param {string} worktreeId
+   * @returns {Set<string>}
+   */
+  getActiveAnalyses(worktreeId) {
+    return new Set(this._analyses.get(worktreeId) || []);
+  }
+
+  /**
+   * Forcefully clear all tracking state for a single worktree.
+   *
+   * Removes sessions, analyses, and grace timers.
+   * Does NOT fire the onIdle callback — the caller is assumed to be
+   * handling the worktree lifecycle directly (e.g., deleting or
+   * marking it available in the pool).
+   *
+   * @param {string} worktreeId - Pool worktree ID to purge
+   */
+  clearWorktree(worktreeId) {
+    this._sessions.delete(worktreeId);
+    this._analyses.delete(worktreeId);
+    const timer = this._graceTimers.get(worktreeId);
+    if (timer) {
+      clearTimeout(timer);
+      this._graceTimers.delete(worktreeId);
+    }
+    logger.debug(`Cleared all tracking state for pool worktree ${worktreeId}`);
+  }
+
+  /**
+   * Internal: check if a worktree has become idle and start grace period or fire callback.
+   * @param {string} worktreeId
+   * @private
+   */
+  _checkIdle(worktreeId) {
+    const hasSessions = (this._sessions.get(worktreeId)?.size || 0) > 0;
+    const hasAnalyses = (this._analyses.get(worktreeId)?.size || 0) > 0;
+
+    if (hasSessions || hasAnalyses) return; // still in use
+
+    // Already have a grace timer? Let it run
+    if (this._graceTimers.has(worktreeId)) return;
+
+    // Start grace period
+    logger.debug(`Starting ${GRACE_PERIOD_MS / 1000}s grace period for pool worktree ${worktreeId}`);
+    const timer = setTimeout(async () => {
+      this._graceTimers.delete(worktreeId);
+      // Double-check still idle (a new session could have connected during grace)
+      const stillHasSessions = (this._sessions.get(worktreeId)?.size || 0) > 0;
+      const stillHasAnalyses = (this._analyses.get(worktreeId)?.size || 0) > 0;
+      if (!stillHasSessions && !stillHasAnalyses) {
+        logger.info(`Pool worktree ${worktreeId} idle after grace period`);
+        if (this.onIdle) {
+          try {
+            await this.onIdle(worktreeId);
+          } catch (err) {
+            logger.error(`onIdle callback failed for ${worktreeId}: ${err.message}`);
+          }
+        }
+      }
+    }, GRACE_PERIOD_MS);
+
+    // Don't hold the process open for grace timers
+    if (timer.unref) timer.unref();
+    this._graceTimers.set(worktreeId, timer);
+  }
+
+  /**
+   * Clear all tracking state. Useful for testing.
+   */
+  reset() {
+    this._sessions.clear();
+    this._analyses.clear();
+    for (const timer of this._graceTimers.values()) {
+      clearTimeout(timer);
+    }
+    this._graceTimers.clear();
+    this.onIdle = null;
+  }
+}
+
+module.exports = { WorktreePoolUsageTracker, GRACE_PERIOD_MS };

package/src/git/worktree.js

@@ -241,10 +241,12 @@ class GitWorktreeManager {
    *   When set, worktree is created with --no-checkout from the main git root (no sparse-checkout inheritance),
    *   and the script is executed before checkout with PR context as environment variables.
    * @param {number} [options.checkoutTimeout] - Timeout in ms for checkout script (default: 300000 = 5 minutes)
-   * @returns {Promise<string>} Path to created worktree
+   * @param {string} [options.explicitId] - When provided, use this ID for the worktrees-table record
+   *   instead of generating one. Used by the worktree pool to align the worktrees-table ID with the pool ID.
+   * @returns {Promise<{ path: string, id: string }>} Path and database ID of created worktree
    */
   async createWorktreeForPR(prInfo, prData, repositoryPath, options = {}) {
-    const { worktreeSourcePath, checkoutScript, checkoutTimeout } = options;
+    const { worktreeSourcePath, checkoutScript, checkoutTimeout, explicitId } = options;
     // Check if worktree already exists in DB
     const repository = normalizeRepository(prInfo.owner, prInfo.repo);
     let worktreePath;
@@ -265,7 +267,25 @@ class GitWorktreeManager {
         // Try to reuse existing worktree by refreshing it
         console.log(`Found existing worktree for PR #${prInfo.number} at ${worktreePath}`);
         try {
-          return await this.refreshWorktree(worktreeRecord, prInfo.number, prData, prInfo);
+          const refreshedPath = await this.refreshWorktree(worktreeRecord, prInfo.number, prData, prInfo);
+          let returnId = worktreeRecord.id;
+
+          // If explicitId is provided and differs from the existing record's ID,
+          // migrate the worktrees-table row to use the pool ID. This happens when
+          // pool mode is enabled for a repo that already has legacy worktree records.
+          if (explicitId && worktreeRecord.id !== explicitId && this.worktreeRepo) {
+            const migrated = await this.worktreeRepo.getOrCreate({
+              prNumber: prInfo.number,
+              repository,
+              branch: prData.head_branch || prData.base_branch,
+              path: refreshedPath,
+              explicitId,
+            });
+            returnId = migrated.id;
+            console.log(`Migrated worktree record from ${worktreeRecord.id} to ${explicitId}`);
+          }
+
+          return { path: refreshedPath, id: returnId };
         } catch (refreshError) {
           // If refresh fails due to uncommitted changes, propagate that error
           if (refreshError.message.includes('uncommitted changes')) {
@@ -289,20 +309,23 @@ class GitWorktreeManager {
       if (legacyExists && await this.isValidGitWorktree(legacyPath)) {
         console.log(`Found legacy worktree for PR #${prInfo.number} at ${legacyPath}, adopting it`);
 
-        // Create DB record for the legacy worktree
+        // Create DB record for the legacy worktree — pass explicitId so the record
+        // is created with the pool ID when pool mode is active
         if (this.worktreeRepo) {
           worktreeRecord = await this.worktreeRepo.getOrCreate({
             prNumber: prInfo.number,
             repository,
             branch: prData.head_branch || prData.base_branch,
-            path: legacyPath
+            path: legacyPath,
+            explicitId,
           });
           console.log(`Created database record for legacy worktree`);
         }
 
         // Try to refresh and reuse the legacy worktree
         try {
-          return await this.refreshWorktree({ path: legacyPath, id: worktreeRecord?.id }, prInfo.number, prData, prInfo);
+          const refreshedPath = await this.refreshWorktree({ path: legacyPath, id: worktreeRecord?.id }, prInfo.number, prData, prInfo);
+          return { path: refreshedPath, id: worktreeRecord?.id };
         } catch (refreshError) {
           // If refresh fails due to uncommitted changes, propagate that error
           if (refreshError.message.includes('uncommitted changes')) {
@@ -438,8 +461,14 @@ class GitWorktreeManager {
       }
 
       // Checkout to PR head commit
-      console.log(`Checking out to PR head commit ${prData.head_sha}...`);
-      await worktreeGit.checkout([`${remote}/pr-${prInfo.number}`]);
+      const targetSha = prData.head_sha;
+      if (targetSha) {
+        console.log(`Checking out to PR head commit ${targetSha}...`);
+        await worktreeGit.checkout([targetSha]);
+      } else {
+        console.log(`Checking out to PR head ref ${remote}/pr-${prInfo.number}...`);
+        await worktreeGit.checkout([`${remote}/pr-${prInfo.number}`]);
+      }
       
       // Verify we're at the correct commit
       const currentCommit = await worktreeGit.revparse(['HEAD']);
@@ -448,18 +477,21 @@ class GitWorktreeManager {
       }
 
       // Store/update worktree record in database
+      let worktreeDbId;
       if (this.worktreeRepo) {
-        await this.worktreeRepo.getOrCreate({
+        const record = await this.worktreeRepo.getOrCreate({
           prNumber: prInfo.number,
           repository,
           branch: prData.head_branch || prData.base_branch,
-          path: worktreePath
+          path: worktreePath,
+          explicitId,
         });
+        worktreeDbId = record.id;
         console.log(`Worktree record stored in database`);
       }
 
       console.log(`Worktree created successfully at ${worktreePath}`);
-      return worktreePath;
+      return { path: worktreePath, id: worktreeDbId };
 
     } catch (error) {
       console.error('Error creating worktree:', error);
@@ -503,9 +535,10 @@ class GitWorktreeManager {
       // Resolve which remote points to the PR's base repository (handles forks)
       const remote = await this.resolveRemoteForPR(worktreeGit, prData, prInfo);
 
-      // Fetch the latest from the resolved remote
+      // Fetch the latest from the resolved remote (--prune removes stale
+      // tracking refs that would otherwise block fetch on ref hierarchy conflicts)
       console.log(`Fetching latest changes from ${remote}...`);
-      await worktreeGit.fetch([remote]);
+      await worktreeGit.fetch([remote, '--prune']);
 
       // Fetch the PR head using GitHub's pull request refs
       console.log(`Fetching PR #${number} head...`);