kilocode/packages/kilo-vscode/tests/unit/worktree-state-manager.test.ts

import { describe, it, expect, beforeEach, afterEach } from "bun:test"
import * as fs from "fs"
import * as path from "path"
import * as os from "os"
import { WorktreeStateManager } from "../../src/agent-manager/WorktreeStateManager"
import { restoreWorktrees } from "../../src/agent-manager/state-recovery"
import type { WorktreeInfo } from "../../src/agent-manager/WorktreeManager"

describe("WorktreeStateManager", () => {
  let root: string
  let manager: WorktreeStateManager
  const logs: string[] = []

  beforeEach(() => {
    root = fs.mkdtempSync(path.join(os.tmpdir(), "wtsm-test-"))
    // Pre-create .kilo dir so fire-and-forget saves don't race on mkdir
    fs.mkdirSync(path.join(root, ".kilo"), { recursive: true })
    logs.length = 0
    manager = new WorktreeStateManager(root, (msg) => logs.push(msg))
  })

  afterEach(async () => {
    await manager.flush()
    fs.rmSync(root, { recursive: true, force: true })
  })

  describe("worktree CRUD", () => {
    it("adds and retrieves worktrees", () => {
      const wt = manager.addWorktree({ branch: "fix-123", path: "/tmp/wt", parentBranch: "main" })
      expect(wt.id).toMatch(/^wt-/)
      expect(wt.branch).toBe("fix-123")
      expect(wt.createdAt).toBeTruthy()

      expect(manager.getWorktrees()).toHaveLength(1)
      expect(manager.getWorktree(wt.id)).toEqual(wt)
    })

    it("finds worktree by path", () => {
      manager.addWorktree({ branch: "a", path: "/tmp/a", parentBranch: "main" })
      const b = manager.addWorktree({ branch: "b", path: "/tmp/b", parentBranch: "main" })

      expect(manager.findWorktreeByPath("/tmp/b")?.id).toBe(b.id)
      expect(manager.findWorktreeByPath("/tmp/c")).toBeUndefined()
    })

    it("removes worktree and deletes its sessions", () => {
      const wt = manager.addWorktree({ branch: "fix", path: "/tmp/fix", parentBranch: "main" })
      manager.addSession("s1", wt.id)
      manager.addSession("s2", wt.id)

      const orphaned = manager.removeWorktree(wt.id)
      expect(orphaned).toHaveLength(2)
      expect(manager.getWorktrees()).toHaveLength(0)
      // Sessions are removed from state
      expect(manager.getSession("s1")).toBeUndefined()
      expect(manager.getSession("s2")).toBeUndefined()
      expect(manager.getSessions()).toHaveLength(0)
    })

    it("returns empty array when removing nonexistent worktree", () => {
      expect(manager.removeWorktree("nonexistent")).toHaveLength(0)
    })

    it("tracks one automatic rename without changing branch ownership", () => {
      const wt = manager.addWorktree({
        branch: "quiet-river",
        path: "/tmp/wt",
        parentBranch: "main",
        branchOwned: true,
      })
      manager.addSession("session-1", wt.id)
      manager.armAutoName(wt.id, "session-1")

      expect(manager.getWorktree(wt.id)?.autoNameSessionId).toBe("session-1")
      expect(manager.renameOwnedBranch(wt.id, "quiet-river", "fix-token-refresh")).toBe(true)
      expect(manager.getWorktree(wt.id)).toMatchObject({
        branch: "fix-token-refresh",
        branchOwned: true,
        autoNameSessionId: undefined,
        originalBranch: undefined,
      })
    })

    it("treats an observed branch change as manual and cancels automatic naming", () => {
      const wt = manager.addWorktree({
        branch: "quiet-river",
        path: "/tmp/wt",
        parentBranch: "main",
        branchOwned: true,
      })
      manager.armAutoName(wt.id, "session-1")

      expect(manager.updateWorktreeBranch(wt.id, "my-manual-name")).toBe(true)
      expect(manager.getWorktree(wt.id)).toMatchObject({
        branch: "my-manual-name",
        originalBranch: "quiet-river",
        autoNameSessionId: undefined,
      })
    })

    it("cancels automatic naming when a worktree gains another session", () => {
      const wt = manager.addWorktree({
        branch: "quiet-river",
        path: "/tmp/wt",
        parentBranch: "main",
        branchOwned: true,
      })
      manager.addSession("session-1", wt.id)
      manager.armAutoName(wt.id, "session-1")
      manager.addSession("session-2", wt.id)
      expect(manager.getWorktree(wt.id)?.autoNameSessionId).toBeUndefined()
    })

    it("never arms imported branches for automatic naming", () => {
      const wt = manager.addWorktree({
        branch: "existing-feature",
        path: "/tmp/wt",
        parentBranch: "main",
        branchOwned: false,
      })
      manager.armAutoName(wt.id, "session-1")
      expect(manager.getWorktree(wt.id)?.autoNameSessionId).toBeUndefined()
    })
  })

  describe("session CRUD", () => {
    it("adds and retrieves sessions", () => {
      const wt = manager.addWorktree({ branch: "fix", path: "/tmp/fix", parentBranch: "main" })
      const s = manager.addSession("sess-1", wt.id)

      expect(s.id).toBe("sess-1")
      expect(s.worktreeId).toBe(wt.id)
      expect(manager.getSession("sess-1")).toEqual(s)
    })

    it("adds session with null worktreeId", () => {
      const s = manager.addSession("local-1", null)
      expect(s.worktreeId).toBeNull()
    })

    it("drops obsolete session prefs while loading state", async () => {
      const file = path.join(root, ".kilo", "agent-manager.json")
      fs.writeFileSync(
        file,
        JSON.stringify({
          worktrees: {},
          sessions: {
            "local-1": {
              worktreeId: null,
              createdAt: new Date().toISOString(),
              prefs: { agent: "code" },
            },
          },
        }),
      )

      await manager.load()

      expect(manager.getSession("local-1")).not.toHaveProperty("prefs")
    })

    it("filters sessions by worktreeId", () => {
      const wt1 = manager.addWorktree({ branch: "a", path: "/tmp/a", parentBranch: "main" })
      const wt2 = manager.addWorktree({ branch: "b", path: "/tmp/b", parentBranch: "main" })
      manager.addSession("s1", wt1.id)
      manager.addSession("s2", wt1.id)
      manager.addSession("s3", wt2.id)

      expect(manager.getSessions(wt1.id)).toHaveLength(2)
      expect(manager.getSessions(wt2.id)).toHaveLength(1)
      expect(manager.getSessions()).toHaveLength(3)
    })

    it("moves session to a different worktree", () => {
      const wt1 = manager.addWorktree({ branch: "a", path: "/tmp/a", parentBranch: "main" })
      const wt2 = manager.addWorktree({ branch: "b", path: "/tmp/b", parentBranch: "main" })
      manager.addSession("s1", wt1.id)

      manager.moveSession("s1", wt2.id)
      expect(manager.getSession("s1")?.worktreeId).toBe(wt2.id)
    })

    it("moves session back to local (null worktreeId)", () => {
      const wt = manager.addWorktree({ branch: "a", path: "/tmp/a", parentBranch: "main" })
      manager.addSession("s1", wt.id)
      expect(manager.getSession("s1")?.worktreeId).toBe(wt.id)

      manager.moveSession("s1", null)
      expect(manager.getSession("s1")?.worktreeId).toBeNull()
    })

    it("moveSession is a no-op for nonexistent session", () => {
      manager.moveSession("nonexistent", "wt-1")
      expect(manager.getSessions()).toHaveLength(0)
    })

    it("removes session", () => {
      manager.addSession("s1", null)
      manager.removeSession("s1")
      expect(manager.getSession("s1")).toBeUndefined()
    })
  })

  describe("directoryFor", () => {
    it("returns worktree path for worktree session", () => {
      const wt = manager.addWorktree({ branch: "fix", path: "/tmp/fix", parentBranch: "main" })
      manager.addSession("s1", wt.id)
      expect(manager.directoryFor("s1")).toBe("/tmp/fix")
    })

    it("returns undefined for local session", () => {
      manager.addSession("s1", null)
      expect(manager.directoryFor("s1")).toBeUndefined()
    })

    it("returns undefined for unknown session", () => {
      expect(manager.directoryFor("nonexistent")).toBeUndefined()
    })
  })

  describe("worktreeSessionIds", () => {
    it("returns only session IDs that belong to worktrees", () => {
      const wt = manager.addWorktree({ branch: "fix", path: "/tmp/fix", parentBranch: "main" })
      manager.addSession("s1", wt.id)
      manager.addSession("s2", null)
      manager.addSession("s3", wt.id)

      const ids = manager.worktreeSessionIds()
      expect(ids.size).toBe(2)
      expect(ids.has("s1")).toBe(true)
      expect(ids.has("s3")).toBe(true)
      expect(ids.has("s2")).toBe(false)
    })
  })

  describe("persistence", () => {
    it("saves and loads state, preserving local sessions and pruning orphaned sessions", async () => {
      const wt = manager.addWorktree({ branch: "fix", path: "/tmp/fix", parentBranch: "main" })
      manager.addSession("s1", wt.id)
      manager.addSession("s2", null)
      manager.addSession("s3", "missing")
      // Flush fire-and-forget saves from mutations, then do a final save
      await manager.flush()
      await manager.save()

      const loaded = new WorktreeStateManager(root, () => {})
      const result = await loaded.load()

      expect(result.status).toBe("loaded")
      expect(loaded.getWorktrees()).toHaveLength(1)
      expect(loaded.getWorktrees()[0].branch).toBe("fix")
      expect(loaded.getSessions()).toHaveLength(2)
      expect(loaded.getSession("s1")?.worktreeId).toBe(wt.id)
      expect(loaded.getSession("s2")?.worktreeId).toBeNull()
      expect(loaded.getSession("s3")).toBeUndefined()
    })

    it("load is a no-op when file does not exist", async () => {
      const result = await manager.load()
      expect(result.status).toBe("missing")
      expect(manager.getWorktrees()).toHaveLength(0)
      expect(manager.getSessions()).toHaveLength(0)
    })

    it("flush waits for saves queued during an in-flight write", async () => {
      const file = path.join(root, ".kilo", "agent-manager.json")
      const api = fs.promises as unknown as { writeFile: (...args: unknown[]) => Promise<void> }
      const original = api.writeFile
      const gate = {
        release: () => {},
        promise: Promise.resolve(),
      }
      gate.promise = new Promise<void>((resolve) => {
        gate.release = resolve
      })
      const state = { blocked: false }

      api.writeFile = async (...args: unknown[]) => {
        const target = typeof args[0] === "string" ? args[0] : ""
        if (!state.blocked && target.includes("agent-manager.json.")) {
          state.blocked = true
          await gate.promise
        }
        await original(...args)
      }

      try {
        manager.addSession("first", null)
        await Promise.resolve()
        manager.addSession("second", null)
        gate.release()
        await manager.flush()
      } finally {
        api.writeFile = original
      }

      const data = JSON.parse(fs.readFileSync(file, "utf-8")) as { sessions: Record<string, unknown> }
      expect(data.sessions.first).toBeDefined()
      expect(data.sessions.second).toBeDefined()
    })

    it("creates .kilo directory if missing", async () => {
      const fresh = path.join(root, "subdir")
      const mgr = new WorktreeStateManager(fresh, () => {})
      mgr.addWorktree({ branch: "test", path: "/tmp/test", parentBranch: "main" })
      await mgr.flush()
      await mgr.save()

      expect(fs.existsSync(path.join(fresh, ".kilo", "agent-manager.json"))).toBe(true)
    })

    it("does not overwrite a corrupt state file after load fails", async () => {
      const file = path.join(root, ".kilo", "agent-manager.json")
      fs.writeFileSync(file, "{", "utf-8")

      const result = await manager.load()
      manager.addSession("local-after-failure", null)
      await manager.flush()
      await manager.save()

      expect(result.status).toBe("failed")
      expect(fs.readFileSync(file, "utf-8")).toBe("{")
      expect(logs.some((l) => l.includes("Skipping save because state failed to load"))).toBe(true)
    })

    it("backs up a corrupt state file before recovery saves", async () => {
      const file = path.join(root, ".kilo", "agent-manager.json")
      fs.writeFileSync(file, "{", "utf-8")

      const result = await manager.load()
      const recovered = await manager.prepareRecovery()
      manager.addSession("local-after-recovery", null)
      await manager.flush()

      const files = fs.readdirSync(path.join(root, ".kilo"))
      expect(result.status).toBe("failed")
      expect(recovered).toBe(true)
      expect(files.some((item) => item.startsWith("agent-manager.json.corrupt-"))).toBe(true)
      expect(JSON.parse(fs.readFileSync(file, "utf-8")).sessions["local-after-recovery"].worktreeId).toBeNull()
    })

    it("allows saves after a later missing-file reload", async () => {
      const file = path.join(root, ".kilo", "agent-manager.json")
      fs.writeFileSync(file, "{", "utf-8")

      await manager.load()
      fs.rmSync(file)
      await manager.load()
      manager.addSession("local-after-missing", null)
      await manager.flush()

      expect(JSON.parse(fs.readFileSync(file, "utf-8")).sessions["local-after-missing"].worktreeId).toBeNull()
    })
  })

  describe("recovery", () => {
    it("restores discovered worktrees and metadata sessions", async () => {
      const infos: WorktreeInfo[] = [
        {
          branch: "fix-recovered",
          path: "/tmp/recovered",
          parentBranch: "main",
          remote: "origin",
          createdAt: Date.UTC(2026, 0, 1),
          sessionId: "sess-recovered",
        },
      ]

      const result = restoreWorktrees(manager, infos)
      await manager.flush()

      const worktree = manager.findWorktreeByPath("/tmp/recovered")
      expect(result).toEqual({ worktrees: 1, sessions: 1 })
      expect(worktree?.branch).toBe("fix-recovered")
      expect(worktree?.remote).toBe("origin")
      expect(manager.getSession("sess-recovered")?.worktreeId).toBe(worktree?.id)
    })
  })

  describe("tab order", () => {
    it("sets and gets tab order for a key", () => {
      manager.setTabOrder("wt-1", ["s1", "s2", "s3"])
      expect(manager.getTabOrder()["wt-1"]).toEqual(["s1", "s2", "s3"])
    })

    it("overwrites existing tab order", () => {
      manager.setTabOrder("wt-1", ["s1", "s2"])
      manager.setTabOrder("wt-1", ["s2", "s1"])
      expect(manager.getTabOrder()["wt-1"]).toEqual(["s2", "s1"])
    })

    it("removes tab order for a key", () => {
      manager.setTabOrder("wt-1", ["s1"])
      manager.removeTabOrder("wt-1")
      expect(manager.getTabOrder()["wt-1"]).toBeUndefined()
    })

    it("removeTabOrder is a no-op for missing key", () => {
      manager.removeTabOrder("nonexistent")
      expect(Object.keys(manager.getTabOrder())).toHaveLength(0)
    })

    it("cleans up tab order when worktree is removed", () => {
      const wt = manager.addWorktree({ branch: "fix", path: "/tmp/fix", parentBranch: "main" })
      manager.addSession("s1", wt.id)
      manager.setTabOrder(wt.id, ["s1"])

      manager.removeWorktree(wt.id)
      expect(manager.getTabOrder()[wt.id]).toBeUndefined()
    })

    it("removes session from tab order arrays when session is removed", () => {
      const wt = manager.addWorktree({ branch: "fix", path: "/tmp/fix", parentBranch: "main" })
      manager.addSession("s1", wt.id)
      manager.addSession("s2", wt.id)
      manager.setTabOrder(wt.id, ["s1", "s2"])

      manager.removeSession("s1")
      expect(manager.getTabOrder()[wt.id]).toEqual(["s2"])
    })

    it("removes tab order entry when last session in order is removed", () => {
      manager.addSession("s1", null)
      manager.setTabOrder("local", ["s1"])

      manager.removeSession("s1")
      expect(manager.getTabOrder()["local"]).toBeUndefined()
    })

    it("persists and loads tab order", async () => {
      const wt = manager.addWorktree({ branch: "fix", path: "/tmp/fix", parentBranch: "main" })
      manager.setTabOrder(wt.id, ["s2", "s1"])
      manager.setTabOrder("local", ["s3", "s4"])
      await manager.flush()
      await manager.save()

      const loaded = new WorktreeStateManager(root, () => {})
      await loaded.load()

      expect(loaded.getTabOrder()[wt.id]).toEqual(["s2", "s1"])
      expect(loaded.getTabOrder()["local"]).toEqual(["s3", "s4"])
    })

    it("does not persist empty tab order", async () => {
      manager.addWorktree({ branch: "fix", path: "/tmp/fix", parentBranch: "main" })
      await manager.flush()
      await manager.save()

      const content = fs.readFileSync(path.join(root, ".kilo", "agent-manager.json"), "utf-8")
      const data = JSON.parse(content)
      expect(data.tabOrder).toBeUndefined()
    })
  })

  describe("sessionsCollapsed", () => {
    it("defaults to true when state is missing", async () => {
      await manager.load()

      expect(manager.getSessionsCollapsed()).toBe(true)
    })

    it("preserves the expanded default from legacy state", async () => {
      const file = path.join(root, ".kilo", "agent-manager.json")
      fs.writeFileSync(file, JSON.stringify({ worktrees: {}, sessions: {} }))

      await manager.load()

      expect(manager.getSessionsCollapsed()).toBe(false)
    })

    it("sets and gets collapsed state", () => {
      manager.setSessionsCollapsed(true)
      expect(manager.getSessionsCollapsed()).toBe(true)

      manager.setSessionsCollapsed(false)
      expect(manager.getSessionsCollapsed()).toBe(false)
    })

    it("persists and loads collapsed state", async () => {
      manager.setSessionsCollapsed(true)
      await manager.flush()
      await manager.save()

      const loaded = new WorktreeStateManager(root, () => {})
      await loaded.load()
      expect(loaded.getSessionsCollapsed()).toBe(true)
    })

    it("persists and loads expanded state", async () => {
      manager.setSessionsCollapsed(false)
      await manager.flush()
      await manager.save()

      const content = fs.readFileSync(path.join(root, ".kilo", "agent-manager.json"), "utf-8")
      const data = JSON.parse(content)
      expect(data.sessionsCollapsed).toBe(false)

      const loaded = new WorktreeStateManager(root, () => {})
      await loaded.load()
      expect(loaded.getSessionsCollapsed()).toBe(false)
    })
  })

  describe("sidebarCollapsed", () => {
    it("defaults to false", () => {
      expect(manager.getSidebarCollapsed()).toBe(false)
    })

    it("sets and gets collapsed state", () => {
      manager.setSidebarCollapsed(true)
      expect(manager.getSidebarCollapsed()).toBe(true)

      manager.setSidebarCollapsed(false)
      expect(manager.getSidebarCollapsed()).toBe(false)
    })

    it("persists and loads collapsed state", async () => {
      manager.setSidebarCollapsed(true)
      await manager.flush()
      await manager.save()

      const loaded = new WorktreeStateManager(root, () => {})
      await loaded.load()
      expect(loaded.getSidebarCollapsed()).toBe(true)
    })

    it("does not persist when false", async () => {
      manager.setSidebarCollapsed(false)
      await manager.flush()
      await manager.save()

      const content = fs.readFileSync(path.join(root, ".kilo", "agent-manager.json"), "utf-8")
      const data = JSON.parse(content)
      expect(data.sidebarCollapsed).toBeUndefined()
    })
  })

  describe("validate", () => {
    it("removes worktrees whose directories do not exist and prunes their sessions", async () => {
      const existing = path.join(root, "wt-exists")
      fs.mkdirSync(existing, { recursive: true })

      manager.addWorktree({ branch: "exists", path: existing, parentBranch: "main" })
      const gone = manager.addWorktree({ branch: "gone", path: path.join(root, "wt-gone"), parentBranch: "main" })
      manager.addSession("s1", gone.id)

      await manager.validate(root)

      expect(manager.getWorktrees()).toHaveLength(1)
      expect(manager.getWorktrees()[0].branch).toBe("exists")
      // Session removed along with its worktree
      expect(manager.getSession("s1")).toBeUndefined()
    })

    it("preserves local sessions and prunes missing worktree references on validate", async () => {
      const existing = path.join(root, "wt-exists")
      fs.mkdirSync(existing, { recursive: true })

      const wt = manager.addWorktree({ branch: "exists", path: existing, parentBranch: "main" })
      manager.addSession("s1", wt.id)
      manager.addSession("s2", null)
      manager.addSession("s3", "missing")

      await manager.validate(root)

      expect(manager.getSession("s1")).toBeTruthy()
      expect(manager.getSession("s2")?.worktreeId).toBeNull()
      expect(manager.getSession("s3")).toBeUndefined()
    })

    it("resolves relative paths against root", async () => {
      const relative = ".kilo/worktrees/test-branch"
      const absolute = path.join(root, relative)
      fs.mkdirSync(absolute, { recursive: true })

      manager.addWorktree({ branch: "test", path: relative, parentBranch: "main" })
      await manager.validate(root)

      expect(manager.getWorktrees()).toHaveLength(1)
    })
  })

  describe("concurrent save serialization", () => {
    it("rapid mutations do not lose data after flush", async () => {
      // Fire many mutations without awaiting saves individually
      for (let i = 0; i < 20; i++) {
        manager.addWorktree({ branch: `b-${i}`, path: `/tmp/b-${i}`, parentBranch: "main" })
      }
      const wts = manager.getWorktrees()
      for (let i = 0; i < 20; i++) {
        manager.addSession(`s-${i}`, wts[i]!.id)
      }

      // Wait for all fire-and-forget saves to settle
      await manager.flush()
      await manager.save()

      // Reload from disk and verify all data persisted
      const loaded = new WorktreeStateManager(root, () => {})
      await loaded.load()

      expect(loaded.getWorktrees()).toHaveLength(20)
      expect(loaded.getSessions()).toHaveLength(20)
      for (let i = 0; i < 20; i++) {
        expect(loaded.getWorktrees().find((w) => w.branch === `b-${i}`)).toBeTruthy()
        expect(loaded.getSession(`s-${i}`)).toBeTruthy()
      }
    })

    it("interleaved add and remove persists correctly", async () => {
      const wt1 = manager.addWorktree({ branch: "keep", path: "/tmp/keep", parentBranch: "main" })
      const wt2 = manager.addWorktree({ branch: "remove", path: "/tmp/remove", parentBranch: "main" })
      manager.addSession("s1", wt1.id)
      manager.addSession("s2", wt2.id)
      manager.removeWorktree(wt2.id)
      manager.addSession("s3", wt1.id)

      await manager.flush()
      await manager.save()

      const loaded = new WorktreeStateManager(root, () => {})
      await loaded.load()

      expect(loaded.getWorktrees()).toHaveLength(1)
      expect(loaded.getWorktrees()[0].branch).toBe("keep")
      // s2 was removed when wt2 was removed, s1 and s3 belong to wt1
      expect(loaded.getSession("s1")?.worktreeId).toBe(wt1.id)
      expect(loaded.getSession("s2")).toBeUndefined()
      expect(loaded.getSession("s3")?.worktreeId).toBe(wt1.id)
    })

    it("concurrent save() calls resolve without data loss", async () => {
      manager.addWorktree({ branch: "first", path: "/tmp/first", parentBranch: "main" })

      // Trigger multiple saves concurrently — the second should queue behind the first
      const p1 = manager.save()
      manager.addWorktree({ branch: "second", path: "/tmp/second", parentBranch: "main" })
      const p2 = manager.save()
      await Promise.all([p1, p2])

      const loaded = new WorktreeStateManager(root, () => {})
      await loaded.load()
      expect(loaded.getWorktrees()).toHaveLength(2)
    })

    it("flush resolves after in-flight save completes", async () => {
      manager.addWorktree({ branch: "flush-test", path: "/tmp/flush", parentBranch: "main" })
      // Don't await — let save fire in background
      void manager.save()
      // flush must wait for it
      await manager.flush()

      const loaded = new WorktreeStateManager(root, () => {})
      await loaded.load()
      expect(loaded.getWorktrees().find((w) => w.branch === "flush-test")).toBeTruthy()
    })
  })

  describe("load with corrupt data", () => {
    it("handles malformed JSON gracefully", async () => {
      const file = path.join(root, ".kilo", "agent-manager.json")
      fs.writeFileSync(file, "not-valid-json{{{", "utf-8")

      await manager.load()

      // State should be empty — no crash
      expect(manager.getWorktrees()).toHaveLength(0)
      expect(manager.getSessions()).toHaveLength(0)
      // Should have logged an error
      expect(logs.some((l) => l.includes("Failed to load state"))).toBe(true)
    })

    it("handles partial data with missing sessions key", async () => {
      const file = path.join(root, ".kilo", "agent-manager.json")
      fs.writeFileSync(
        file,
        JSON.stringify({
          worktrees: { "wt-1": { branch: "a", path: "/a", parentBranch: "main", createdAt: new Date().toISOString() } },
        }),
        "utf-8",
      )

      await manager.load()

      expect(manager.getWorktrees()).toHaveLength(1)
      expect(manager.getWorktrees()[0].branch).toBe("a")
      expect(manager.getSessions()).toHaveLength(0)
    })

    it("handles partial data with missing worktrees key and local sessions", async () => {
      const file = path.join(root, ".kilo", "agent-manager.json")
      fs.writeFileSync(
        file,
        JSON.stringify({ sessions: { "s-1": { worktreeId: null, createdAt: new Date().toISOString() } } }),
        "utf-8",
      )

      await manager.load()

      expect(manager.getWorktrees()).toHaveLength(0)
      expect(manager.getSession("s-1")?.worktreeId).toBeNull()
    })
  })

  describe("legacy .kilocode migration", () => {
    it("migrates and loads state from .kilocode when .kilo is absent", async () => {
      // Remove the .kilo dir created in beforeEach
      fs.rmSync(path.join(root, ".kilo"), { recursive: true, force: true })

      // Write state to legacy .kilocode dir (migration will move it to .kilo)
      const legacyDir = path.join(root, ".kilocode")
      fs.mkdirSync(legacyDir, { recursive: true })
      fs.writeFileSync(
        path.join(legacyDir, "agent-manager.json"),
        JSON.stringify({
          worktrees: {
            "wt-legacy": {
              branch: "legacy-branch",
              path: "/tmp/legacy",
              parentBranch: "main",
              createdAt: new Date().toISOString(),
            },
          },
          sessions: {},
        }),
        "utf-8",
      )

      await manager.load()

      expect(manager.getWorktrees()).toHaveLength(1)
      expect(manager.getWorktrees()[0].branch).toBe("legacy-branch")
    })

    it("skips migration when .kilo state already exists", async () => {
      // .kilo state already present — migration should skip agent-manager.json
      fs.writeFileSync(
        path.join(root, ".kilo", "agent-manager.json"),
        JSON.stringify({
          worktrees: {
            "wt-new": {
              branch: "new-branch",
              path: "/tmp/new",
              parentBranch: "main",
              createdAt: new Date().toISOString(),
            },
          },
          sessions: {},
        }),
        "utf-8",
      )

      // Legacy .kilocode state
      const legacyDir = path.join(root, ".kilocode")
      fs.mkdirSync(legacyDir, { recursive: true })
      fs.writeFileSync(
        path.join(legacyDir, "agent-manager.json"),
        JSON.stringify({
          worktrees: {
            "wt-old": {
              branch: "old-branch",
              path: "/tmp/old",
              parentBranch: "main",
              createdAt: new Date().toISOString(),
            },
          },
          sessions: {},
        }),
        "utf-8",
      )

      await manager.load()

      expect(manager.getWorktrees()).toHaveLength(1)
      expect(manager.getWorktrees()[0].branch).toBe("new-branch")
    })

    it("rewrites stale .kilocode paths in worktree entries (unix)", async () => {
      fs.writeFileSync(
        path.join(root, ".kilo", "agent-manager.json"),
        JSON.stringify({
          worktrees: {
            "wt-stale": {
              branch: "fix",
              path: "/repo/.kilocode/worktrees/fix",
              parentBranch: "main",
              createdAt: new Date().toISOString(),
            },
          },
          sessions: {},
        }),
        "utf-8",
      )

      await manager.load()

      expect(manager.getWorktrees()[0].path).toBe(`/repo/.kilo/worktrees/fix`)
    })

    it("rewrites stale .kilocode paths with backslashes (windows)", async () => {
      fs.writeFileSync(
        path.join(root, ".kilo", "agent-manager.json"),
        JSON.stringify({
          worktrees: {
            "wt-win": {
              branch: "fix",
              path: "C:\\.kilocode\\worktrees\\fix",
              parentBranch: "main",
              createdAt: new Date().toISOString(),
            },
          },
          sessions: {},
        }),
        "utf-8",
      )

      await manager.load()

      // Separator style from the stored path is preserved (backslashes stay as backslashes)
      expect(manager.getWorktrees()[0].path).toBe("C:\\.kilo\\worktrees\\fix")
    })
  })
})