refactor(ui): track rendered blocks in RenderState

Add a RenderedBlock type and maintain block metadata maps on the RenderState
(instance-level tables: _blocks, _message_blocks, _part_blocks). Introduce
APIs to set/get/remove blocks and to associate blocks with messages and parts:
set_block, set_message_block, set_part_block, get_block,
get_message_block, get_part_block, and remove_block. Ensure block
line ranges stay in sync when messages/parts are updated, shifted, or removed.
Update reset, shift and removal logic accordingly and add unit tests for block
metadata handling.

Author: sudo-tee

lua/opencode/ui/render_state.lua

@@ -11,9 +11,22 @@
 ---@field actions table[] Actions associated with this part
 ---@field has_extmarks boolean? Whether the part currently has extmarks applied
 
+---@class RenderedBlock
+---@field block RenderBlock Direct reference to a cached/rendered block
+---@field key string Stable render block key
+---@field kind string Block kind (message/part/overlay)
+---@field message_id string? Parent message ID when applicable
+---@field part_id string? Parent part ID when applicable
+---@field line_start integer? Line where the block starts when mounted
+---@field line_end integer? Line where the block ends when mounted
+---@field line_count integer? Cached block line count
+
 ---@class RenderState
 ---@field _messages table<string, RenderedMessage> Message ID -> rendered message
 ---@field _parts table<string, RenderedPart> Part ID -> rendered part
+---@field _blocks table<string, RenderedBlock> Block key -> rendered block
+---@field _message_blocks table<string, string> Message ID -> primary block key
+---@field _part_blocks table<string, string> Part ID -> primary block key
 ---@field _part_ranges {[1]: integer, [2]: integer, [3]: string}[] Sorted [line_start, line_end, part_id] for binary search
 ---@field _message_ranges {[1]: integer, [2]: integer, [3]: string}[] Sorted [line_start, line_end, message_id] for binary search
 ---@field _ranges_valid boolean Whether range arrays are sorted and up-to-date
@@ -32,6 +45,9 @@ end
 function RenderState:reset()
   self._messages = {}
   self._parts = {}
+  self._blocks = {}
+  self._message_blocks = {}
+  self._part_blocks = {}
   self._part_ranges = {}
   self._message_ranges = {}
   self._ranges_valid = false
@@ -44,6 +60,92 @@ function RenderState:reset()
   self._snapshot_id_index = {} -- snapshot_id -> OpencodeMessagePart
 end
 
+---@param mapping table<string, string>
+---@param id string?
+---@param key string?
+local function clear_block_reference(mapping, id, key)
+  if not id or not key then
+    return
+  end
+
+  if mapping[id] == key then
+    mapping[id] = nil
+  end
+end
+
+---@param block RenderBlock?
+---@return string?
+local function get_block_key(block)
+  if not block or not block.key or block.key == '' then
+    return nil
+  end
+  return block.key
+end
+
+---@param message string|OpencodeMessage|nil
+---@return string?
+local function resolve_message_id(message)
+  if type(message) == 'string' then
+    return message ~= '' and message or nil
+  end
+
+  return message and message.info and message.info.id or nil
+end
+
+---@param part string|OpencodeMessagePart|nil
+---@return string?
+local function resolve_part_id(part)
+  if type(part) == 'string' then
+    return part ~= '' and part or nil
+  end
+
+  return part and part.id or nil
+end
+
+---@param block_key string
+---@param line_start integer?
+---@param line_end integer?
+---@param message_id string?
+---@param part_id string?
+---@return RenderedBlock?
+function RenderState:_upsert_block_entry(block_key, line_start, line_end, message_id, part_id)
+  local entry = self._blocks[block_key]
+  if not entry then
+    return nil
+  end
+
+  local previous_message_id = entry.message_id
+  local previous_part_id = entry.part_id
+
+  entry.message_id = message_id or entry.message_id
+  entry.part_id = part_id or entry.part_id
+  entry.line_count = entry.block and entry.block.line_count or entry.line_count
+  entry.kind = entry.block and entry.block.kind or entry.kind
+
+  if line_start ~= nil then
+    entry.line_start = line_start
+  end
+  if line_end ~= nil then
+    entry.line_end = line_end
+  end
+
+  if previous_message_id ~= entry.message_id then
+    clear_block_reference(self._message_blocks, previous_message_id, block_key)
+  end
+  if previous_part_id ~= entry.part_id then
+    clear_block_reference(self._part_blocks, previous_part_id, block_key)
+  end
+
+  if entry.kind == 'message' and entry.message_id then
+    self._message_blocks[entry.message_id] = block_key
+  end
+  if entry.kind == 'part' and entry.part_id then
+    self._part_blocks[entry.part_id] = block_key
+  end
+
+  return entry
+end
+
 function RenderState:_recompute_max_line_end()
   local max_line_end = 0
 
@@ -59,6 +161,12 @@ function RenderState:_recompute_max_line_end()
     end
   end
 
+  for _, block_data in pairs(self._blocks) do
+    if block_data.line_end and block_data.line_end > max_line_end then
+      max_line_end = block_data.line_end
+    end
+  end
+
   self._max_line_end = max_line_end
   self._max_line_end_valid = true
   return max_line_end
@@ -205,6 +313,44 @@ function RenderState:get_message(message_id)
   return self._messages[message_id]
 end
 
