feat(renderer): make undo/redo work

Author: sudo-tee

lua/opencode/api.lua

@@ -582,17 +582,13 @@ function M.undo()
     return
   end
 
-  -- ui.render_output(true)
   state.api_client
     :revert_message(state.active_session.id, {
       messageID = last_user_message.info.id,
     })
     :and_then(function(response)
-      state.active_session.revert = response.revert
       vim.schedule(function()
-        -- FIXME: shouldn't require a full re-render
-        vim.notify('Last message undone successfully', vim.log.levels.INFO)
-        require('opencode.ui.renderer').render_full_session()
+        vim.cmd('checktime')
       end)
     end)
     :catch(function(err)
@@ -612,11 +608,8 @@ function M.redo()
   state.api_client
     :unrevert_messages(state.active_session.id)
     :and_then(function(response)
-      state.active_session.revert = response.revert
       vim.schedule(function()
-        -- FIXME: shouldn't require a full re-render
-        vim.notify('Last message reverted successfully', vim.log.levels.INFO)
-        require('opencode.ui.renderer').render_full_session()
+        vim.cmd('checktime')
       end)
     end)
     :catch(function(err)

lua/opencode/session.lua

@@ -150,7 +150,7 @@ end
 
 ---Get messages for a session
 ---@param session Session
----@return Promise<{info: MessageInfo, parts: MessagePart[]}[]?>
+---@return Promise<OpencodeMessage[]?>
 function M.get_messages(session)
   local state = require('opencode.state')
   if not session then

lua/opencode/types.lua

@@ -423,3 +423,8 @@
 ---@field desc string|nil Description of the command
 ---@field fn fun(args:string[]):nil Function to execute the command
 ---@field args boolean Whether the command accepts arguments
+
+---@class OpencodeRevertSummary
+---@field messages number Number of messages reverted
+---@field tool_calls number Number of tool calls reverted
+---@field files table<string, {additions: number, deletions: number}> Summary of file changes reverted

lua/opencode/ui/formatter.lua

@@ -107,12 +107,15 @@ function M._calculate_revert_stats(messages, revert_index, revert_info)
 end
 
 ---Format the revert callout with statistics
----@param output Output Output object to write to
----@param stats {messages: number, tool_calls: number, files: table<string, {additions: number, deletions: number}>}
-function M._format_revert_message(output, stats)
+---@param session_data OpencodeMessage[] All messages in the session
+---@param start_idx number Index of the message where revert occurred
+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)
   local message_text = stats.messages == 1 and 'message' or 'messages'
   local tool_text = stats.tool_calls == 1 and 'tool call' or 'tool calls'
 
+  output:add_lines(M.separator)
   output:add_line(
     string.format('> %d %s reverted, %d %s reverted', stats.messages, message_text, stats.tool_calls, tool_text)
   )
@@ -146,6 +149,7 @@ function M._format_revert_message(output, stats)
       end
     end
   end
+  return output
 end
 
 ---@param output Output Output object to write to

lua/opencode/ui/renderer.lua

@@ -106,39 +106,36 @@ 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)
 end
 
 function M._render_full_session_data(session_data)
   M.reset()
+  local revert_index = nil
 
   for i, msg in ipairs(session_data) do
     -- output:add_lines(M.separator)
     -- state.current_message = msg
 
     if state.active_session.revert and state.active_session.revert.messageID == msg.info.id then
-      ---@type {messages: number, tool_calls: number, files: table<string, {additions: number, deletions: number}>}
-      local revert_stats = M._calculate_revert_stats(state.messages, i, state.active_session.revert)
-      local output = require('opencode.ui.output'):new()
-      formatter._format_revert_message(output, revert_stats)
-      M.render_output(output)
-
-      -- FIXME: how does reverting work? why is it breaking out of the message reading loop?
-      break
+      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 })
+    M.on_message_updated({ info = msg.info }, revert_index)
 
     for j, part in ipairs(msg.parts or {}) do
-      M.on_part_updated({ part = part })
+      M.on_part_updated({ part = part }, revert_index)
     end
   end
 
+  if revert_index then
+    M._write_formatted_data(formatter._format_revert_message(state.messages, revert_index))
+  end
   M._scroll_to_bottom()
 end
 
@@ -344,7 +341,8 @@ end
 ---Event handler for message.updated events
 ---Creates new message or updates existing message info
 ---@param message {info: MessageInfo} Event properties
-function M.on_message_updated(message)
+---@param revert_index? integer Revert index in session, if applicable
+function M.on_message_updated(message, revert_index)
   ---@type OpencodeMessage
   if not message or not message.info or not message.info.id or not message.info.sessionID then
     return
@@ -358,14 +356,22 @@ function M.on_message_updated(message)
   local rendered_message = M._render_state:get_message(message.info.id)
   local found_msg = rendered_message and rendered_message.message
 
