fix(agent-manager): stabilize PR review comments

packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts

@@ -1417,7 +1417,8 @@ export class AgentManagerProvider implements Disposable {
     // already excludes worktrees in collapsed sections.
     this.syncPollerSkips()
     this.statsPoller.setEnabled(worktrees.length > 0 || this.panel !== undefined)
-    this.prBridge.poller.setEnabled(worktrees.length > 0)
+    // Start PR polling during state hydration so persisted badges get live status.
+    this.prBridge.poller.setEnabled(this.panel !== undefined)
   }
 
   /** Push empty state when the folder is not a git repo or has no folder open. */

packages/kilo-vscode/src/agent-manager/PRStatusPoller.ts

@@ -14,6 +14,7 @@ import {
   parseReviewers,
 } from "./pr/am-pr-utils"
 import type { PRResult, GhThread, GhReviewRequest, GhReview } from "./pr/am-pr-types"
+import { withContext } from "./pr/pr-comment-context"
 
 interface PRStatusPollerOptions {
   getWorktrees: () => Worktree[]
@@ -248,7 +249,7 @@ export class PRStatusPoller {
       const pr = await this.cachedFetchPR(wt.branch, wt.path)
       if (!pr || this.stale(generation)) {
         if (this.stale(generation)) return
-        const hash = `${worktreeId}:none`
+        const hash = `${worktreeId}:${wt.branch}:none`
         if (this.lastHash.get(worktreeId) === hash) return
         this.lastHash.set(worktreeId, hash)
         this.options.onStatus(worktreeId, null)
@@ -482,10 +483,14 @@ export class PRStatusPoller {
     }
   }
 
+  /**
+   * Undefined on failure, never an empty thread list: the panel keeps the
+   * comments it already shows instead of collapsing the section mid-review.
+   */
   private async fetchComments(
     prNumber: number,
     cwd: string,
-  ): Promise<{ total: number; unresolved: number; comments: PRComment[] }> {
+  ): Promise<{ total: number; unresolved: number; comments: PRComment[] } | undefined> {
     try {
       const repo = await this.getRepoInfo(cwd)
       const query = `query($owner: String!, $repo: String!, $number: Int!) {
@@ -533,12 +538,12 @@ export class PRStatusPoller {
       )
       const pr = JSON.parse(stdout)?.data?.repository?.pullRequest
       const threads = pr?.reviewThreads
-      const comments = parseComments((threads?.nodes ?? []) as GhThread[])
+      const comments = await withContext(cwd, parseComments((threads?.nodes ?? []) as GhThread[]))
       const totalCount = threads?.totalCount ?? comments.length
       return { total: totalCount, unresolved: comments.filter((c) => !c.resolved).length, comments }
     } catch (err) {
       this.options.log("Failed to fetch PR comments:", err)
-      return { total: 0, unresolved: 0, comments: [] }
+      return undefined
     }
   }
 }

packages/kilo-vscode/src/agent-manager/pr-status-bridge.ts

@@ -10,7 +10,7 @@ import type { Disposable } from "./host"
 import type { Semaphore } from "./semaphore"
 import { PRStatusPoller } from "./PRStatusPoller"
 import { resolveComment, unresolveComment } from "./pr/PRActions"
-import { ghErrorReason } from "./pr/am-pr-utils"
+import { ghErrorReason, mergePRStatus } from "./pr/am-pr-utils"
 
 interface PRBridgeHost {
   getWorktrees(): Worktree[]
@@ -32,6 +32,8 @@ interface PanelLike {
 export class PRStatusBridge {
   readonly poller: PRStatusPoller
   private readonly cache = new Map<string, AgentManagerOutMessage>()
+  /** Branch each cached PR was found on, so a branch switch still clears it. */
+  private readonly branches = new Map<string, string>()
   private readonly host: PRBridgeHost
   private lastErrorNotified: "gh_missing" | "gh_auth" | "fetch_failed" | undefined
 
@@ -137,11 +139,13 @@ export class PRStatusBridge {
   /** Remove cached status for a deleted worktree. */
   remove(worktreeId: string): void {
     this.cache.delete(worktreeId)
+    this.branches.delete(worktreeId)
   }
 
   reset(): void {
     this.poller.stop()
     this.cache.clear()
+    this.branches.clear()
     this.lastErrorNotified = undefined
   }
 
@@ -160,29 +164,57 @@ function bridgePollerOpts(bridge: PRStatusBridge, host: PRBridgeHost) {
     semaphore: host.semaphore,
     onStatus: (id: string, pr: PRStatus | null, err?: "gh_missing" | "gh_auth" | "fetch_failed") => {
       if (err) {
-        // Don't forward errors to the webview when we have prior PR data
-        // (in-memory cache or persisted prNumber) — that would overwrite
-        // the live badge with pr:null. Only forward when there's truly no
-        // prior data (first poll failed, nothing persisted).
-        if (!bridge["cache"].has(id) && !host.hasPersistedPR(id))
-          host.postToWebview({
-            type: "agentManager.prStatus",
-            worktreeId: id,
-            pr: null,
-            error: err,
-          } as AgentManagerOutMessage)
-        // Always forward auth/missing errors so the webview can show a toast,
-        // regardless of whether prior data exists. Deduplicate per error type
-        // so multiple failing worktrees don't produce multiple toasts.
-        if (err === "gh_auth" || err === "gh_missing") bridge.notifyError(err)
+        reportError(bridge, host, id, err)
         return
       }
-      const msg = { type: "agentManager.prStatus", worktreeId: id, pr, error: err } as AgentManagerOutMessage
-      bridge["cache"].set(id, msg)
-      bridge["lastErrorNotified"] = undefined
-      host.postToWebview(msg)
-      host.updateWorktreePR(id, pr?.number, pr?.url, pr?.state)
+      accept(bridge, host, id, pr)
     },
     log: (...args: unknown[]) => host.log(...args),
   }
 }
+
+function reportError(
+  bridge: PRStatusBridge,
+  host: PRBridgeHost,
+  id: string,
+  err: "gh_missing" | "gh_auth" | "fetch_failed",
+): void {
+  // Don't forward errors to the webview when we have prior PR data
+  // (in-memory cache or persisted prNumber) — that would overwrite
+  // the live badge with pr:null. Only forward when there's truly no
+  // prior data (first poll failed, nothing persisted).
+  if (!bridge["cache"].has(id) && !host.hasPersistedPR(id))
+    host.postToWebview({
+      type: "agentManager.prStatus",
+      worktreeId: id,
+      pr: null,
+      error: err,
+    } as AgentManagerOutMessage)
+  // Always forward auth/missing errors so the webview can show a toast,
+  // regardless of whether prior data exists. Deduplicate per error type
+  // so multiple failing worktrees don't produce multiple toasts.
+  if (err === "gh_auth" || err === "gh_missing") bridge.notifyError(err)
+}
+
+function accept(bridge: PRStatusBridge, host: PRBridgeHost, id: string, pr: PRStatus | null): void {
+  const cached = bridge["cache"].get(id)
+  const prev = cached?.type === "agentManager.prStatus" ? cached.pr : null
+  const branch = host.getWorktrees().find((w: Worktree) => w.id === id)?.branch
+  // `gh` answers "no pull request" for a rate limit, a network blip, or an
+  // unresolvable fork ref exactly as it does for a branch that never had one. A
+  // PR cannot leave a branch, so on the same branch the known PR is kept:
+  // forwarding pr:null would unmount the panel and throw away the comment the
+  // user is reading.
+  if (!pr && prev && branch !== undefined && bridge["branches"].get(id) === branch) {
+    host.log(`PR status: keeping PR #${prev.number} for ${id}, empty result on ${branch}`)
+    return
+  }
+  const merged = pr && prev ? mergePRStatus(prev, pr) : pr
+  const msg = { type: "agentManager.prStatus", worktreeId: id, pr: merged } as AgentManagerOutMessage
+  bridge["cache"].set(id, msg)
+  if (pr && branch !== undefined) bridge["branches"].set(id, branch)
+  if (!pr) bridge["branches"].delete(id)
+  bridge["lastErrorNotified"] = undefined
+  host.postToWebview(msg)
+  host.updateWorktreePR(id, pr?.number, pr?.url, pr?.state)
+}

packages/kilo-vscode/src/agent-manager/pr/am-pr-utils.ts

@@ -1,5 +1,5 @@
 import { createHash } from "node:crypto"
-import type { CheckStatus, PRComment, PRReviewer, ReviewerState } from "../types"
+import type { CheckStatus, PRComment, PRReviewer, PRStatus, ReviewerState } from "../types"
 import type { PRResult, GhThread, GhReviewRequest, GhReview } from "./am-pr-types"
 
 export function parsePRResult(json: string): PRResult | null {
@@ -128,6 +128,17 @@ export function ghErrorReason(message: string): string {
   return (last ?? message.trim()).replace(/^gh:\s*/, "").slice(0, 200)
 }
 
+/**
+ * Carry review threads across a status that has none. Only the selected worktree
+ * fetches comments, and that fetch can fail, so a plain replace would collapse
+ * the open comment list in the panel while the user is reading it.
+ */
+export function mergePRStatus(prev: PRStatus | undefined, next: PRStatus): PRStatus {
+  if (next.comments || !prev?.comments) return next
+  if (prev.number !== next.number) return next
+  return { ...next, comments: prev.comments }
+}
+
 /**
  * Signature of the comment threads, for poll deduplication. Thread and
  * unresolved counts alone hide edits and new replies, which the panel renders.

packages/kilo-vscode/src/agent-manager/pr/pr-comment-context.ts

@@ -0,0 +1,70 @@
+/**
+ * GitHub truncates `diffHunk` at the commented line, so a review comment about
+ * what happens *after* that line has no code to point at. The worktree holds
+ * the file, so the lines below the comment are read from disk and attached to
+ * the thread.
+ */
+import { readFile, realpath, stat } from "node:fs/promises"
+import path from "node:path"
+import type { PRComment } from "../types"
+
+/** Lines of trailing context, matching the window the panel renders. */
+const AFTER = 3
+/** A generated or vendored file is not worth reading for three lines. */
+const SIZE = 2_000_000
+/** Bound the mtime cache so a long session cannot grow it without limit. */
+const CACHE = 200
+
+const cache = new Map<string, { mtime: number; lines: string[] }>()
+
+export function clearContextCache(): void {
+  cache.clear()
+}
+
+/**
+ * The last line of a hunk, without its diff marker. GitHub always ends the
+ * hunk at the commented line, so this is the text the comment refers to.
+ */
+function anchor(hunk: string): string | undefined {
+  const lines = hunk.split("\n").filter((line) => line.length > 0 && !line.startsWith("\\"))
+  const last = lines.at(-1)
+  return lines.length > 1 && last ? last.slice(1) : undefined
+}
+
+async function lines(dir: string, file: string): Promise<string[] | undefined> {
+  const root = await realpath(dir).catch(() => undefined)
+  const full = await realpath(path.resolve(dir, file)).catch(() => undefined)
+  if (!root || !full) return undefined
+  const rel = path.relative(root, full)
+  if (!rel || rel.startsWith("..") || path.isAbsolute(rel)) return undefined
+  const info = await stat(full).catch(() => undefined)
+  if (!info?.isFile() || info.size > SIZE) return undefined
+  const hit = cache.get(full)
+  if (hit && hit.mtime === info.mtimeMs) return hit.lines
+  const text = await readFile(full, "utf8").catch(() => undefined)
+  if (text === undefined) return undefined
+  const value = text.split("\n")
+  if (cache.size >= CACHE) cache.clear()
+  cache.set(full, { mtime: info.mtimeMs, lines: value })
+  return value
+}
+
+/**
+ * Attach trailing context to every thread whose file still matches its hunk.
+ * The anchor check is what keeps this honest: once the agent edits the file,
+ * the commented line no longer matches and the thread keeps hunk-only context
+ * instead of showing unrelated code.
+ */
+export async function withContext(dir: string, comments: PRComment[]): Promise<PRComment[]> {
+  return Promise.all(
+    comments.map(async (item) => {
+      if (!item.file || !item.line || !item.diffHunk || item.outdated) return item
+      const text = anchor(item.diffHunk)
+      if (text === undefined) return item
+      const source = await lines(dir, item.file)
+      if (!source || source[item.line - 1] !== text) return item
+      const after = source.slice(item.line, item.line + AFTER)
+      return after.length > 0 ? { ...item, after } : item
+    }),
+  )
+}

packages/kilo-vscode/src/agent-manager/types.ts

@@ -77,6 +77,8 @@ export interface PRComment {
   outdated: boolean
   createdAt?: number
   diffHunk?: string
+  /** Lines after the commented line, read from the worktree: a hunk has none. */
+  after?: string[]
   replies?: PRCommentReply[]
 }