feat(restore-point): make restore points work

Author: sudo-tee

lua/opencode/event_manager.lua

@@ -88,7 +88,10 @@ local state = require('opencode.state')
 
 --- @class ServerStoppedEvent
 
---- @alias EventName
+--- @class RestorePointCreatedEvent
+--- @field restore_point RestorePoint
+
+--- @alias OpencodeEventName
 --- | "installation.updated"
 --- | "lsp.client.diagnostics"
 --- | "message.updated"
@@ -109,6 +112,7 @@ local state = require('opencode.state')
 --- | "custom.server_starting"
 --- | "custom.server_ready"
 --- | "custom.server_stopped"
+--- | "custom.restore_point.created"
 
 --- @class EventManager
 --- @field events table<string, function[]> Event listener registry
@@ -150,7 +154,8 @@ end
 --- @overload fun(self: EventManager, event_name: "custom.server_starting", callback: fun(data: ServerStartingEvent): nil)
 --- @overload fun(self: EventManager, event_name: "custom.server_ready", callback: fun(data: ServerReadyEvent): nil)
 --- @overload fun(self: EventManager, event_name: "custom.server_stopped", callback: fun(data: ServerStoppedEvent): nil)
---- @param event_name EventName The event name to listen for
+--- @overload fun(self: EventManager, event_name: "custom.restore_point.created", callback: fun(data: RestorePointCreatedEvent): nil)
+--- @param event_name OpencodeEventName The event name to listen for
 --- @param callback function Callback function to execute when event is triggered
 function EventManager:subscribe(event_name, callback)
   if not self.events[event_name] then
@@ -180,7 +185,8 @@ end
 --- @overload fun(self: EventManager, event_name: "custom.server_starting", callback: fun(data: ServerStartingEvent): nil)
 --- @overload fun(self: EventManager, event_name: "custom.server_ready", callback: fun(data: ServerReadyEvent): nil)
 --- @overload fun(self: EventManager, event_name: "custom.server_stopped", callback: fun(data: ServerStoppedEvent): nil)
---- @param event_name EventName The event name
+--- @overload fun(self: EventManager, event_name: "custom.restore_point.created", callback: fun(data: RestorePointCreatedEvent): nil)
+--- @param event_name OpencodeEventName The event name
 --- @param callback function The callback function to remove
 function EventManager:unsubscribe(event_name, callback)
   local listeners = self.events[event_name]
@@ -197,14 +203,21 @@ function EventManager:unsubscribe(event_name, callback)
 end
 
 --- Emit an event to all subscribers
---- @param event_name EventName The event name
+--- @param event_name OpencodeEventName The event name
 --- @param data any Data to pass to event listeners
 function EventManager:emit(event_name, data)
   local listeners = self.events[event_name]
   if not listeners then
     return
   end
 
+  local event = { type = event_name, properties = data }
+
+  if require('opencode.config').debug.capture_streamed_events then
+    table.insert(self.captured_events, vim.deepcopy(event))
+  end
+
+  -- schedule events to allow for similar pieces of state to be updated
   for _, callback in ipairs(listeners) do
     pcall(callback, data)
   end
@@ -269,22 +282,11 @@ function EventManager:_subscribe_to_server_events(server)
   local api_client = state.api_client
 
   local emitter = function(event)
-    -- schedule events to allow for similar pieces of state to be updated
     vim.schedule(function()
       self:emit(event.type, event.properties)
     end)
   end
 
-  if require('opencode.config').debug.capture_streamed_events then
-    local _emitter = emitter
-    emitter = function(event)
-      -- make a deepcopy to make sure we're saving a clean copy
-      -- (we modify event in renderer)
-      table.insert(self.captured_events, vim.deepcopy(event))
-      _emitter(event)
-    end
-  end
-
   self.server_subscription = api_client:subscribe_to_events(nil, emitter)
 end
 
@@ -312,7 +314,7 @@ function EventManager:get_event_names()
 end
 
 --- Get number of subscribers for an event
---- @param event_name EventName The event name
+--- @param event_name OpencodeEventName The event name
 --- @return number Number of subscribers
 function EventManager:get_subscriber_count(event_name)
   local listeners = self.events[event_name]

lua/opencode/snapshot.lua

@@ -93,7 +93,7 @@ function M.save_restore_point(snapshot_id, from_snapshot_id, deleted_files)
     return nil
   end
 
-  state.append('restore_points', snapshot)
+  state.event_manager:emit('custom.restore_point.created', { restore_point = snapshot })
   return snapshot
 end
 