+  if revert_index then
+    if not found_msg then
+      table.insert(state.messages, message)
+    end
+    M._render_state:set_message(message.info.id, message, 0, 0)
+    return
+  end
+
   if found_msg then
     -- see if an error was added (or removed). have to check before we set
     -- found_msg.info = message.info below
     local rerender_message = not vim.deep_equal(found_msg.info.error, message.info.error)
 
     found_msg.info = message.info
 
-    if rerender_message then
+    if rerender_message and not revert_index then
       local header_data = formatter.format_message_header(found_msg)
       M._replace_message_in_buffer(message.info.id, header_data)
     end
@@ -380,7 +386,6 @@ function M.on_message_updated(message)
     end
 
     state.current_message = message
-
     if message.info.role == 'user' then
       state.last_user_message = message
     end
@@ -394,7 +399,8 @@ end
 ---Event handler for message.part.updated events
 ---Inserts new parts or replaces existing parts in buffer
 ---@param properties {part: MessagePart} Event properties
-function M.on_part_updated(properties)
+---@param revert_index? integer Revert index in session, if applicable
+function M.on_part_updated(properties, revert_index)
   if not properties or not properties.part then
     return
   end
@@ -445,6 +451,10 @@ function M.on_part_updated(properties)
 
   local formatted = formatter.format_part(part, message.info.role)
 
+  if revert_index and is_new_part then
+    return
+  end
+
   if is_new_part then
     M._insert_part_to_buffer(part.id, formatted)
   else
@@ -533,6 +543,10 @@ end
 ---@param properties {info: Session}
 function M.on_session_updated(properties)
   require('opencode.ui.topbar').render()
+  if not vim.deep_equal(state.active_session.revert, properties.info.revert) then
+    state.active_session.revert = properties.info.revert
+    M._render_full_session_data(state.messages)
+  end
 end
 
 ---Event handler for session.error events

lua/opencode/ui/session_picker.lua

@@ -18,6 +18,7 @@ local function format_session(session)
     table.insert(parts, modified)
   end
 
+  table.insert(parts, 'ID: ' .. (session.id or 'N/A'))
   return table.concat(parts, ' ~ ')
 end
 

lua/opencode/ui/ui.lua

