diff --git a/lua/opencode/commands/handlers/session.lua b/lua/opencode/commands/handlers/session.lua index a2150c1c..f6829212 100644 --- a/lua/opencode/commands/handlers/session.lua +++ b/lua/opencode/commands/handlers/session.lua @@ -70,10 +70,18 @@ end ---@param request_promise Promise ---@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 diff --git a/lua/opencode/state/renderer.lua b/lua/opencode/state/renderer.lua index a2242b86..f2e7c0ed 100644 --- a/lua/opencode/state/renderer.lua +++ b/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', {}) diff --git a/lua/opencode/state/store.lua b/lua/opencode/state/store.lua index f5c91334..b7816280 100644 --- a/lua/opencode/state/store.lua +++ b/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, diff --git a/lua/opencode/ui/input_window.lua b/lua/opencode/ui/input_window.lua index c9c63937..ae96cbb7 100644 --- a/lua/opencode/ui/input_window.lua +++ b/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 diff --git a/lua/opencode/ui/renderer.lua b/lua/opencode/ui/renderer.lua index 852b239c..25b1a837 100644 --- a/lua/opencode/ui/renderer.lua +++ b/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 diff --git a/lua/opencode/ui/renderer/events.lua b/lua/opencode/ui/renderer/events.lua index e250d3c0..2c901368 100644 --- a/lua/opencode/ui/renderer/events.lua +++ b/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) diff --git a/lua/opencode/util.lua b/lua/opencode/util.lua index 3daa540f..ee4344e0 100644 --- a/lua/opencode/util.lua +++ b/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 diff --git a/tests/unit/input_window_spec.lua b/tests/unit/input_window_spec.lua index 7443d198..1d7435f5 100644 --- a/tests/unit/input_window_spec.lua +++ b/tests/unit/input_window_spec.lua @@ -589,4 +589,210 @@ describe('input_window', function() config.values.child_readonly = orig_readonly end) end) + + local function make_message(parts) + return { + info = { id = 'msg_1', sessionID = 'ses_1', role = 'user' }, + parts = parts, + } + end + + describe('build_prompt_from_message', function() + it('returns nil when given a nil message', function() + local prompt = input_window.build_prompt_from_message(nil) + assert.is_nil(prompt) + end) + + it('returns nil when the message has no parts', function() + local prompt = input_window.build_prompt_from_message(make_message({})) + assert.is_nil(prompt) + end) + + it('emits the raw text from a single non-synthetic text part', function() + local prompt = input_window.build_prompt_from_message(make_message({ + { type = 'text', text = 'hello world' }, + })) + assert.same({ 'hello world' }, prompt.lines) + assert.same({}, prompt.mention_paths) + end) + + it('skips synthetic text parts', function() + local prompt = input_window.build_prompt_from_message(make_message({ + { type = 'text', synthetic = true, text = 'should be dropped' }, + { type = 'text', text = 'keep me' }, + })) + assert.same({ 'keep me' }, prompt.lines) + end) + + it('emits @ tokens for file parts using filename', function() + local prompt = input_window.build_prompt_from_message(make_message({ + { type = 'text', text = 'look at' }, + { type = 'file', filename = 'lua/opencode/foo.lua' }, + { type = 'text', text = 'thanks' }, + })) + assert.same({ 'look at', '@lua/opencode/foo.lua ', 'thanks' }, prompt.lines) + assert.same({ 'lua/opencode/foo.lua' }, prompt.mention_paths) + end) + + it('falls back to source.path when filename is missing', function() + local prompt = input_window.build_prompt_from_message(make_message({ + { type = 'file', source = { path = 'src/main.lua' } }, + })) + assert.same({ '@src/main.lua ' }, prompt.lines) + assert.same({ 'src/main.lua' }, prompt.mention_paths) + end) + + it('emits @ tokens for agent parts', function() + local prompt = input_window.build_prompt_from_message(make_message({ + { type = 'text', text = 'use' }, + { type = 'agent', name = 'build' }, + { type = 'text', text = 'to compile' }, + })) + assert.same({ 'use', '@build ', 'to compile' }, prompt.lines) + assert.same({ 'build' }, prompt.mention_paths) + end) + + it('skips tool, step-start, and patch parts', function() + local prompt = input_window.build_prompt_from_message(make_message({ + { type = 'text', text = 'first' }, + { type = 'tool', text = 'should be dropped' }, + { type = 'step-start' }, + { type = 'patch', text = 'also dropped' }, + { type = 'text', text = 'last' }, + })) + assert.same({ 'first', 'last' }, prompt.lines) + assert.same({}, prompt.mention_paths) + end) + + it('splits text parts on embedded newlines into separate lines', function() + local prompt = input_window.build_prompt_from_message(make_message({ + { type = 'text', text = 'line1\nline2' }, + })) + assert.same({ 'line1', 'line2' }, prompt.lines) + end) + + it('splits text parts on embedded newlines interleaved with mentions', function() + local prompt = input_window.build_prompt_from_message(make_message({ + { type = 'text', text = 'before' }, + { type = 'file', filename = 'a.lua' }, + { type = 'text', text = 'middle\nmore' }, + { type = 'agent', name = 'build' }, + { type = 'text', text = 'after' }, + })) + assert.same({ + 'before', + '@a.lua ', + 'middle', + 'more', + '@build ', + 'after', + }, prompt.lines) + assert.same({ 'a.lua', 'build' }, prompt.mention_paths) + end) + + it('handles nil and non-string fields defensively', function() + local prompt = input_window.build_prompt_from_message(make_message({ + { type = 'text', text = nil }, + { type = 'text' }, + { type = 'text', text = 'safe' }, + { type = 'file', filename = nil }, + { type = 'agent', name = '' }, + })) + assert.same({ 'safe' }, prompt.lines) + assert.same({}, prompt.mention_paths) + end) + end) + + describe('refill_prompt_from_message', function() + local function open_input_window() + local input_buf = vim.api.nvim_create_buf(false, true) + local output_buf = vim.api.nvim_create_buf(false, true) + local output_win = vim.api.nvim_open_win(output_buf, false, { + relative = 'editor', + width = 80, + height = 5, + row = 0, + col = 0, + }) + local input_win = vim.api.nvim_open_win(input_buf, true, { + relative = 'editor', + width = 80, + height = 5, + row = 6, + col = 0, + }) + state.ui.set_windows({ + input_buf = input_buf, + input_win = input_win, + output_buf = output_buf, + output_win = output_win, + }) + vim.api.nvim_set_current_win(output_win) + input_window._hidden = false + return input_buf, input_win, output_buf, output_win + end + + local function cleanup(input_buf, input_win, output_buf, output_win) + pcall(vim.api.nvim_win_close, input_win, true) + pcall(vim.api.nvim_win_close, output_win, true) + pcall(vim.api.nvim_buf_delete, input_buf, { force = true }) + pcall(vim.api.nvim_buf_delete, output_buf, { force = true }) + state.ui.clear_windows() + end + + it('parks the cursor at the end of the refilled text', function() + local input_buf, input_win, output_buf, output_win = open_input_window() + local message = make_message({ + { type = 'text', text = 'refactor this' }, + }) + input_window.refill_prompt_from_message(message) + local lines = vim.api.nvim_buf_get_lines(input_buf, 0, -1, false) + local cursor = vim.api.nvim_win_get_cursor(input_win) + assert.equals(#lines, cursor[1]) + assert.equals(#lines[#lines] - 1, cursor[2]) + assert.same({ 'refactor this' }, lines) + cleanup(input_buf, input_win, output_buf, output_win) + end) + + it('parks the cursor on the last line of a multi-line refill', function() + local input_buf, input_win, output_buf, output_win = open_input_window() + local message = make_message({ + { type = 'text', text = 'line1' }, + { type = 'text', text = 'line2' }, + { type = 'text', text = 'line3' }, + }) + input_window.refill_prompt_from_message(message) + local lines = vim.api.nvim_buf_get_lines(input_buf, 0, -1, false) + local cursor = vim.api.nvim_win_get_cursor(input_win) + assert.equals(#lines, cursor[1]) + assert.equals(#lines[#lines] - 1, cursor[2]) + assert.same({ 'line1', 'line2', 'line3' }, lines) + cleanup(input_buf, input_win, output_buf, output_win) + end) + + it('parks the cursor after the mention token when a file is attached', function() + local input_buf, input_win, output_buf, output_win = open_input_window() + local message = make_message({ + { type = 'text', text = 'look at' }, + { type = 'file', filename = 'lua/opencode/foo.lua' }, + { type = 'text', text = 'thanks' }, + }) + input_window.refill_prompt_from_message(message) + local lines = vim.api.nvim_buf_get_lines(input_buf, 0, -1, false) + local cursor = vim.api.nvim_win_get_cursor(input_win) + assert.equals(#lines, cursor[1]) + assert.equals(#lines[#lines] - 1, cursor[2]) + assert.equals('thanks', lines[#lines]) + cleanup(input_buf, input_win, output_buf, output_win) + end) + + it('returns false and does not touch the buffer when there is nothing to refill', function() + local input_buf, input_win, output_buf, output_win = open_input_window() + vim.api.nvim_buf_set_lines(input_buf, 0, -1, false, { 'untouched' }) + local filled = input_window.refill_prompt_from_message(make_message({})) + assert.is_false(filled) + assert.same({ 'untouched' }, vim.api.nvim_buf_get_lines(input_buf, 0, -1, false)) + cleanup(input_buf, input_win, output_buf, output_win) + end) + end) end)