@@ -254,7 +254,7 @@ function M.restore_file(snapshot_id, file_path)
 end
 
 ---@param from_snapshot_id string
----@return RestorePoint[]
+---@return RestorePoint[]|nil
 function M.get_restore_points_by_parent(from_snapshot_id)
   local restore_points = M.get_restore_points()
   restore_points = vim.tbl_filter(function(item)
@@ -263,6 +263,9 @@ function M.get_restore_points_by_parent(from_snapshot_id)
   table.sort(restore_points, function(a, b)
     return a.created_at > b.created_at
   end)
+  if #restore_points == 0 then
+    return nil
+  end
   return restore_points
 end
 

lua/opencode/state.lua

@@ -20,7 +20,7 @@ local config = require('opencode.config')
 ---@field last_output number
 ---@field last_sent_context any
 ---@field active_session Session|nil
----@field restore_points table<string, any>
+---@field restore_points RestorePoint[]
 ---@field current_model string|nil
 ---@field current_model_info table|nil
 ---@field messages OpencodeMessage[]|nil

lua/opencode/ui/formatter.lua

@@ -156,7 +156,7 @@ end
 ---@param output Output Output object to write to
 ---@param part MessagePart
 function M._format_patch(output, part)
-  local restore_points = snapshot.get_restore_points_by_parent(part.hash)
+  local restore_points = snapshot.get_restore_points_by_parent(part.hash) or {}
   M._format_action(output, icons.get('snapshot') .. ' Created Snapshot', vim.trim(part.hash:sub(1, 8)))
   local snapshot_header_line = output:get_line_count()
 
@@ -200,15 +200,15 @@ function M._format_patch(output, part)
       output:add_action({
         text = 'Restore [A]ll',
         type = 'diff_restore_snapshot_all',
-        args = { part.hash },
+        args = { restore_point.id },
         key = 'A',
         display_line = restore_line,
         range = { from = restore_line, to = restore_line },
       })
       output:add_action({
         text = '[R]estore file',
         type = 'diff_restore_snapshot_file',
-        args = { part.hash },
+        args = { restore_point.id },
         key = 'R',
         display_line = restore_line,
         range = { from = restore_line, to = restore_line },

lua/opencode/ui/output_window.lua

@@ -51,11 +51,6 @@ function M.setup(windows)
 
   M.update_dimensions(windows)
   M.setup_keymaps(windows)
-  state.subscribe('restore_points', function(_, new_val, old_val)
-    -- FIXME: restore points
-    -- local outout_renderer = require('opencode.ui.output_renderer')
-    -- outout_renderer.render(state.windows, true)
-  end)
 end
 
 function M.update_dimensions(windows)

lua/opencode/ui/render_state.lua

@@ -73,6 +73,20 @@ function RenderState:get_part_by_call_id(call_id, message_id)
   return nil
 end
 
+---Get part ID by snapshot_id and message ID
+---@param snapshot_id string Call ID
+---@return MessagePart? part Part if found
+function RenderState:get_part_by_snapshot_id(snapshot_id)
+  for _, rendered_message in pairs(self._messages) do
+    for _, part in ipairs(rendered_message.message.parts) do
+      if part.type == 'patch' and part.hash == snapshot_id then
+        return part
+      end
+    end
+  end
+  return nil
+end
+
 ---Ensure line index is up to date
 function RenderState:_ensure_line_index()
   if not self._line_index_valid then

lua/opencode/ui/renderer.lua

@@ -77,6 +77,7 @@ function M._setup_event_subscriptions(subscribe)
   state.event_manager[method](state.event_manager, 'permission.updated', M.on_permission_updated)
   state.event_manager[method](state.event_manager, 'permission.replied', M.on_permission_replied)
   state.event_manager[method](state.event_manager, 'file.edited', M.on_file_edited)
+  state.event_manager[method](state.event_manager, 'custom.restore_point.created', M.on_restore_points)
 end
 
 ---Unsubscribe from local state and server subscriptions
@@ -635,6 +636,18 @@ function M.on_file_edited(properties)
   vim.cmd('checktime')
 end
 
+---@param properties RestorePointCreatedEvent
+function M.on_restore_points(properties)
+  state.append('restore_points', properties.restore_point)
+  if not properties or not properties.restore_point or not properties.restore_point.from_snapshot_id then
+    return
+  end
+  local part = M._render_state:get_part_by_snapshot_id(properties.restore_point.from_snapshot_id)
+  if part then
+    M.on_part_updated({ part = part })
+  end
+end
+
 ---Find part ID by call ID and message ID
 ---Useful for finding a part for a permission
 ---@param call_id string Call ID to search for