@@ -199,7 +199,7 @@ function M.select_session(sessions, cb)
     vim.ui.select(sessions, {
       prompt = '',
       format_item = function(session)
-        local parts = {}
+        local parts = { { session.id } }
 
         if session.description then
           table.insert(parts, session.description)

tests/data/revert.expected.json

@@ -0,0 +1 @@
+{"actions":[{"range":{"to":55,"from":55},"text":"[R]evert file","display_line":55,"key":"R","args":["c410b2b4024de020aea223c5248eec89216de53f"],"type":"diff_revert_selected_file"},{"range":{"to":55,"from":55},"text":"Revert [A]ll","display_line":55,"key":"A","args":["c410b2b4024de020aea223c5248eec89216de53f"],"type":"diff_revert_all"},{"range":{"to":55,"from":55},"text":"[D]iff","display_line":55,"key":"D","args":["c410b2b4024de020aea223c5248eec89216de53f"],"type":"diff_open"}],"extmarks":[[1,2,0,{"virt_text_repeat_linebreak":false,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["▌󰭻 ","OpencodeMessageRoleUser"],[" "],["USER","OpencodeMessageRoleUser"],["","OpencodeHint"],[" (2025-10-19 17:50:43)","OpencodeHint"],[" [msg_9fd985573001fk1Xlot7uyDgTo]","OpencodeHint"]],"virt_text_win_col":-3,"priority":10}],[2,3,0,{"virt_text_repeat_linebreak":true,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["▌","OpencodeMessageRoleUser"]],"virt_text_win_col":-3,"priority":4096}],[3,4,0,{"virt_text_repeat_linebreak":true,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["▌","OpencodeMessageRoleUser"]],"virt_text_win_col":-3,"priority":4096}],[4,5,0,{"virt_text_repeat_linebreak":true,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["▌","OpencodeMessageRoleUser"]],"virt_text_win_col":-3,"priority":4096}],[5,6,0,{"virt_text_repeat_linebreak":true,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["▌","OpencodeMessageRoleUser"]],"virt_text_win_col":-3,"priority":4096}],[6,9,0,{"virt_text_repeat_linebreak":false,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[[" ","OpencodeMessageRoleAssistant"],[" "],["BUILD","OpencodeMessageRoleAssistant"],[" gpt-4.1","OpencodeHint"],[" (2025-10-19 17:50:44)","OpencodeHint"],[" [msg_9fd985a4d001wOX3Op7CpFiCTq]","OpencodeHint"]],"virt_text_win_col":-3,"priority":10}],[7,27,0,{"virt_text_repeat_linebreak":false,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["▌󰭻 ","OpencodeMessageRoleUser"],[" "],["USER","OpencodeMessageRoleUser"],["","OpencodeHint"],[" (2025-10-19 17:50:57)","OpencodeHint"],[" [msg_9fd988c92001w0IZCVPQsN6xa9]","OpencodeHint"]],"virt_text_win_col":-3,"priority":10}],[8,28,0,{"virt_text_repeat_linebreak":true,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["▌","OpencodeMessageRoleUser"]],"virt_text_win_col":-3,"priority":4096}],[9,29,0,{"virt_text_repeat_linebreak":true,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["▌","OpencodeMessageRoleUser"]],"virt_text_win_col":-3,"priority":4096}],[10,32,0,{"virt_text_repeat_linebreak":false,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[[" ","OpencodeMessageRoleAssistant"],[" "],["BUILD","OpencodeMessageRoleAssistant"],[" gpt-4.1","OpencodeHint"],[" (2025-10-19 17:50:57)","OpencodeHint"],[" [msg_9fd988ca7001lgaGttpI4YeGSA]","OpencodeHint"]],"virt_text_win_col":-3,"priority":10}],[11,38,0,{"virt_text_repeat_linebreak":true,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["▌","OpencodeToolBorder"]],"virt_text_win_col":-1,"priority":4096}],[12,39,0,{"virt_text_repeat_linebreak":true,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["▌","OpencodeToolBorder"]],"virt_text_win_col":-1,"priority":4096}],[13,40,0,{"virt_text_repeat_linebreak":true,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["▌","OpencodeToolBorder"]],"virt_text_win_col":-1,"priority":4096}],[14,41,0,{"virt_text_repeat_linebreak":true,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["▌","OpencodeToolBorder"]],"virt_text_win_col":-1,"priority":4096}],[15,42,0,{"virt_text_repeat_linebreak":true,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["▌","OpencodeToolBorder"]],"virt_text_win_col":-1,"priority":4096}],[16,43,0,{"virt_text_repeat_linebreak":true,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["▌","OpencodeToolBorder"]],"virt_text_win_col":-1,"priority":4096}],[17,44,0,{"virt_text_repeat_linebreak":true,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["▌","OpencodeToolBorder"]],"virt_text_win_col":-1,"priority":4096}],[18,45,0,{"virt_text_repeat_linebreak":true,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["▌","OpencodeToolBorder"]],"virt_text_win_col":-1,"priority":4096}],[19,46,0,{"virt_text_repeat_linebreak":true,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["▌","OpencodeToolBorder"]],"virt_text_win_col":-1,"priority":4096}],[20,47,0,{"virt_text_repeat_linebreak":true,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["▌","OpencodeToolBorder"]],"virt_text_win_col":-1,"priority":4096}],[21,48,0,{"virt_text_repeat_linebreak":true,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["▌","OpencodeToolBorder"]],"virt_text_win_col":-1,"priority":4096}],[22,49,0,{"virt_text_repeat_linebreak":true,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["▌","OpencodeToolBorder"]],"virt_text_win_col":-1,"priority":4096}],[23,50,0,{"virt_text_repeat_linebreak":true,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["▌","OpencodeToolBorder"]],"virt_text_win_col":-1,"priority":4096}],[24,51,0,{"virt_text_repeat_linebreak":true,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["▌","OpencodeToolBorder"]],"virt_text_win_col":-1,"priority":4096}],[25,52,0,{"virt_text_repeat_linebreak":true,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["▌","OpencodeToolBorder"]],"virt_text_win_col":-1,"priority":4096}],[26,57,0,{"virt_text_repeat_linebreak":false,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[[" ","OpencodeMessageRoleAssistant"],[" "],["BUILD","OpencodeMessageRoleAssistant"],[" gpt-4.1","OpencodeHint"],[" (2025-10-19 17:50:59)","OpencodeHint"],[" [msg_9fd98942d001elqd2sd8CZeOoA]","OpencodeHint"]],"virt_text_win_col":-3,"priority":10}],[27,67,0,{"virt_text_repeat_linebreak":false,"right_gravity":true,"ns_id":3,"virt_text_hide":false,"virt_text_pos":"win_col","virt_text":[["-20","OpencodeDiffDeleteText"]],"virt_text_win_col":11,"priority":1000}]],"lines":["","----","","","write 10 random words","","[poem.md](poem.md)","","----","","","Here are 10 random words:","","1. Lantern  ","2. Whisper  ","3. Velvet  ","4. Orbit  ","5. Timber  ","6. Quiver  ","7. Mosaic  ","8. Ember  ","9. Spiral  ","10. Glimmer","","Let me know if you need them in a specific format or want to use them in a file!","","----","","","write 10 random words to the file","","----","","","I will write 10 random words to poem.md, each on a new line.","","Proceeding to update the file now.","","**  write** `poem.md`","","```md","Lantern","Whisper","Velvet","Orbit","Timber","Quiver","Mosaic","Ember","Spiral","Glimmer","","```","","**󰻛  Created Snapshot** `c410b2b4`","","----","","","The file poem.md has been updated with 10 random words, each on a new line. Task complete! If you need anything else, let me know.","","----","","> 2 messages reverted, 4 tool calls reverted",">","> type `/redo` to restore.",""," poem.md: -20"],"timestamp":1760961403}