fix/fill input

lua/opencode/commands/handlers/session.lua

@@ -70,10 +70,18 @@ end
 
 ---@param request_promise Promise<any>
 ---@param error_prefix string
-local function run_api_action_with_checktime(request_promise, error_prefix)
-  request_promise:and_then(schedule_checktime):catch(function(err)
-    notify_error(error_prefix, err)
-  end)
+---@param on_success? fun(...)
+local function run_api_action_with_checktime(request_promise, error_prefix, on_success)
+  request_promise
+    :and_then(function(...)
+      schedule_checktime()
+      if on_success then
+        on_success(...)
+      end
+    end)
+    :catch(function(err)
+      notify_error(error_prefix, err)
+    end)
 end
 
 function M.actions.open_input_new_session()
@@ -409,20 +417,55 @@ function M.actions.rename_session(current_session, new_title)
   end)(current_session, new_title)
 end
 
----@param messageId? string
-function M.actions.undo(messageId)
+---@param state_obj OpencodeState
+---@param target_id string
+---@return OpencodeMessage|nil
+local function find_message_in_state(state_obj, target_id)
+  for _, m in ipairs(state_obj.messages or {}) do
+    if m.info and m.info.id == target_id then
+      return m
+    end
+  end
+  return nil
+end
+
+---@param state_obj OpencodeState
+---@return OpencodeMessage|nil
+local function find_last_user_message(state_obj)
+  local messages = state_obj.messages or {}
+  local revert = state_obj.active_session and state_obj.active_session.revert
+
+  local revert_index = revert
+    and require('opencode.util').find_index_of(messages, function(m)
+      return m.info and m.info.id == revert.messageID
+    end)
+
+  for i = revert_index and revert_index - 1 or #messages, 1, -1 do
+    local m = messages[i]
+    if m.info and m.info.role == 'user' then
+      return m
+    end
+  end
+  return nil
+end
+
+---@param message_id? string
+function M.actions.undo(message_id)
   return with_active_session('No active session to undo', function(state_obj)
-    local message_to_revert = messageId or (state_obj.last_user_message and state_obj.last_user_message.info.id)
-    if not message_to_revert then
+    local target = message_id and find_message_in_state(state_obj, message_id) or find_last_user_message(state_obj)
+    if not target then
       vim.notify('No user message to undo', vim.log.levels.WARN)
       return
     end
 
     run_api_action_with_checktime(
       state_obj.api_client:revert_message(state_obj.active_session.id, {
-        messageID = message_to_revert,
+        messageID = target.info.id,
       }),
-      'Failed to undo last message: '
+      'Failed to undo last message: ',
+      function()
+        require('opencode.ui.input_window').refill_prompt_from_message(target)
+      end
     )
   end)
 end
