From c83ff8d5982e6ebf92623911e232f1dbd0b0a00c Mon Sep 17 00:00:00 2001
From: Folke Lemaitre <folke.lemaitre@gmail.com>
Date: Thu, 6 Nov 2025 12:00:29 +0100
Subject: [PATCH] feat(gh): add inline review comment annotations to diff
 viewer
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Refactors the diff renderer to support displaying GitHub review comments
inline with the relevant diff lines. Links review comments to specific
file positions using reviewThreads data from the GraphQL API, enabling a
more intuitive code review experience directly within the diff view.

- Refactored gh/api.lua to async fetch review comments and threads
- Added annotation support to diff renderer with context object pattern
- Separated diff parsing from formatting for better maintainability
- Integrated review comment positioning via reviewThreads linkage

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
---
 lua/snacks/gh/api.lua                |  77 ++++--
 lua/snacks/gh/buf.lua                |   2 +-
 lua/snacks/gh/render/init.lua        | 108 +++++++--
 lua/snacks/gh/types.lua              |   6 +
 lua/snacks/picker/preview.lua        |   5 +-
 lua/snacks/picker/source/diff.lua    |   2 +
 lua/snacks/picker/source/gh.lua      |  39 +--
 lua/snacks/picker/util/async.lua     |   3 +
 lua/snacks/picker/util/diff.lua      | 349 +++++++++++++++++++--------
 lua/snacks/picker/util/highlight.lua |   4 +-
 lua/snacks/util/spawn.lua            |   2 +-
 11 files changed, 442 insertions(+), 155 deletions(-)

diff --git a/lua/snacks/gh/api.lua b/lua/snacks/gh/api.lua
index 5074e3d1..928cd467 100644
--- a/lua/snacks/gh/api.lua
+++ b/lua/snacks/gh/api.lua
@@ -320,7 +320,7 @@ function M.view(cb, item, opts)
     api_opts.fields = vim.list_extend(api_opts.fields, opts.fields)
   end
 
-  item = not Item.is(item) and cache_get(item) or item
+  item = M.get_cached(item)
   local todo = Item.is(item) and item:need(api_opts.fields) or api_opts.fields
   if opts.force or item.dirty then
     todo = api_opts.fields
@@ -335,13 +335,20 @@ function M.view(cb, item, opts)
   local args = { item.type, "view", tostring(item.number) }
   local need_reviews = item.type == "pr" and vim.tbl_contains(todo, "comments")
   local it ---@type snacks.gh.Item?
-  local pending = need_reviews and 2 or 1
+  local completed = 0
+  local fetch_comments = false
+  local procs = {} ---@type snacks.spawn.Proc[]
 
   ---@param data? snacks.gh.Item|{}
   local function handler(data)
     it = data and vim.tbl_extend("force", it or {}, data or {}) or it