+---@param block_key string
+---@return RenderedBlock?
+function RenderState:get_block(block_key)
+  return self._blocks[block_key]
+end
+
+---@param message string|OpencodeMessage
+---@return RenderedBlock?
+function RenderState:get_message_block(message)
+  local message_id = resolve_message_id(message)
+  if not message_id then
+    return nil
+  end
+
+  local block_key = self._message_blocks[message_id]
+  if not block_key then
+    block_key = 'message:' .. message_id
+  end
+
+  return self._blocks[block_key]
+end
+
+---@param part string|OpencodeMessagePart
+---@return RenderedBlock?
+function RenderState:get_part_block(part)
+  local part_id = resolve_part_id(part)
+  if not part_id then
+    return nil
+  end
+
+  local block_key = self._part_blocks[part_id]
+  if not block_key then
+    block_key = 'part:' .. part_id
+  end
+
+  return self._blocks[block_key]
+end
+
 ---@param line integer 1-indexed
 ---@return RenderedMessage?
 function RenderState:get_message_at_line(line)
@@ -347,6 +493,58 @@ function RenderState:set_message(message, line_start, line_end)
       self._max_line_end = line_end
     end
   end
+
+  local block_key = self._message_blocks[message_id] or ('message:' .. message_id)
+  if self._blocks[block_key] then
+    self:_upsert_block_entry(block_key, line_start, line_end, message_id, nil)
+  end
+end
+
+---@param block RenderBlock
+---@param line_start integer?
+---@param line_end integer?
+---@return RenderedBlock?
+function RenderState:set_block(block, line_start, line_end)
+  local block_key = get_block_key(block)
+  if not block_key then
+    return nil
+  end
+
+  local existing = self._blocks[block_key]
+  if not existing then
+    self._blocks[block_key] = {
+      block = block,
+      key = block_key,
+      kind = block.kind,
+      message_id = block.message_id,
+      part_id = block.part_id,
+      line_start = line_start,
+      line_end = line_end,
+      line_count = block.line_count,
+    }
+  else
+    existing.block = block
+    existing.key = block_key
+    existing.kind = block.kind
+    existing.line_count = block.line_count
+  end
+
+  return self:_upsert_block_entry(block_key, line_start, line_end, block.message_id, block.part_id)
+end
+
+---@param message string|OpencodeMessage
+---@param block RenderBlock
+---@param line_start integer?
+---@param line_end integer?
+---@return RenderedBlock?
+function RenderState:set_message_block(message, block, line_start, line_end)
+  local message_id = resolve_message_id(message)
+  local entry = self:set_block(block, line_start, line_end)
+  if not entry then
+    return nil
+  end
+
+  return self:_upsert_block_entry(entry.key, line_start, line_end, message_id, entry.part_id)
 end
 
 ---@param part OpencodeMessagePart
@@ -394,6 +592,26 @@ function RenderState:set_part(part, line_start, line_end)
   end
 
   self:_index_task_part_child_session(part_id, part)
+
+  local block_key = self._part_blocks[part_id] or ('part:' .. part_id)
+  if self._blocks[block_key] then
+    self:_upsert_block_entry(block_key, line_start, line_end, message_id, part_id)
+  end
+end
+
+---@param part string|OpencodeMessagePart
+---@param block RenderBlock
+---@param line_start integer?
+---@param line_end integer?
+---@return RenderedBlock?
+function RenderState:set_part_block(part, block, line_start, line_end)
+  local part_id = resolve_part_id(part)
+  local entry = self:set_block(block, line_start, line_end)
+  if not entry then
+    return nil
+  end
+
+  return self:_upsert_block_entry(entry.key, line_start, line_end, entry.message_id, part_id)
 end
 
 ---@param part_id string
@@ -419,6 +637,11 @@ function RenderState:update_part_lines(part_id, new_line_start, new_line_end)
   part_data.line_end = new_line_end
   self._ranges_valid = false
 
+  local block_key = self._part_blocks[part_id] or ('part:' .. part_id)
+  if self._blocks[block_key] then
+    self:_upsert_block_entry(block_key, new_line_start, new_line_end, part_data.message_id, part_id)
+  end
+
   if self._max_line_end_valid then
     if old_line_end == self._max_line_end and new_line_end < old_line_end then
       self._max_line_end_valid = false
@@ -462,6 +685,8 @@ function RenderState:remove_part(part_id)
     return false
   end
 
+  self:remove_block(self._part_blocks[part_id] or ('part:' .. part_id))
+
   if part_data.part and part_data.part.type == 'patch' and part_data.part.hash then
     self._snapshot_id_index[part_data.part.hash] = nil
   end
@@ -486,6 +711,20 @@ function RenderState:remove_part(part_id)
   return true
 end
 
+---@param block_key string
+---@return boolean
+function RenderState:remove_block(block_key)
+  local entry = self._blocks[block_key]
+  if not entry then
+    return false
+  end
+
+  clear_block_reference(self._message_blocks, entry.message_id, block_key)
+  clear_block_reference(self._part_blocks, entry.part_id, block_key)
+  self._blocks[block_key] = nil
+  return true
+end
+
 ---@param message_id string
 ---@return boolean
 function RenderState:remove_message(message_id)
@@ -494,6 +733,8 @@ function RenderState:remove_message(message_id)
     return false
   end
 
+  self:remove_block(self._message_blocks[message_id] or ('message:' .. message_id))
+
   local line_count = msg_data.line_end - msg_data.line_start + 1
   local shift_from = msg_data.line_end + 1
 
@@ -577,6 +818,14 @@ function RenderState:shift_all(from_line, delta)
     end
   end
 
+  for _, block_data in pairs(self._blocks) do
+    if block_data.line_start and block_data.line_end and block_data.line_start >= from_line then
+      block_data.line_start = block_data.line_start + delta
+      block_data.line_end = block_data.line_end + delta
+      shifted = true
+    end
+  end
+
   if shifted then
     self._ranges_valid = false
     if self._max_line_end_valid then