@@ -526,7 +569,12 @@ end
 ---@param message_id? string
 function M.actions.fork_session(message_id)
   return with_active_session('No active session to fork', function(state_obj)
-    local message_to_fork = message_id or state_obj.last_user_message and state_obj.last_user_message.info.id
+    local target = message_id and find_message_in_state(state_obj, message_id) or find_last_user_message(state_obj)
+    if not target then
+      vim.notify('No user message to fork from', vim.log.levels.WARN)
+      return
+    end
+    local message_to_fork = target.info.id
     if not message_to_fork then
       vim.notify('No user message to fork from', vim.log.levels.WARN)
       return

lua/opencode/state/renderer.lua

@@ -13,11 +13,6 @@ function M.set_current_message(message)
   return store.set('current_message', message)
 end
 
----@param message OpencodeMessage|nil
-function M.set_last_user_message(message)
-  return store.set('last_user_message', message)
-end
-
 ---@param permissions OpencodePermission[]
 function M.set_pending_permissions(permissions)
   return store.set('pending_permissions', permissions)
@@ -27,6 +22,7 @@ end
 function M.update_pending_permissions(mutator)
   return store.mutate('pending_permissions', mutator)
 end
+
 ---@param cost number
 function M.set_cost(cost)
   if not cost or cost <= 0 then
@@ -55,7 +51,6 @@ function M.reset()
   return store.batch(function()
     store.set('messages', {})
     store.set('current_message', nil)
-    store.set('last_user_message', nil)
     store.set('tokens_count', 0)
     store.set('cost', 0)
     store.set('pending_permissions', {})

lua/opencode/state/store.lua

@@ -27,7 +27,6 @@ local M = {}
 ---@field current_variant string|nil
 ---@field messages OpencodeMessage[]|nil
 ---@field current_message OpencodeMessage|nil
----@field last_user_message OpencodeMessage|nil
 ---@field pending_permissions OpencodePermission[]
 ---@field cost number
 ---@field tokens_count number
@@ -72,7 +71,6 @@ local _state = {
   current_variant = nil,
   messages = nil,
   current_message = nil,
-  last_user_message = nil,
   pending_permissions = {},
   cost = 0,
   tokens_count = 0,

lua/opencode/ui/input_window.lua

@@ -426,6 +426,72 @@ function M.set_content(text, windows)
   vim.api.nvim_buf_set_lines(windows.input_buf, 0, -1, false, lines)
 end
 
+---@param message OpencodeMessage|nil
+---@return { lines: string[], mention_paths: string[] }|nil
+function M.build_prompt_from_message(message)
+  if not message or not message.parts then
+    return nil
+  end
+
+  local lines = {}
+  local mention_paths = {}
+
+  for _, part in ipairs(message.parts) do
+    if type(part) == 'table' then
+      if part.type == 'text' then
+        if not part.synthetic and type(part.text) == 'string' and part.text ~= '' then
+          for _, sub in ipairs(vim.split(part.text, '\n', { plain = true })) do
+            lines[#lines + 1] = sub
+          end
+        end
+      elseif part.type == 'file' then
+        local name = part.filename or (part.source and part.source.path) or part.name
+        if type(name) == 'string' and name ~= '' then
+          lines[#lines + 1] = '@' .. name .. ' '
+          table.insert(mention_paths, name)
+        end
+      elseif part.type == 'agent' then
+        local name = part.name or (part.source and part.source.path)
+        if type(name) == 'string' and name ~= '' then
+          lines[#lines + 1] = '@' .. name .. ' '
+          table.insert(mention_paths, name)
+        end
+      end
+    end
+  end
+
+  if #lines == 0 then
+    return nil
+  end
+
+  return { lines = lines, mention_paths = mention_paths }
+end
+
+---@param message OpencodeMessage|nil
+---@return boolean
+function M.refill_prompt_from_message(message)
+  local prompt = M.build_prompt_from_message(message)
+  if not prompt then
+    return false
+  end
+  if not M.mounted() then
+    return false
+  end
+  ---@cast state.windows { input_win: integer, input_buf: integer }
+
+  M.set_content(prompt.lines)
+  require('opencode.ui.mention').restore_mentions(state.windows.input_buf)
+
+  -- nvim_win_set_cursor clamps col to the line's last byte, which is the
+  -- canonical end-of-line position both for 'a' (append) and 'i' (insert).
+  pcall(vim.api.nvim_win_set_cursor, state.windows.input_win, {
+    math.max(#prompt.lines, 1),
+    #(prompt.lines[#prompt.lines] or ''),
+  })
+
+  return true
+end
+
 function M.set_current_line(text, windows)
   windows = windows or state.windows
   if not M.mounted(windows) then

lua/opencode/ui/renderer.lua

@@ -333,7 +333,6 @@ local function fetch_session()
   if not session or session == '' then
     return Promise.new():resolve(nil)
   end
-  state.renderer.set_last_user_message(nil)
   return require('opencode.session').get_messages(session)
 end
 

lua/opencode/ui/renderer/events.lua

@@ -288,12 +288,9 @@ function M.on_message_updated(message, revert_index)
     state.renderer.set_current_message(msg)
   end
 
-  if msg.info.role == 'user' then
-    state.renderer.set_last_user_message(msg)
-    if not found_before then
-      local local_submit_pending = (state.user_message_count or {})[msg.info.sessionID] or 0
-      scroll(local_submit_pending > 0)
-    end
+  if msg.info.role == 'user' and not found_before then
+    local local_submit_pending = (state.user_message_count or {})[msg.info.sessionID] or 0
+    scroll(local_submit_pending > 0)
   end
 
   update_stats(msg)

lua/opencode/util.lua

@@ -297,6 +297,9 @@ function M.index_of(tbl, value)
   return nil
 end
 
+---@generic T
+---@param tbl T[]
+---@param predicate function(item: T): boolean
 function M.find_index_of(tbl, predicate)
   for i, v in ipairs(tbl) do
     if predicate(v) then