Kilo-Org · kirillk · Aug 26, 2026 · Aug 25, 2026 · Aug 25, 2026 · Aug 26, 2026
diff --git a/.changeset/jetbrains-worktree-pr-detection.md b/.changeset/jetbrains-worktree-pr-detection.md
@@ -0,0 +1,5 @@
+---
+"@kilocode/kilo-jetbrains": patch
+---
+
+Detect a worktree's pull request reliably in Agent Manager. Imported PRs — including PRs from forks — hand-made worktrees, and locally renamed branches now show their PR badge, the current repository row gets one too, and a freshly imported PR no longer waits out the status poll. Imported PR branches also get proper git tracking, so `git push` and `git pull` work in the new worktree.
diff --git a/.changeset/jetbrains-worktree-tab-artifacts.md b/.changeset/jetbrains-worktree-tab-artifacts.md
@@ -0,0 +1,5 @@
+---
+"@kilocode/kilo-jetbrains": patch
+---
+
+Stop the New Worktree dialog from flashing the previous tab's content when switching between New, From PR, and From Branch.
diff --git a/.../kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/rpc/KiloWorktreeRpcApiImpl.kt b/.../kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/rpc/KiloWorktreeRpcApiImpl.kt
diff --git a/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/rpc/PrResolver.kt b/packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/rpc/PrResolver.kt
@@ -0,0 +1,102 @@
+package ai.kilocode.backend.rpc
+
+import ai.kilocode.rpc.dto.GhAvailability
+import ai.kilocode.rpc.dto.WorktreePrDto
+import kotlinx.serialization.json.Json
+import kotlinx.serialization.json.JsonArray
+import kotlinx.serialization.json.JsonObject
+import kotlinx.serialization.json.jsonPrimitive
+import java.nio.file.Path
+
+/** Result of running a `git`/`gh` command. */
+internal data class CmdOut(val exit: Int, val stdout: String, val stderr: String) {
+    val ok get() = exit == 0
+}
+
+/** PR for one checkout, plus the gh availability observed while resolving it. */
+internal data class PrLookup(val pr: WorktreePrDto? = null, val availability: GhAvailability = GhAvailability.OK)
+
+internal const val PR_FIELDS = "number,state,isDraft,url,title"
+
+/**
+ * Resolves the pull request a checkout belongs to. A worktree can reach a PR in several ways —
+ * Kilo's PR import, `gh pr checkout`, a hand-made `git worktree add`, a branch renamed locally, a
+ * fork PR — so identity is resolved by branch config or head commit rather than by branch name
+ * alone, in increasing order of cost:
+ *
+ * 1. `gh pr view` with no selector. The only form that honours `branch.<name>.merge`, so it
+ *    resolves `refs/pull/N/head` branches by PR number and fork PRs through the push remote.
+ * 2. `gh pr view <branch>`. Matches same-repo branches pushed to origin, no branch config needed.
+ *    Cannot match a fork PR: gh compares against `owner:branch` for cross-repository heads.
+ * 3. `gh pr list --search "<HEAD sha>"`, accepting only an exact `headRefOid` match.
+ *
+ * Commands are injected so the strategy ladder is testable without `gh` or network access.
+ */
+internal class PrResolver(
+    private val gh: (Path, List<String>) -> CmdOut,
+    private val git: (Path, List<String>) -> CmdOut,
+) {
+    /**
+     * Resolves the PR for the checkout at [path] on [branch]. [base] is the repository's base
+     * branch; a PR headed by it is not worth a search query, so strategy 3 is skipped there.
+     */
+    fun resolve(path: String, branch: String, base: String?): PrLookup {
+        val dir = Path.of(path).normalize()
+        view(dir, path, null)?.let { return it }
+        view(dir, path, branch)?.let { return it }
+        if (branch == base) return PrLookup()
+        return search(dir, path) ?: PrLookup()
+    }
+
+    /** Null means "no PR here, keep looking"; a value is terminal (a PR, or gh being unusable). */
+    private fun view(dir: Path, path: String, branch: String?): PrLookup? {
+        val args = buildList {
+            add("pr")
+            add("view")
+            branch?.let { add(it) }
+            add("--json")
+            add(PR_FIELDS)
+        }
+        val out = gh(dir, args)
+        if (!out.ok) return unusable(out.stderr)
+        return parsePr(path, out.stdout)?.let { PrLookup(it) }
+    }
+
+    private fun search(dir: Path, path: String): PrLookup? {
+        val head = git(dir, listOf("rev-parse", "HEAD")).stdout.trim()
+        if (head.isEmpty()) return null
+        val out = gh(
+            dir,
+            listOf("pr", "list", "--state", "all", "--search", "$head is:pr", "--limit", "5", "--json", "$PR_FIELDS,headRefOid"),
+        )
+        if (!out.ok) return unusable(out.stderr)
+        val items = runCatching { json.parseToJsonElement(out.stdout) as? JsonArray }.getOrNull() ?: return null
+        for (item in items) {
+            val obj = item as? JsonObject ?: continue
+            // The search matches commit mentions too, so only an exact head match is our PR.
+            if (obj["headRefOid"]?.jsonPrimitive?.content != head) continue
+            parsePr(path, obj.toString())?.let { return PrLookup(it) }
+        }
+        return null
+    }
+
+    private fun unusable(stderr: String): PrLookup? {
+        val status = prError(stderr)
+        return if (status == GhAvailability.OK) null else PrLookup(availability = status)
+    }
+}
+
+/**
+ * Classifies a failing `gh pr` command. A missing PR is the normal case, so anything that is not a
+ * recognised authorization failure counts as OK — a missing `gh` binary is caught by the upfront
+ * availability probe instead.
+ */
+internal fun prError(stderr: String): GhAvailability {
+    val text = stderr.lowercase()
+    if (text.contains("not logged") || text.contains("gh auth login") || text.contains("authentication")) {
+        return GhAvailability.UNAUTH
+    }
+    return GhAvailability.OK
+}
+
+private val json = Json { ignoreUnknownKeys = true }
diff --git a/...o-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/rpc/KiloWorktreeRpcApiImplTest.kt b/...o-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/rpc/KiloWorktreeRpcApiImplTest.kt
@@ -1,5 +1,6 @@
 package ai.kilocode.backend.rpc
 
