refactor(render_state): clean up set/update apis

Don't need to pass id, object, and possibly other things to
set_part/update_part_data/set_message, just pass a part or message.

Author: cameronr

lua/opencode/ui/formatter.lua

@@ -109,6 +109,7 @@ end
 ---Format the revert callout with statistics
 ---@param session_data OpencodeMessage[] All messages in the session
 ---@param start_idx number Index of the message where revert occurred
+---@return Output output object representing the lines, extmarks, and actions
 function M._format_revert_message(session_data, start_idx)
   local output = Output.new()
   local stats = M._calculate_revert_stats(session_data, start_idx, state.active_session.revert)

lua/opencode/ui/render_state.lua

@@ -129,20 +129,24 @@ function RenderState:get_actions_at_line(line)
 end
 
 ---Set or update message render data
----@param message_id string Message ID
----@param message_ref OpencodeMessage Direct reference to message
+---@param message OpencodeMessage Direct reference to message
 ---@param line_start integer? Line where message header starts
 ---@param line_end integer? Line where message header ends
-function RenderState:set_message(message_id, message_ref, line_start, line_end)
+function RenderState:set_message(message, line_start, line_end)
+  if not message or not message.info or not message.info.id then
+    return
+  end
+  local message_id = message.info.id
+
   if not self._messages[message_id] then
     self._messages[message_id] = {
-      message = message_ref,
+      message = message,
       line_start = line_start,
       line_end = line_end,
     }
   else
     local msg_data = self._messages[message_id]
-    msg_data.message = message_ref
+    msg_data.message = message
     if line_start then
       msg_data.line_start = line_start
     end
@@ -157,12 +161,16 @@ function RenderState:set_message(message_id, message_ref, line_start, line_end)
 end
 
 ---Set or update part render data
----@param part_id string Part ID
----@param part MessagePart Direct reference to part
----@param message_id string Parent message ID
+---@param part MessagePart Direct reference to part (must include id/messageID)
 ---@param line_start integer? Line where part starts
 ---@param line_end integer? Line where part ends
-function RenderState:set_part(part_id, part, message_id, line_start, line_end)
+function RenderState:set_part(part, line_start, line_end)
+  if not part or not part.id then
+    return
+  end
+  local part_id = part.id
+  local message_id = part.messageID
+
   if not self._parts[part_id] then
     self._parts[part_id] = {
       part = part,
@@ -174,7 +182,9 @@ function RenderState:set_part(part_id, part, message_id, line_start, line_end)
   else
     local render_part = self._parts[part_id]
     render_part.part = part
-    render_part.message_id = message_id
+    if message_id then
+      render_part.message_id = message_id
+    end
     if line_start then
       render_part.line_start = line_start
     end
@@ -218,10 +228,12 @@ function RenderState:update_part_lines(part_id, new_line_start, new_line_end)
 end
 
 ---Update part data reference
----@param part_id string Part ID
----@param part_ref MessagePart New part reference
-function RenderState:update_part_data(part_id, part_ref)
-  local part_data = self._parts[part_id]
+---@param part_ref MessagePart New part reference (must include id)
+function RenderState:update_part_data(part_ref)
+  if not part_ref or not part_ref.id then
+    return
+  end
+  local part_data = self._parts[part_ref.id]
   if not part_data then
     return
   end

lua/opencode/ui/renderer.lua

@@ -106,7 +106,7 @@ function M.render_full_session()
 
   if config.debug.enabled then
     -- TODO: I want to track full renders for now, remove at some point
-    -- vim.notify('rendering full session\n' .. debug.traceback(), vim.log.levels.WARN)
+    vim.notify('rendering full session\n' .. debug.traceback(), vim.log.levels.WARN)
   end
 
   fetch_session():and_then(M._render_full_session_data)
@@ -124,11 +124,9 @@ function M._render_full_session_data(session_data)
       revert_index = i
     end
 
-    -- only pass in the info so, the parts will be processed as part of the loop
-    -- TODO: remove part processing code in formatter
     M.on_message_updated({ info = msg.info }, revert_index)
 
-    for j, part in ipairs(msg.parts or {}) do
+    for _, part in ipairs(msg.parts or {}) do
       M.on_part_updated({ part = part }, revert_index)
     end
   end
@@ -235,7 +233,7 @@ function M._insert_part_to_buffer(part_id, formatted_data)
     return false
   end
 
-  M._render_state:set_part(part_id, cached.part, cached.message_id, range.line_start, range.line_end)
+  M._render_state:set_part(cached.part, range.line_start, range.line_end)
   return true
 end
 
@@ -328,7 +326,7 @@ function M._replace_message_in_buffer(message_id, formatted_data)
   local old_line_end = cached.line_end
   local new_line_end = cached.line_start + new_line_count - 1
 
-  M._render_state:set_message(message_id, cached.message, cached.line_start, new_line_end)
+  M._render_state:set_message(cached.message, cached.line_start, new_line_end)
 
   local delta = new_line_end - old_line_end
   if delta ~= 0 then
@@ -360,7 +358,7 @@ function M.on_message_updated(message, revert_index)
     if not found_msg then
       table.insert(state.messages, message)
     end
-    M._render_state:set_message(message.info.id, message, 0, 0)
+    M._render_state:set_message(message, 0, 0)
     return
   end
 
@@ -382,7 +380,7 @@ function M.on_message_updated(message, revert_index)
     local range = M._write_formatted_data(header_data)
 
     if range then
-      M._render_state:set_message(message.info.id, message, range.line_start, range.line_end)
+      M._render_state:set_message(message, range.line_start, range.line_end)
     end
 
     state.current_message = message
@@ -444,9 +442,9 @@ function M.on_part_updated(properties, revert_index)
   end
 
   if is_new_part then
-    M._render_state:set_part(part.id, part, part.messageID)
+    M._render_state:set_part(part)
   else
-    M._render_state:update_part_data(part.id, part)
+    M._render_state:update_part_data(part)
   end
 
   local formatted = formatter.format_part(part, message.info.role)