@in-the-loop-labs/pair-review 3.2.3 → 3.3.1

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

@@ -170,6 +170,97 @@ class GitWorktreeManager {
     return this.resolveRemoteForRepo(git, cloneUrl, sshUrl);
   }
 
+  /**
+   * Extract a PR number from either { number } or { prNumber } shapes.
+   * @param {Object|null} prInfo
+   * @returns {number|null}
+   */
+  getPRNumber(prInfo) {
+    if (!prInfo) return null;
+    return prInfo.number || prInfo.prNumber || null;
+  }
+
+  /**
+   * Extract the PR head branch name from either REST or stored PR metadata.
+   * @param {Object|null} prData
+   * @returns {string}
+   */
+  getPRHeadBranch(prData) {
+    return prData?.head?.ref || prData?.head_branch || '';
+  }
+
+  /**
+   * Extract the PR head SHA from either REST or stored PR metadata.
+   * @param {Object|null} prData
+   * @returns {string}
+   */
+  getPRHeadSha(prData) {
+    return prData?.head?.sha || prData?.head_sha || '';
+  }
+
+  /**
+   * Detect whether a fetch failed because the remote does not expose a PR ref.
+   * @param {Error} error
+   * @returns {boolean}
+   */
+  isMissingRemoteRefError(error) {
+    const message = String(error?.message || '').toLowerCase();
+    return message.includes('couldn\'t find remote ref') ||
+      message.includes('could not find remote ref') ||
+      message.includes('remote ref does not exist') ||
+      message.includes('fatal: invalid refspec');
+  }
+
+  /**
+   * Fetch a PR head into a stable tracking ref, falling back from GitHub PR
+   * refs to a direct SHA fetch when the git transport does not expose
+   * refs/pull/* (for example, alternate internal fetch backends).
+   *
+   * @param {Object} git - simple-git instance
+   * @param {Object} prInfo - PR info { owner, repo, number } or { prNumber }
+   * @param {Object} prData - PR data from GitHub API or stored metadata
+   * @param {Object} [options={}]
+   * @param {string|null} [options.remote] - Base repository remote to use for PR-ref fetch
+   * @returns {Promise<{remote: string, trackingRef: string|null, checkoutTarget: string}>}
+   */
+  async fetchPRHead(git, prInfo, prData, options = {}) {
+    const prNumber = this.getPRNumber(prInfo);
+    const headSha = this.getPRHeadSha(prData);
+    const baseRemote = options.remote || await this.resolveRemoteForPR(git, prData, prInfo);
+
+    if (!prNumber) {
+      throw new Error('Cannot fetch PR head without a PR number');
+    }
+
+    const prTrackingRef = `refs/remotes/${baseRemote}/pr-${prNumber}`;
+
+    try {
+      await git.fetch([baseRemote, `+refs/pull/${prNumber}/head:${prTrackingRef}`]);
+      return {
+        remote: baseRemote,
+        trackingRef: prTrackingRef,
+        checkoutTarget: prTrackingRef
+      };
+    } catch (prRefError) {
+      if (!this.isMissingRemoteRefError(prRefError)) {
+        throw prRefError;
+      }
+
+      console.warn(`PR ref fetch unavailable for PR #${prNumber} on remote ${baseRemote}, falling back to SHA fetch: ${prRefError.message}`);
+
+      if (!headSha) {
+        throw prRefError;
+      }
+
+      await git.raw(['fetch', baseRemote, headSha]);
+      return {
+        remote: baseRemote,
+        trackingRef: null,
+        checkoutTarget: headSha
+      };
+    }
+  }
+
   /**
    * Execute a user-provided checkout script in the worktree.
    * The script receives PR context as environment variables and is responsible
@@ -241,10 +332,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 +358,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 +400,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')) {
@@ -404,20 +518,21 @@ class GitWorktreeManager {
         console.log(`Base SHA fetch not needed or already available: ${fetchError.message}`);
       }
 
-      // Fetch the PR head using GitHub's pull request refs (more reliable than branch names)
+      // Fetch the PR head using PR refs when available, with a branch/SHA fallback
       console.log(`Fetching PR #${prInfo.number} head...`);
-      await worktreeGit.fetch([remote, `+refs/pull/${prInfo.number}/head:refs/remotes/${remote}/pr-${prInfo.number}`]);
+      const fetchedHead = await this.fetchPRHead(worktreeGit, prInfo, prData, { remote });
 
       // Execute checkout script if configured (before checkout so sparse-checkout is set up)
       if (checkoutScript) {
         // Fetch the actual head branch by name (for checkout scripts that expect branch refs)
         // This may fail for fork PRs where the branch is in a different repo - that's okay
-        if (prData.head_branch) {
+        const headBranch = this.getPRHeadBranch(prData);
+        if (headBranch) {
           try {
-            console.log(`Fetching head branch ${prData.head_branch}...`);
-            await worktreeGit.fetch([remote, `+refs/heads/${prData.head_branch}:refs/remotes/${remote}/${prData.head_branch}`]);
+            console.log(`Fetching head branch ${headBranch}...`);
+            await worktreeGit.fetch([remote, `+refs/heads/${headBranch}:refs/remotes/${remote}/${headBranch}`]);
             // Create/update a local branch pointing to the fetched ref so tooling can reference it by name
-            await worktreeGit.branch(['-f', prData.head_branch, `${remote}/${prData.head_branch}`]);
+            await worktreeGit.branch(['-f', headBranch, `${remote}/${headBranch}`]);
           } catch (branchFetchError) {
             // Expected for fork PRs - the branch exists in the fork, not the base repo
             console.log(`Could not fetch head branch (may be from a fork): ${branchFetchError.message}`);
@@ -427,9 +542,9 @@ class GitWorktreeManager {
         console.log(`Executing checkout script: ${checkoutScript}`);
         const scriptEnv = {
           BASE_BRANCH: prData.base_branch,
-          HEAD_BRANCH: prData.head_branch,
+          HEAD_BRANCH: headBranch,
           BASE_SHA: prData.base_sha,
-          HEAD_SHA: prData.head_sha,
+          HEAD_SHA: this.getPRHeadSha(prData),
           PR_NUMBER: String(prInfo.number),
           WORKTREE_PATH: worktreePath
         };
@@ -438,8 +553,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 = this.getPRHeadSha(prData);
+      if (targetSha) {
+        console.log(`Checking out to PR head commit ${targetSha}...`);
+        await worktreeGit.checkout([targetSha]);
+      } else {
+        console.log(`Checking out to PR head ref ${fetchedHead.checkoutTarget}...`);
+        await worktreeGit.checkout([fetchedHead.checkoutTarget]);
+      }
       
       // Verify we're at the correct commit
       const currentCommit = await worktreeGit.revparse(['HEAD']);
@@ -448,18 +569,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,17 +627,18 @@ 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
+      // Fetch the PR head using PR refs when available, with a branch/SHA fallback
       console.log(`Fetching PR #${number} head...`);
-      await worktreeGit.fetch([remote, `+refs/pull/${number}/head:refs/remotes/${remote}/pr-${number}`]);
+      const fetchedHead = await this.fetchPRHead(worktreeGit, prInfo, prData, { remote });
 
       // Checkout to PR head commit
       console.log(`Checking out to PR head commit ${headSha}...`);
-      await worktreeGit.checkout([`${remote}/pr-${number}`]);
+      await worktreeGit.checkout([fetchedHead.checkoutTarget]);
 
       // Verify we're at the correct commit
       const currentCommit = await worktreeGit.revparse(['HEAD']);
@@ -886,11 +1011,11 @@ class GitWorktreeManager {
 
       // Fetch the latest PR head from remote
       console.log(`Fetching PR #${prNumber} head from ${remote}...`);
-      await git.fetch([remote, `pull/${prNumber}/head`]);
+      const fetchedHead = await this.fetchPRHead(git, prInfo || { number: prNumber }, prData, { remote });
 
       // Reset to the fetched PR head
       console.log(`Resetting worktree to PR head...`);
-      await git.raw(['reset', '--hard', 'FETCH_HEAD']);
+      await git.raw(['reset', '--hard', fetchedHead.checkoutTarget]);
 
       // Update last_accessed_at in database
       if (this.worktreeRepo) {
@@ -944,13 +1069,13 @@ class GitWorktreeManager {
         ? await this.resolveRemoteForPR(git, prData, prInfo)
         : defaultRemote;
 
-      // 3. Fetch PR head into a persistent ref
-      console.log(`Fetching PR #${prNumber} head from ${remote} into refs/remotes/${remote}/pr-${prNumber}...`);
-      await git.fetch([remote, `+refs/pull/${prNumber}/head:refs/remotes/${remote}/pr-${prNumber}`]);
+      // 3. Fetch PR head into a persistent ref (or by SHA when refs are unavailable)
+      console.log(`Fetching PR #${prNumber} head from ${remote}...`);
+      const fetchedHead = await this.fetchPRHead(git, prInfo || { number: prNumber }, prData, { remote });
 
       // 4. Reset worktree to the fetched ref
-      console.log(`Resetting worktree to refs/remotes/${remote}/pr-${prNumber}...`);
-      await git.raw(['reset', '--hard', `refs/remotes/${remote}/pr-${prNumber}`]);
+      console.log(`Resetting worktree to ${fetchedHead.checkoutTarget}...`);
+      await git.raw(['reset', '--hard', fetchedHead.checkoutTarget]);
 
       // 5. Return the new HEAD SHA
       const headSha = (await git.revparse(['HEAD'])).trim();
@@ -1186,4 +1311,4 @@ class GitWorktreeManager {
   }
 }
 
-module.exports = { GitWorktreeManager };
+module.exports = { GitWorktreeManager };