+import ai.kilocode.rpc.parsePrUrl
 import ai.kilocode.rpc.dto.CreateWorktreeRequestDto
 import ai.kilocode.rpc.dto.GhAvailability
 import ai.kilocode.rpc.dto.GhState
@@ -22,11 +23,13 @@ import kotlin.test.assertTrue
 
 class KiloWorktreeRpcApiImplTest {
     private val repo: Path = Files.createTempDirectory("kilo-worktree")
+    private val remote: Path = Files.createTempDirectory("kilo-origin")
     private val api = KiloWorktreeRpcApiImpl()
 
     @AfterTest
     fun tearDown() {
         delete(repo)
+        delete(remote)
     }
 
     @Test
@@ -593,9 +596,118 @@ class KiloWorktreeRpcApiImplTest {
     }
 
     @Test
-    fun `parsePrHeadRef reads headRefName`() {
-        assertEquals("feature/login", parsePrHeadRef("""{"headRefName":"feature/login","title":"x"}"""))
-        assertEquals("", parsePrHeadRef("not json"))
+    fun `parsePrHead reads head branch and repository`() {
+        val same = parsePrHead("""{"headRefName":"feature/login","title":"x","isCrossRepository":false}""")
+        assertEquals("feature/login", same.ref)
+        assertFalse(same.cross)
+
+        val fork = parsePrHead(
+            """{"headRefName":"patch-1","isCrossRepository":true,"headRepositoryOwner":{"login":"Contributor"}}""",
+        )
+        assertEquals("patch-1", fork.ref)
+        assertTrue(fork.cross)
+        assertEquals("Contributor", fork.owner)
+
+        assertEquals(PrHead(), parsePrHead("not json"))
+    }
+
+    @Test
+    fun `prBranchName prefixes fork heads and falls back to the pr number`() {
+        assertEquals("feature/login", prBranchName(PrHead("feature/login"), 7))
+        assertEquals("contributor/patch-1", prBranchName(PrHead("patch-1", cross = true, owner = "Contributor"), 7))
+        // A cross-repo PR whose owner gh did not report still needs a usable branch name.
+        assertEquals("patch-1", prBranchName(PrHead("patch-1", cross = true), 7))
+        assertEquals("pr-7", prBranchName(PrHead(), 7))
+    }
+
+    @Test
+    fun `prTargets keeps the main tree and drops detached and prunable entries`() {
+        val items = listOf(
+            WorktreeDto("/repo", "repo", "main", "/repo", main = true),
+            WorktreeDto("/repo/.kilo/worktrees/a", "a", "feature/a", "/repo/.kilo/worktrees/a"),
+            WorktreeDto("/repo/.kilo/worktrees/detached", "detached", "(detached)", "/repo/.kilo/worktrees/detached"),
+            WorktreeDto("/repo/.kilo/worktrees/gone", "gone", "feature/gone", "/repo/.kilo/worktrees/gone", prunable = true),
+        )
+
+        assertEquals(listOf("/repo", "/repo/.kilo/worktrees/a"), prTargets(items).map { it.path })
+    }
+
+    @Test
+    fun `baseBranch reads the main tree branch and ignores a detached one`() {
+        val main = WorktreeDto("/repo", "repo", "main", "/repo", main = true)
+        val linked = WorktreeDto("/repo/.kilo/worktrees/a", "a", "feature/a", "/repo/.kilo/worktrees/a")
+
+        assertEquals("main", baseBranch(listOf(main, linked)))
+        assertNull(baseBranch(listOf(main.copy(branch = "(detached)"), linked)))
+        assertNull(baseBranch(listOf(linked)))
+    }
+
+    @Test
+    fun `fetchPrBranch tracks the head branch for a same-repo pull request`() {
+        initRepo()
+        val origin = originWith(pull = 7, head = "feature/login")
+
+        val failure = fetchPrBranch(runner(repo), 7, PrHead("feature/login"), "feature/login")
+
+        assertNull(failure, "same-repo import should succeed")
+        assertEquals("origin", config("branch.feature/login.remote"))
+        assertEquals("refs/heads/feature/login", config("branch.feature/login.merge"))
+        assertEquals(
+            head(origin, "refs/heads/feature/login"),
+            head(repo, "refs/heads/feature/login"),
+            "local branch should point at the fetched head",
+        )
+    }
+
+    @Test
+    fun `fetchPrBranch falls back to the pull ref when the head branch is gone`() {
+        initRepo()
+        val origin = originWith(pull = 7, head = "feature/login")
+        git(origin, "update-ref", "-d", "refs/heads/feature/login")
+
+        val failure = fetchPrBranch(runner(repo), 7, PrHead("feature/login"), "feature/login")
+
+        assertNull(failure, "import should fall back to the pull ref")
+        assertEquals("refs/pull/7/head", config("branch.feature/login.merge"))
+        assertEquals(head(origin, "refs/pull/7/head"), head(repo, "refs/heads/feature/login"))
+    }
+
+    @Test
+    fun `fetchPrBranch tracks the pull ref for a fork pull request`() {
+        initRepo()
+        val origin = originWith(pull = 7, head = "patch-1")
+        // A fork head is not on origin at all; only the pull ref can reach it.
+        git(origin, "update-ref", "-d", "refs/heads/patch-1")
+        val fork = PrHead("patch-1", cross = true, owner = "contributor")
+
+        val failure = fetchPrBranch(runner(repo), 7, fork, prBranchName(fork, 7))
+
+        assertNull(failure, "fork import should succeed")
+        assertEquals("origin", config("branch.contributor/patch-1.remote"))
+        assertEquals("refs/pull/7/head", config("branch.contributor/patch-1.merge"))
+        assertEquals(head(origin, "refs/pull/7/head"), head(repo, "refs/heads/contributor/patch-1"))
+    }
+
+    @Test
+    fun `fetchPrBranch force updates a branch left by an earlier import`() {
+        initRepo()
+        val origin = originWith(pull = 7, head = "feature/login")
+        git(repo, "branch", "feature/login")
+
+        val failure = fetchPrBranch(runner(repo), 7, PrHead("feature/login"), "feature/login")
+
+        assertNull(failure, "re-import should refresh the stale branch")
+        assertEquals(head(origin, "refs/heads/feature/login"), head(repo, "refs/heads/feature/login"))
+    }
+
+    @Test
+    fun `fetchPrBranch reports the failing command`() {
+        initRepo()
+
+        val failure = fetchPrBranch(runner(repo), 7, PrHead("feature/login"), "feature/login")
+
+        assertNotNull(failure, "a repo without origin cannot fetch a pull request")
+        assertFalse(failure.ok)
     }
 
     @Test
@@ -752,6 +864,39 @@ class KiloWorktreeRpcApiImplTest {
         git(repo, "commit", "-m", "init")
     }
 
+    /**
+     * Builds an "origin" repository holding [head] plus a `refs/pull/<pull>/head` ref pointing at it,
+     * the shape GitHub exposes for a pull request, and registers it as [repo]'s origin.
+     */
+    private fun originWith(pull: Int, head: String): Path {
+        git(remote, "init")
+        git(remote, "config", "user.email", "test@kilo.ai")
+        git(remote, "config", "user.name", "Kilo Test")
+        Files.writeString(remote.resolve("README.md"), "origin")
+        git(remote, "add", "README.md")
+        git(remote, "commit", "-m", "init")
+        val base = output(remote, "branch", "--show-current").trim()
+        git(remote, "checkout", "-b", head)
+        Files.writeString(remote.resolve("pr.txt"), "pr work\n")
+        git(remote, "add", "pr.txt")
+        git(remote, "commit", "-m", "pr work")
+        git(remote, "update-ref", "refs/pull/$pull/head", "refs/heads/$head")
+        // Leave the PR head unchecked out so tests can delete it to emulate a deleted branch.
+        git(remote, "checkout", base)
+        git(repo, "remote", "add", "origin", remote.toString())
+        return remote
+    }
+
+    private fun runner(dir: Path): (List<String>) -> CmdOut = { args ->
+        val cmd = GeneralCommandLine(listOf("git") + args).withWorkDirectory(dir.toFile())
+        val out = CapturingProcessHandler(cmd).runProcess(30_000)
+        CmdOut(if (out.isTimeout) -1 else out.exitCode, out.stdout, out.stderr)
+    }
+
+    private fun config(key: String): String = output(repo, "config", "--get", key).trim()
+
+    private fun head(dir: Path, ref: String): String = output(dir, "rev-parse", ref).trim()
+
     private fun git(dir: Path, vararg args: String) {
         val cmd = GeneralCommandLine(listOf("git") + args).withWorkDirectory(dir.toFile())
         val out = CapturingProcessHandler(cmd).runProcess(30_000)