-    pending = pending - 1
-    if pending > 0 then
+    if fetch_comments then
+      fetch_comments = false
+      item.repo = it and Item.get_repo(it.url) or nil
+      procs[#procs + 1] = M.comments(item, handler)
+    end
+    completed = completed + 1
+    if completed < #procs then
       return
     end
     if not it then
@@ -356,21 +363,52 @@ function M.view(cb, item, opts)
     todo = vim.tbl_filter(function(f)
       return f ~= "comments" and f ~= "reviews"
     end, todo)
-    M.comments(item, handler)
+    if item.repo then
+      procs[#procs + 1] = M.comments(item, handler)
+    else
+      -- fetch comments once we fetched the item
+      fetch_comments = true
+    end
   end
 
-  ---@param data? snacks.gh.Item
-  return M.fetch(function(_, data)
-    handler(data)
-  end, {
-    args = args,
-    fields = todo,
-    repo = item.repo or api_opts.repo,
-  })
+  if #todo > 0 then
+    ---@param data? snacks.gh.Item
+    procs[#procs + 1] = M.fetch(function(_, data)
+      handler(data)
+    end, {
+      args = args,
+      fields = todo,
+      repo = item.repo or api_opts.repo,
+    })
+  end
+
+  ---@type snacks.picker.Waitable
+  return {
+    ---@async
+    wait = function()
+      for _, proc in ipairs(procs) do
+        proc:wait()
+      end
+    end,
+  }
 end
 
 ---@param item snacks.gh.api.View
-function M.get(item)
+---@param opts? { fields?: string[], force?: boolean }
+---@async
+function M.get(item, opts)
+  local ret ---@type snacks.picker.gh.Item?
+  local procs = M.view(function(it)
+    ret = it
+  end, item, opts)
+  if procs then
+    procs:wait()
+  end
+  return ret
+end
+
+---@param item snacks.gh.api.View
+function M.get_cached(item)
   return not Item.is(item) and cache_get(item) or item
 end
 
@@ -406,6 +444,17 @@ function M.comments(item, cb)
       query($owner: String!, $name: String!, $number: Int!) {
         repository(owner: $owner, name: $name) {
           pullRequest(number: $number) {
+            reviewThreads(first: 100) {
+              nodes {
+                id
+                diffSide
+                comments(first: 50) {
+                  nodes {
+                    id
+                  }
+                }
+              }
+            }
             reviews(first: 100) {
               nodes {
                 id
diff --git a/lua/snacks/gh/buf.lua b/lua/snacks/gh/buf.lua
index 95a6c38d..e92f467d 100644
--- a/lua/snacks/gh/buf.lua
+++ b/lua/snacks/gh/buf.lua
@@ -92,7 +92,7 @@ function M:render(opts)
     return
   end
   opts = opts or {}
-  self.item = Api.get(self.item)
+  self.item = Api.get_cached(self.item)
 
   self:bo()
   self:wo()
diff --git a/lua/snacks/gh/render/init.lua b/lua/snacks/gh/render/init.lua
index 4f4bdbe0..d29c9eac 100644
--- a/lua/snacks/gh/render/init.lua
+++ b/lua/snacks/gh/render/init.lua
@@ -9,11 +9,13 @@ local util = Snacks.picker.util
 -- 1. As top-level review.comments
 -- 2. As replies in the thread tree
 ---@class snacks.gh.render.ctx
----@field buf number
 ---@field item snacks.picker.gh.Item
 ---@field opts snacks.gh.Config
 ---@field comment_skip table<string, boolean>
 ---@field is_review? boolean
+---@field diff? boolean render diffs (defaults to true)
+---@field markdown? boolean render in a markdown buffer (defaults to true)
+---@field annotations? snacks.diff.Annotation[]
 
 ---@param field string
 local function time_prop(field)
@@ -250,7 +252,6 @@ function M.render(buf, item, opts)
 
   ---@type snacks.gh.render.ctx
   local ctx = {
-    buf = buf,
     item = item,
     opts = opts,
     comment_skip = {},
@@ -418,9 +419,11 @@ function M.comment_body(body, ctx)
 end
 
 ---@param lines snacks.picker.Highlight[][]
-function M.indent(lines)
+---@param ctx snacks.gh.render.ctx
+function M.indent(lines, ctx)
+  -- indent guides for lines after the first
   local indent = {} ---@type snacks.picker.Highlight[]
-  -- virtual overlay showing indent guides
+  indent[#indent + 1] = { "   ", "Normal" }
   indent[#indent + 1] = {
     col = 0,
     virt_text = {
@@ -432,23 +435,24 @@ function M.indent(lines)
     hl_mode = "combine",
     virt_text_repeat_linebreak = true,
   }
-  -- actual indent space
-  local first = vim.deepcopy(indent)
-  first = {
-    {
-      col = 0,
-      end_col = 3,
-      conceal = "",
-      priority = 1000,
-    },
-    { " * ", "Normal" },
-  }
-  local other = vim.deepcopy(indent)
-  table.insert(other, 1, { "   ", "Normal" })
+
+  --- first indent. In a markdown buffer, we need proper structure,
+  --- so we conceal the list marker
+  ---@type snacks.picker.Highlight[]
+  local first = ctx.markdown == false and {}
+    or {
+      {
+        col = 0,
+        end_col = 3,
+        conceal = "",
+        priority = 1000,
+      },
+      { " * ", "Normal" },
+    }
 
   local ret = {} ---@type snacks.picker.Highlight[][]
   for l, line in ipairs(lines) do
-    local new = vim.deepcopy(l == 1 and first or other)
+    local new = vim.deepcopy(l == 1 and first or indent)
     extend(new, line)
     ret[l] = new
   end
@@ -467,6 +471,7 @@ function M.comment_diff(comment, ctx)
     count = originalLine - comment.originalStartLine + 1
   end
   count = math.max(ctx.opts.diff.min, math.abs(count))
+
   local Diff = require("snacks.picker.util.diff")
   local diff = ("diff --git a/%s b/%s\n%s"):format(comment.path, comment.path, comment.diffHunk)
   local ret = Diff.format(diff, {
@@ -478,6 +483,33 @@ function M.comment_diff(comment, ctx)
   return ret
 end
 
+---@param comment snacks.gh.Comment
+---@param ctx snacks.gh.render.ctx
+function M.annotate(comment, ctx)
+  if not comment.path or not comment.diffHunk then
+    return
+  end
+  local side = "right"
+  for _, thread in ipairs(ctx.item.reviewThreads or {}) do
+    for _, c in ipairs(thread.comments or {}) do
+      if c.id == comment.id then
+        side = (thread.diffSide or "RIGHT"):lower()
+        break
+      end
+    end
+  end
+  ---@type snacks.diff.Annotation
+  local ret = {
+    side = side,
+    file = comment.path,
+    line = comment.line or comment.originalLine or 1,
+    text = {},
+  }
+  ctx.annotations = ctx.annotations or {}
+  table.insert(ctx.annotations, ret)
+  return ret
+end
+
 ---@param comment snacks.gh.Comment
 ---@param ctx snacks.gh.render.ctx
 function M.comment(comment, ctx)
@@ -487,12 +519,16 @@ function M.comment(comment, ctx)
   extend(header, M.comment_header(comment, {}, ctx))
   ret[#ret + 1] = header
 
+  local annotation ---@type snacks.diff.Annotation?
   if not comment.replyTo then
-    -- add diff hunk for top-level comments
-    local diff = M.comment_diff(comment, ctx)
-    if #diff > 0 then
-      vim.list_extend(ret, diff)
-      ret[#ret + 1] = {} -- empty line between diff and body
+    annotation = M.annotate(comment, ctx)
+    if ctx.diff ~= false then
+      -- add diff hunk for top-level comments
+      local diff = M.comment_diff(comment, ctx)
+      if #diff > 0 then
+        vim.list_extend(ret, diff)
+        ret[#ret + 1] = {} -- empty line between diff and body
+      end
     end
   end
 
@@ -511,7 +547,11 @@ function M.comment(comment, ctx)
       end
     end
   end
-  return M.indent(ret)
+  ret = M.indent(ret, ctx)
+  if annotation then
+    annotation.text = vim.deepcopy(ret)
+  end
+  return ret
 end
 
 ---@param id string
@@ -559,7 +599,25 @@ function M.review(review, ctx)
     ret[#ret + 1] = {} -- empty line between review and comments
     vim.list_extend(ret, M.comment(comment, ctx))
   end
-  return M.indent(ret)
+  return M.indent(ret, ctx)
+end
+
+---@param pr snacks.picker.gh.Item
+function M.annotations(pr)
+  ---@type snacks.gh.render.ctx
+  local ctx = {
+    item = pr,
+    opts = Snacks.gh.config(),
+    comment_skip = {},
+    is_review = true,
+    diff = false,
+    markdown = false,
+  }
+  for _, review in ipairs(pr.reviews or {}) do
+    review.created = review.submitted
+    M.review(review, ctx)
+  end
+  return ctx.annotations
 end
 
 return M
diff --git a/lua/snacks/gh/types.lua b/lua/snacks/gh/types.lua
index e21d1010..b284751d 100644
--- a/lua/snacks/gh/types.lua
+++ b/lua/snacks/gh/types.lua
@@ -90,6 +90,11 @@
 ---@field context? string
 ---@field state? "SUCCESS" | "FAILURE" | "PENDING"
 
+---@class snacks.gh.review.Thread
+---@field id string
+---@field diffSide "LEFT" | "RIGHT"
+---@field comments {id: string}[]
+
 ---@class snacks.gh.Review
 ---@field id string
 ---@field author snacks.gh.User
@@ -128,6 +133,7 @@
 ---@field headRefName? string
 ---@field isDraft? boolean
 ---@field reviews? snacks.gh.Review[]
+---@field reviewThreads? snacks.gh.review.Thread[]
 
 ---@class snacks.gh.Commit
 ---@field oid string
diff --git a/lua/snacks/picker/preview.lua b/lua/snacks/picker/preview.lua
index ec76f41a..e2227f57 100644
--- a/lua/snacks/picker/preview.lua
+++ b/lua/snacks/picker/preview.lua
@@ -185,7 +185,10 @@ end
 ---@param ft "diff"|"git"
 ---@param ctx snacks.picker.preview.ctx
 local function fancy_diff(diff, ft, ctx)
-  require("snacks.picker.util.diff").render(ctx.preview:scratch(), ns, diff, { ft = ft })
+  require("snacks.picker.util.diff").render(ctx.preview:scratch(), ns, diff, {
+    ft = ft,
+    annotations = ctx.item.annotations or ctx.picker.opts.annotations,
+  })
   Snacks.util.wo(ctx.win, ctx.picker.opts.previewers.diff.wo or {})
 end
 
diff --git a/lua/snacks/picker/source/diff.lua b/lua/snacks/picker/source/diff.lua
index a5f6afd4..f8bf1120 100644
--- a/lua/snacks/picker/source/diff.lua
+++ b/lua/snacks/picker/source/diff.lua
@@ -4,6 +4,7 @@ local M = {}
 ---@field cmd? string optional since diff can be passed as string
 ---@field group? boolean Group hunks by file
 ---@field diff? string|number diff string or buffer number
+---@field annotations? snacks.diff.Annotation[]
 
 ---@class snacks.picker.diff.hunk.Pos
 ---@field line number
@@ -83,6 +84,7 @@ function M.diff(opts, ctx)
         file = file,
         cwd = cwd,
         rename = block.rename and block.rename.from or nil,
+        annotations = opts.annotations,
         block = block,
         pos = { line, 0 },
       })
diff --git a/lua/snacks/picker/source/gh.lua b/lua/snacks/picker/source/gh.lua
index e8eea087..572dc876 100644
--- a/lua/snacks/picker/source/gh.lua
+++ b/lua/snacks/picker/source/gh.lua
@@ -92,13 +92,6 @@ function M.get_actions(opts, ctx)
         return
       end
       item = Api.get({ type = opts.type or "pr", repo = opts.repo, number = opts.number })
-      local proc = Api.view(function(it)
-        item = it
-      end, item)
-
-      if proc then
-        proc:wait()
-      end
       if not item then
         Snacks.notify.error("snacks.picker.gh.get_actions: Failed to get item")
         return
@@ -148,14 +141,28 @@ function M.diff(opts, ctx)
   if opts.repo then
     vim.list_extend(args, { "--repo", opts.repo })
   end
-  return require("snacks.picker.source.diff").diff(
-    ctx:opts({
-      cmd = "gh",
-      args = args,
-      cwd = cwd,
-    }),
-    ctx
-  )
+  local Render = require("snacks.gh.render")
+  local Diff = require("snacks.picker.source.diff")
+  ---@async
+  return function(cb)
+    local item = Api.get({ type = "pr", repo = opts.repo, number = opts.pr })
+    local annotations ---@type snacks.diff.Annotation[]
+    ctx.async:schedule(function()
+      annotations = Render.annotations(item)
+    end)
+
+    Diff.diff(
+      ctx:opts({
+        cmd = "gh",
+        args = args,
+        cwd = cwd,
+        annotations = annotations,
+      }),
+      ctx
+    )(function(it)
+      cb(it)
+    end)
+  end
 end
 
 ---@param opts snacks.picker.gh.reactions.Config
@@ -220,7 +227,7 @@ function M.labels(opts, ctx)
       fields = { "labels" },
       args = { "repo", "view", opts.repo },
     })
-    local item = Api.get(opts)
+    local item = Api.get_cached(opts)
     assert(item, "Failed to get item for labels")
     local added = {} ---@type table<string, boolean>
     for _, label in ipairs(item.labels or {}) do
diff --git a/lua/snacks/picker/util/async.lua b/lua/snacks/picker/util/async.lua
index cdc67a74..d3ce4612 100644
--- a/lua/snacks/picker/util/async.lua
+++ b/lua/snacks/picker/util/async.lua
@@ -20,6 +20,9 @@ end
 
 ---@alias snacks.picker.AsyncEvent "done" | "error" | "yield" | "ok" | "abort"
 
+---@class snacks.picker.Waitable
+---@field wait async fun()
+
 ---@class snacks.picker.Async
 ---@field _co? thread
 ---@field _fn fun()
diff --git a/lua/snacks/picker/util/diff.lua b/lua/snacks/picker/util/diff.lua
index 680d4900..c73475ba 100644
--- a/lua/snacks/picker/util/diff.lua
+++ b/lua/snacks/picker/util/diff.lua
@@ -4,6 +4,29 @@ local M = {}
 ---@field max_hunk_lines? number only show last N lines of each hunk (used by GitHub PRs)
 ---@field hunk_header? boolean whether to show hunk header (default: true)
 ---@field ft? "diff" | "git"
+---@field annotations? snacks.diff.Annotation[]
+
+---@class snacks.diff.Annotation
+---@field file string
+---@field side "left" | "right"
+---@field left? number
+---@field right? number
+---@field line number
+---@field text snacks.picker.Highlight[][]
+
+---@class snacks.diff.ctx
+---@field diff snacks.picker.Diff
+---@field opts snacks.diff.Config
+---@field block? snacks.picker.diff.Block
+---@field hunk? snacks.picker.diff.Hunk
+local C = {}
+C.__index = C
+
+---@param ctx snacks.diff.ctx|{}
+---@return snacks.diff.ctx
+function C:extend(ctx)
+  return setmetatable(ctx, { __index = self })
+end
 
 ---@param ... string
 local function diff_linenr(...)
@@ -34,6 +57,8 @@ Snacks.util.set_hl({
   DiffConflictLineNr = diff_linenr("DiagnosticVirtualTextWarn"),
 }, { default = true, prefix = "Snacks" })
 
+local H = Snacks.picker.highlight
+
 ---@param diff string|string[]|snacks.picker.Diff
 function M.get_diff(diff)
   if type(diff) == "string" then
@@ -54,32 +79,33 @@ end
 function M.render(buf, ns, diff, opts)
   diff = M.get_diff(diff)
   local ret = M.format(diff, opts)
-  return Snacks.picker.highlight.render(buf, ns, ret)
+  return H.render(buf, ns, ret)
 end
 
 ---@param diff string|string[]|snacks.picker.Diff
 ---@param opts? snacks.diff.Config
 function M.format(diff, opts)
-  diff = M.get_diff(diff)
+  local ctx = C:extend({
+    diff = M.get_diff(diff),
+    opts = opts or {},
+  })
   local ret = {} ---@type snacks.picker.Highlight[][]
-  vim.list_extend(ret, M.format_header(diff, opts))
-  for _, block in ipairs(diff.blocks) do
-    vim.list_extend(ret, M.format_block(block, opts))
+  vim.list_extend(ret, M.format_header(ctx))
+  for _, block in ipairs(ctx.diff.blocks) do
+    vim.list_extend(ret, M.format_block(ctx:extend({ block = block })))
   end
   return ret
 end
 
----@param diff snacks.picker.Diff
----@param opts? snacks.diff.Config
-function M.format_header(diff, opts)
-  if #(diff.header or {}) == 0 then
+---@param ctx snacks.diff.ctx
+function M.format_header(ctx)
+  if #(ctx.diff.header or {}) == 0 then
     return {}
   end
   local popts = Snacks.picker.config.get({})
-  opts = opts or {}
   local ret = {} ---@type snacks.picker.Highlight[][]
   local msg = {} ---@type string[]
-  for _, line in ipairs(diff.header or {}) do
+  for _, line in ipairs(ctx.diff.header or {}) do
     local hash = line:match("^commit%s+(%S+)$")
     if hash then
       ret[#ret + 1] = {
@@ -110,36 +136,35 @@ function M.format_header(diff, opts)
     ret[#ret + 1] = Snacks.picker.format.commit_message({ msg = subject }, {})
   end
   if #msg > 0 then
-    ret[#ret + 1] = Snacks.picker.highlight.rule()
-    local virt_lines = Snacks.picker.highlight.get_virtual_lines(table.concat(msg, "\n"), { ft = "markdown" })
+    ret[#ret + 1] = H.rule()
+    local virt_lines = H.get_virtual_lines(table.concat(msg, "\n"), { ft = "markdown" })
     for _, vl in ipairs(virt_lines) do
       ret[#ret + 1] = vl
     end
   end
-  ret[#ret + 1] = Snacks.picker.highlight.rule()
+  ret[#ret + 1] = H.rule()
   return ret
 end
 
----@param block snacks.picker.diff.Block
----@param opts? snacks.diff.Config
-function M.format_block(block, opts)
+---@param ctx snacks.diff.ctx
+function M.format_block(ctx)
   local ret = {} ---@type snacks.picker.Highlight[][]
-  vim.list_extend(ret, M.format_block_header(block, opts))
-  for _, hunk in ipairs(block.hunks) do
-    local hunk_lines = M.format_hunk(hunk, block, opts)
-    if opts and opts.max_hunk_lines and #hunk_lines > opts.max_hunk_lines then
-      hunk_lines = vim.list_slice(hunk_lines, #hunk_lines - opts.max_hunk_lines + 1)
+  vim.list_extend(ret, M.format_block_header(ctx))
+  for _, hunk in ipairs(ctx.block.hunks) do
+    local hunk_lines = M.format_hunk(ctx:extend({ hunk = hunk }))
+    if ctx.opts and ctx.opts.max_hunk_lines and #hunk_lines > ctx.opts.max_hunk_lines then
+      hunk_lines = vim.list_slice(hunk_lines, #hunk_lines - ctx.opts.max_hunk_lines + 1)
     end
     vim.list_extend(ret, hunk_lines)
   end
   return ret
 end
 
----@param block snacks.picker.diff.Block
----@param opts? snacks.diff.Config
-function M.format_block_header(block, opts)
+---@param ctx snacks.diff.ctx
+function M.format_block_header(ctx)
+  local block = assert(ctx.block)
   local ret = {} ---@type snacks.picker.Highlight[][]
-  ret[#ret + 1] = Snacks.picker.highlight.add_eol({}, "SnacksDiffHeader")
+  ret[#ret + 1] = H.add_eol({}, "SnacksDiffHeader")
 
   local icon, icon_hl = Snacks.util.icon(block.file)
   local file = {} ---@type snacks.picker.Highlight[]
@@ -154,28 +179,24 @@ function M.format_block_header(block, opts)
   else
     file[#file + 1] = { block.file }
   end
-  Snacks.picker.highlight.insert_hl(file, "SnacksDiffHeader")
-  Snacks.picker.highlight.add_eol(file, "SnacksDiffHeader")
+  H.insert_hl(file, "SnacksDiffHeader")
+  H.add_eol(file, "SnacksDiffHeader")
   ret[#ret + 1] = file
 
-  ret[#ret + 1] = Snacks.picker.highlight.add_eol({}, "SnacksDiffHeader")
+  ret[#ret + 1] = H.add_eol({}, "SnacksDiffHeader")
   return ret
 end
 
----@param hunk snacks.picker.diff.Hunk
----@param block snacks.picker.diff.Block
----@param opts? snacks.diff.Config
-function M.format_hunk(hunk, block, opts)
-  opts = opts or {}
-  local a = Snacks.picker.util.align
-  local ret = {} ---@type snacks.picker.Highlight[][]
+---@param ctx snacks.diff.ctx
+function M.parse_hunk(ctx)
+  local block = assert(ctx.block)
+  local hunk = assert(ctx.hunk)
   local diff = vim.deepcopy(hunk.diff)
-  table.remove(diff, 1) -- remove hunk header line
-  while #diff > 0 and diff[#diff]:match("^%s*$") do
-    table.remove(diff) -- remove trailing empty lines
-  end
-
   local versions = {} ---@type snacks.picker.diff.hunk.Pos[]
+  local unmerged = #versions > 2
+  local lines, prefixes, conflict_markers = {}, {}, {} ---@type string[], string[], table<number, string>
+
+  -- build versions
   versions[#versions + 1] = hunk.left
   vim.list_extend(versions, hunk.parents or {})
   versions[#versions + 1] = hunk.right
@@ -183,9 +204,13 @@ function M.format_hunk(hunk, block, opts)
     versions[#versions + 1] = { line = hunk.line, count = 0 }
   end
 
-  local unmerged = #versions > 2
+  -- setup diff lines
+  table.remove(diff, 1) -- remove hunk header line
+  while #diff > 0 and diff[#diff]:match("^%s*$") do
+    table.remove(diff) -- remove trailing empty lines
+  end
 
-  local code, prefixes, conflict_markers = {}, {}, {} ---@type string[], string[], table<number, string>
+  -- parse diff lines
   for l, line in ipairs(diff) do
     prefixes[#prefixes + 1] = line:sub(1, #versions - 1)
     local code_line = line:sub(#versions)
@@ -193,37 +218,51 @@ function M.format_hunk(hunk, block, opts)
       conflict_markers[l] = code_line
       code_line = ""
     end
-    code[#code + 1] = code_line
+    lines[#lines + 1] = code_line
   end
 
-  table.insert(code, 1, hunk.context or "") -- add hunk context for syntax highlighting
-  local ft = vim.filetype.match({ filename = block.file, contents = code }) or ""
-  local virt_lines = Snacks.picker.highlight.get_virtual_lines(table.concat(code, "\n"), { ft = ft })
-  local context = table.remove(virt_lines, 1) -- remove hunk context virt lines
-  table.remove(code, 1)
+  -- generate virt lines
+  table.insert(lines, 1, hunk.context or "") -- add hunk context for syntax highlighting
+  local ft = vim.filetype.match({ filename = block.file, contents = lines }) or ""
+  local text = H.get_virtual_lines(table.concat(lines, "\n"), { ft = ft })
+  local context = table.remove(text, 1) -- remove hunk context virt lines
+  table.remove(lines, 1) -- remove hunk context code line
 
-  local lines = {} ---@type table<number, string[]>
+  ---@class snacks.diff.hunk.Parse
+  local ret = {
+    len = #diff, -- number of lines in hunk
+    versions = versions, -- positions of each version
+    lines = lines, -- code lines of hunk
+    text = text, -- virt lines of hunk
+    prefixes = prefixes, -- diff prefixes of hunk
+    conflict_markers = conflict_markers, -- conflict markers lines of hunk
+    context = context, -- virt lines of hunk context
+    unmerged = unmerged, -- whether hunk is unmerged
+  }
+  return ret
+end
+
+--- Build hunk line index for each version
+---@param parse snacks.diff.hunk.Parse
+function M.build_hunk_index(parse)
+  local versions = parse.versions
+  local index = {} ---@type table<number, number>[]|{max: number}
   local idx = {} ---@type number[]
   for p, pos in ipairs(versions) do
     idx[p] = idx[p] or ((pos.line or 1) - 1)
   end
   local max = 0
+  for l = 1, parse.len do
+    local prefix = parse.prefixes[l]
+    index[l] = {}
 
-  for l in ipairs(diff) do
-    local prefix = prefixes[l]
-    local line = {} ---@type string[]
-    lines[l] = line
-    for i = 1, #versions do
-      line[i] = ""
-    end
-
-    if not conflict_markers[l] then
+    if not parse.conflict_markers[l] then
       -- Increment parent versions
       for i = 1, #versions - 1 do
         local char = prefix:sub(i, i)
         if char == " " or char == "-" then
           idx[i] = idx[i] + 1
-          line[i] = tostring(idx[i])
+          index[l][i] = idx[i]
           max = math.max(max, #tostring(idx[i]))
         end
       end
@@ -240,40 +279,67 @@ function M.format_hunk(hunk, block, opts)
     end
     if has_working then
       idx[#idx] = idx[#idx] + 1
-      line[#idx] = tostring(idx[#idx])
+      index[l][#idx] = idx[#idx]
       max = math.max(max, #tostring(idx[#idx]))
     end
   end
+  index.max = max
+  return index
+end
 
-  if opts.hunk_header ~= false then
-    local header = {} ---@type snacks.picker.Highlight[]
-    header[#header + 1] = { "  " }
-    header[#header + 1] = { " ", "Special" }
-    header[#header + 1] = { " " }
-    Snacks.picker.highlight.extend(header, context)
-    local context_width = Snacks.picker.highlight.offset(context)
-    ret[#ret + 1] = {
-      { string.rep("─", context_width + 7) .. "┐", "FloatBorder" },
-    }
-    header[#header + 1] = { "  │", "FloatBorder" }
-    ret[#ret + 1] = header
-    ret[#ret + 1] = {
-      { string.rep("─", context_width + 7) .. "┘", "FloatBorder" },
-    }
+---@param parse snacks.diff.hunk.Parse
+function M.format_hunk_header(parse)
+  local ret = {} ---@type snacks.picker.Highlight[][]
+  local header = {} ---@type snacks.picker.Highlight[]
+  header[#header + 1] = { "  " }
+  header[#header + 1] = { " ", "Special" }
+  header[#header + 1] = { " " }
+  H.extend(header, parse.context)
+  local context_width = H.offset(parse.context)
+  ret[#ret + 1] = {
+    { string.rep("─", context_width + 7) .. "┐", "FloatBorder" },
+  }
+  header[#header + 1] = { "  │", "FloatBorder" }
+  ret[#ret + 1] = header
+  ret[#ret + 1] = {
+    { string.rep("─", context_width + 7) .. "┘", "FloatBorder" },
+  }
+  return ret
+end
+
+---@param ctx snacks.diff.ctx
+function M.format_hunk(ctx)
+  local block = assert(ctx.block)
+  local align = Snacks.picker.util.align
+  local ret = {} ---@type snacks.picker.Highlight[][]
+
+  local parse = M.parse_hunk(ctx)
+
+  local annotations = {} ---@type table<string, snacks.picker.Highlight[][]>
+  for _, annotation in ipairs(ctx.opts.annotations or {}) do
+    if annotation.file == block.file then
+      annotations[("%s:%d"):format(annotation.side, annotation.line)] = annotation.text
+    end
+  end
+
+  local index = M.build_hunk_index(parse)
+
+  if ctx.opts.hunk_header ~= false then
+    vim.list_extend(ret, M.format_hunk_header(parse))
   end
 
   local in_conflict = false
-  for l = 1, #diff do
-    local have_left, have_right = lines[l][1] ~= "", lines[l][#versions] ~= ""
-    local hl = (conflict_markers[l] and "SnacksDiffConflict")
+  for l = 1, parse.len do
+    local have_left, have_right = index[l][1] ~= nil, index[l][#parse.versions] ~= nil
+    local hl = (parse.conflict_markers[l] and "SnacksDiffConflict")
       or (have_right and not have_left and "SnacksDiffAdd")
       or (have_left and not have_right and "SnacksDiffDelete")
       or "SnacksDiffContext"
 
-    local prefix = prefixes[l]
-    if unmerged then
+    local prefix = parse.prefixes[l]
+    if parse.unmerged then
       local p = "  "
-      local marker = conflict_markers[l] or ""
+      local marker = parse.conflict_markers[l] or ""
       marker = marker:match("^%s*(%S+)") or ""
       if marker == "<<<<<<<" then
         in_conflict = true
@@ -286,14 +352,14 @@ function M.format_hunk(hunk, block, opts)
       elseif in_conflict then
         p = "│ "
       end
-      prefix = a(p, 2) .. prefix
+      prefix = align(p, 2) .. prefix
     end
 
     local line = {} ---@type snacks.picker.Highlight[]
 
     local line_nr = {} ---@type string[]
-    for i = 1, #versions do
-      line_nr[i] = a(lines[l][i], max, { align = i == #versions and "right" or "left" })
+    for i = 1, #parse.versions do
+      line_nr[i] = align(tostring(index[l][i] or ""), index.max, { align = i == #parse.versions and "right" or "left" })
     end
     local line_col = " " .. table.concat(line_nr, "  ") .. " "
     local prefix_col = " " .. prefix .. " "
@@ -316,10 +382,10 @@ function M.format_hunk(hunk, block, opts)
     }
 
     -- empty prefix overlay that will be used for wrapped lines
-    local ws = (conflict_markers[l] or code[l]):match("^(%s*)") -- add ws for breakindent
+    local ws = (parse.conflict_markers[l] or parse.lines[l]):match("^(%s*)") -- add ws for breakindent
     line[#line + 1] = {
       col = #line_col,
-      virt_text = { { a(prefix_col:gsub("[%-%+]", " "), #ws + #prefix_col), hl } },
+      virt_text = { { align(prefix_col:gsub("[%-%+]", " "), #ws + #prefix_col), hl } },
       virt_text_pos = "overlay",
       hl_mode = "replace",
       virt_text_repeat_linebreak = true,
@@ -333,18 +399,109 @@ function M.format_hunk(hunk, block, opts)
       hl_mode = "replace",
     }
 
-    local vl = Snacks.picker.highlight.indent({}, #line_col + #prefix_col)
-    if conflict_markers[l] then
-      vl[#vl + 1] = { conflict_markers[l], hl }
-    else
-      vim.list_extend(vl, virt_lines[l] or {})
-    end
-    Snacks.picker.highlight.insert_hl(vl, hl)
-    Snacks.picker.highlight.extend(line, vl)
-    Snacks.picker.highlight.add_eol(line, hl)
     ret[#ret + 1] = line
+
+    local annot_left = "left:" .. (index[l][1] or "")
+    local annot_right = "right:" .. (index[l][#parse.versions] or "")
+    local ann = annotations[annot_left] or annotations[annot_right]
+    if ann then
+      vim.list_extend(
+        ret,
+        M.format_annotation(ann, {
+          indent = { line[1] },
+          indent_width = #line_col,
+          hl = hl,
+        })
+      )
+    end
+
+    local vl = H.indent({}, #line_col + #prefix_col)
+    if parse.conflict_markers[l] then
+      vl[#vl + 1] = { parse.conflict_markers[l], hl }
+    else
+      vim.list_extend(vl, parse.text[l] or {})
+    end
+    H.insert_hl(vl, hl)
+    H.extend(line, vl)
+    H.add_eol(line, hl)
   end
   return ret
 end
 
+---@param annotation snacks.picker.Highlight[][]
+---@param ctx {indent: snacks.picker.Highlight[][], indent_width: number, hl: string}
+function M.format_annotation(annotation, ctx)
+  local ret = {} ---@type snacks.picker.Highlight[][]
+  local box, width = M.format_box(annotation)
+
+  local empty = vim.deepcopy(ctx.indent) ---@type snacks.picker.Highlight[]
+  vim.list_extend(empty, H.indent({}, ctx.indent_width + 2, ctx.hl))
+  H.add_eol(empty, ctx.hl)
+
+  ret[#ret + 1] = vim.deepcopy(empty)
+  for _, line in ipairs(box) do
+    for _, chunk in ipairs(line) do
+      if chunk.virt_text_win_col then
+        chunk.virt_text_win_col = chunk.virt_text_win_col + ctx.indent_width + 2
+      end
+    end
+    local al = vim.deepcopy(ctx.indent)
+    local vl = H.indent({}, ctx.indent_width + 2, ctx.hl)
+    H.extend(al, vl)
+    H.extend(al, vim.deepcopy(line))
+    H.add_eol(al, ctx.hl, width + ctx.indent_width + 6)
+    ret[#ret + 1] = al
+  end
+  ret[#ret + 1] = vim.deepcopy(empty)
+  return ret
+end
+
+---@param lines snacks.picker.Highlight[][]
+---@param border_hl? string
+function M.format_box(lines, border_hl)
+  border_hl = border_hl or "FloatBorder"
+  local ret = {} ---@type snacks.picker.Highlight[][]
+  local width = 0
+  for _, line in ipairs(lines) do
+    width = math.max(width, H.offset(line, { char_idx = true }))
+  end
+  width = math.max(width, 50) --[[@as number]]
+
+  ---@param text snacks.picker.Highlight[]
+  ---@param col? number
+  local function vt(text, col)
+    ---@type snacks.picker.Highlight
+    return {
+      col = 0,
+      virt_text_pos = "overlay",
+      virt_text_win_col = col,
+      virt_text = text,
+    }
+  end
+
+  ret[#ret + 1] = {
+    vt({
+      { "┌", border_hl },
+      { string.rep("─", width + 2), border_hl },
+      { "┐", border_hl },
+    }),
+  }
+  for _, line in ipairs(lines) do
+    ret[#ret + 1] = {
+      { "│", border_hl, virtual = true },
+      { " " },
+    }
+    H.extend(ret[#ret], vim.deepcopy(line))
+    table.insert(ret[#ret], vt({ { "│", border_hl } }, width + 3))
+  end
+  ret[#ret + 1] = {
+    vt({
+      { "└", border_hl },
+      { string.rep("─", width + 2), border_hl },
+      { "┘", border_hl },
+    }),
+  }
+  return ret, width
+end
+
 return M
diff --git a/lua/snacks/picker/util/highlight.lua b/lua/snacks/picker/util/highlight.lua
index 625ea6d7..07894d45 100644
--- a/lua/snacks/picker/util/highlight.lua
+++ b/lua/snacks/picker/util/highlight.lua
@@ -378,12 +378,14 @@ end
 
 ---@param line snacks.picker.Highlight[]
 ---@param hl_group string
-function M.add_eol(line, hl_group)
+---@param offset? number
+function M.add_eol(line, hl_group, offset)
   line[#line + 1] = {
     col = M.offset(line),
     virt_text = { { string.rep(" ", 1000), hl_group } },
     virt_text_pos = "overlay",
     hl_mode = "replace",
+    virt_text_win_col = offset,
     virt_text_repeat_linebreak = true,
   }
   return line
diff --git a/lua/snacks/util/spawn.lua b/lua/snacks/util/spawn.lua
index ff1026ca..d7aadc01 100644
--- a/lua/snacks/util/spawn.lua
+++ b/lua/snacks/util/spawn.lua
@@ -20,7 +20,7 @@ local uv = vim.uv or vim.loop
 ---@field cmd? nil
 ---@field on_exit? fun(procs: snacks.spawn.Proc[], err: boolean)
 
----@class snacks.spawn.Proc
+---@class snacks.spawn.Proc: snacks.picker.Waitable
 ---@field opts snacks.spawn.Config
 ---@field handle? uv.uv_process_t
 ---@field stdout uv.uv_pipe_t