feat(ui): stabilize assistant_mode labeling and improve session formatting robustness\n\n- Persist assistant_mode onto latest assistant message after run\n- Backfill assistant_mode for historical assistant messages for stable display names\n- Use assistant_mode (uppercase) in headers instead of generic ASSISTANT\n- Improve revert stats typing and snapshot action anchoring with display_line + range\n- Harden get_message_at_line nil checks to avoid indexing errors\n- Extend MessagePart.type to include patch and step-start; allow OutputExtmark function form\n- Add assistant_mode field to Message type\n- Fix synthetic user message check (part.synthetic ~= true)\n- Improve callout width handling when window invalid; fallback to config width or 80\n- Add type annotations for tool formatting inputs/metadata/output\n- Correct diff virt_text structure and highlight group names\n- General annotation cleanup and stability improvements

Author: Danilo Verde Ribeiro

lua/opencode/core.lua

@@ -142,6 +142,16 @@ function M.after_run(prompt)
   require('opencode.history').write(prompt)
 
   if state.windows then
+    -- Persist assistant_mode on the latest assistant message if available
+    local msgs = state.messages
+    if msgs and #msgs > 0 then
+      local last = msgs[#msgs]
+      if last and last.role == 'assistant' then
+        if not last.assistant_mode or last.assistant_mode == '' then
+          last.assistant_mode = state.current_mode
+        end
+      end
+    end
     ui.render_output()
   end
 end

lua/opencode/types.lua

@@ -230,7 +230,7 @@
 ---@field msg_idx number|nil Message index in session
 ---@field part_idx number|nil Part index in message
 ---@field role 'user'|'assistant'|'system'|nil Message role
----@field type 'text'|'tool'|'header'|nil Message part type
+---@field type 'text'|'tool'|'header'|'patch'|'step-start'|nil Message part type
 ---@field snapshot? string|nil snapshot commit hash
 
 ---@class OutputAction
@@ -241,7 +241,7 @@
 ---@field display_line number Line number to display the action
 ---@field range? { from: number, to: number } Optional range for the action
 
----@alias OutputExtmark vim.api.keyset.set_extmark
+---@alias OutputExtmark vim.api.keyset.set_extmark|fun():vim.api.keyset.set_extmark
 
 ---@class Message
 ---@field id string Unique message identifier
@@ -256,6 +256,7 @@
 ---@field providerID string Provider identifier
 ---@field role 'user'|'assistant'|'system' Role of the message sender
 ---@field system_role string|nil Role defined in system messages
+---@field assistant_mode string|nil Assistant mode active when message was created
 ---@field error table
 
 ---@class RestorePoint

lua/opencode/ui/session_formatter.lua

@@ -32,6 +32,19 @@ function M.format_session(session)
   M.output:add_line('')
   M.output:add_line('')
 
+  -- Backfill assistant_mode for all assistant messages once so names remain stable
+  local last_seen_mode = state.current_mode
+  for _, amsg in ipairs(state.messages) do
+    if amsg.role == 'assistant' then
+      if amsg.assistant_mode and amsg.assistant_mode ~= '' then
+        last_seen_mode = amsg.assistant_mode
+      else
+        amsg.assistant_mode = last_seen_mode or state.current_mode or 'assistant'
+        last_seen_mode = amsg.assistant_mode
+      end
+    end
+  end
+
   for i, msg in ipairs(state.messages) do
     M.output:add_lines(M.separator)
     state.current_message = msg
@@ -49,6 +62,7 @@ function M.format_session(session)
     end
 
     if session.revert and session.revert.messageID == msg.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, session.revert)
       M._format_revert_message(revert_stats)
       break
@@ -61,7 +75,7 @@ function M.format_session(session)
       M.output:add_metadata(M._current)
 
       if part.type == 'text' and part.text then
-        if msg.role == 'user' and not part.synthetic == true then
+        if msg.role == 'user' and part.synthetic ~= true then
           state.last_user_message = msg
           M._format_user_message(vim.trim(part.text), msg)
         elseif msg.role == 'assistant' then
@@ -85,12 +99,18 @@ function M.format_session(session)
 end
 
 ---@param line number Buffer line number
----@return {message: Message, part: MessagePart, type: string, msg_idx: number, part_idx: number}|nil
+---@return {message: Message, part: MessagePart, msg_idx: number, part_idx: number}|nil
 function M.get_message_at_line(line)
   local metadata = M.output:get_nearest_metadata(line)
   if metadata and metadata.msg_idx and metadata.part_idx then
-    local msg = state.messages[metadata.msg_idx]
+    local msg = state.messages and state.messages[metadata.msg_idx]
+    if not msg or not msg.parts then
+      return nil
+    end
     local part = msg.parts[metadata.part_idx]
+    if not part then
+      return nil
+    end
     return {
       message = msg,
       part = part,
@@ -109,7 +129,7 @@ end
 ---@param messages Message[] All messages in the session
 ---@param revert_index number Index of the message where revert occurred
 ---@param revert_info SessionRevertInfo Revert information
----@return {messages: number, tool_calls: number, files: {additions: number, deletions: number}[]}
+---@return {messages: number, tool_calls: number, files: table<string, {additions: number, deletions: number}>}
 function M._calculate_revert_stats(messages, revert_index, revert_info)
   local stats = {
     messages = 0,
@@ -185,10 +205,10 @@ function M._format_revert_message(stats)
       end
       if #file_diff > 0 then
         local line_str = string.format(icons.get('file') .. '%s: %s', file, table.concat(file_diff, ' '))
-        local line_idx = M.output:add_line(line_str)
-        local col = #('  ' .. file .. ': ')
+        local line_idx = M.output:add_line(line_str) ---@type number
+        local col = #('  ' .. file .. ': ') ---@type number
         for _, diff in ipairs(file_diff) do
-          local hl_group = diff:sub(1, 1) == '+' and 'OpencodeDiffAddText' or 'OpencodeDiffDeleteText'
+           local hl_group = diff:sub(1, 1) == '+' and 'OpencodeDiffAddText' or 'OpencodeDiffDeleteText' ---@type string
           M.output:add_extmark(line_idx, {
             virt_text = { { diff, hl_group } },
             virt_text_pos = 'inline',
@@ -206,23 +226,32 @@ function M._format_patch(part)
   local restore_points = snapshot.get_restore_points_by_parent(part.hash)
   M.output:add_empty_line()
   M._format_action(icons.get('snapshot') .. ' **Created Snapshot**', vim.trim(part.hash:sub(1, 8)))
+  local snapshot_header_line = M.output:get_line_count()
+
+  -- Anchor all snapshot-level actions to the snapshot header line
   M.output:add_action({
     text = '[R]evert file',
     type = 'diff_revert_selected_file',
     args = { part.hash },
     key = 'R',
+    display_line = snapshot_header_line,
+    range = { from = snapshot_header_line, to = snapshot_header_line },
   })
   M.output:add_action({
     text = 'Revert [A]ll',
     type = 'diff_revert_all',
     args = { part.hash },
     key = 'A',
+    display_line = snapshot_header_line,
+    range = { from = snapshot_header_line, to = snapshot_header_line },
   })
   M.output:add_action({
     text = '[D]iff',
     type = 'diff_open',
     args = { part.hash },
     key = 'D',
+    display_line = snapshot_header_line,
+    range = { from = snapshot_header_line, to = snapshot_header_line },
   })
 
   if #restore_points > 0 then
@@ -235,17 +264,22 @@ function M._format_patch(part)
           util.time_ago(restore_point.created_at)
         )
       )
+      local restore_line = M.output:get_line_count() ---@type number
       M.output:add_action({
         text = 'Restore [A]ll',
         type = 'diff_restore_snapshot_all',
         args = { part.hash },
         key = 'A',
+        display_line = restore_line,
+        range = { from = restore_line, to = restore_line },
       })
       M.output:add_action({
         text = '[R]estore file',
         type = 'diff_restore_snapshot_file',
         args = { part.hash },
         key = 'R',
+        display_line = restore_line,
+        range = { from = restore_line, to = restore_line },
       })
     end
   end
@@ -258,6 +292,7 @@ function M._format_error(message)
 end
 
 ---@param message Message
+---@param msg_idx number Message index in the session
 function M._format_message_header(message, msg_idx)
   local role = message.role or 'unknown'
   local icon = message.role == 'user' and icons.get('header_user') or icons.get('header_assistant')
@@ -270,14 +305,28 @@ function M._format_message_header(message, msg_idx)
 
   M.output:add_empty_line()
   M.output:add_metadata({ msg_idx = msg_idx, part_idx = 1, role = role, type = 'header' })
+
+  -- Use the assistant_mode stored on the message only (stable label)
+  local display_name
+  if role == 'assistant' then
+    local mode = message.assistant_mode
+    if mode and mode ~= '' then
+      display_name = mode:upper()
+    else
+      display_name = 'ASSISTANT'
+    end
+  else
+    display_name = role:upper()
+  end
+
   M.output:add_extmark(M.output:get_line_count(), {
     virt_text = {
       { icon, role_hl },
       { ' ' },
-      { role:upper(), role_hl },
+      { display_name, role_hl },
       { model_text, 'OpencodeHint' },
-      { time_text, 'OpenCodeHint' },
-      { debug_text, 'OpenCodeHint' },
+      { time_text, 'OpencodeHint' },
+      { debug_text, 'OpencodeHint' },
     },
     virt_text_win_col = -3,
     priority = 10,
@@ -287,9 +336,11 @@ function M._format_message_header(message, msg_idx)
 end
 
 ---@param callout string Callout type (e.g., 'ERROR', 'TODO')
+---@param text string Callout text content
+---@param title? string Optional title for the callout
 function M._format_callout(callout, text, title)
   title = title and title .. ' ' or ''
-  local win_width = vim.api.nvim_win_get_width(state.windows.output_win)
+  local win_width = (state.windows and state.windows.output_win and vim.api.nvim_win_is_valid(state.windows.output_win)) and vim.api.nvim_win_get_width(state.windows.output_win) or config.ui.window_width or 80
   if #text > win_width - 4 then
     local ok, substituted = pcall(vim.fn.substitute, text, '\v(.{' .. (win_width - 8) .. '})', '\1\n', 'g')
     text = ok and substituted or text
@@ -315,7 +366,7 @@ function M._format_user_message(text, message)
     context = context_module.extract_from_opencode_message(message)
   end
 
-  local start_line = M.output:get_line_count() - 1
+   local start_line = M.output:get_line_count() - 1 ---@type number
 
   M.output:add_empty_line()
   M.output:add_lines(vim.split(context.prompt, '\n'))
@@ -333,7 +384,7 @@ function M._format_user_message(text, message)
     M.output:add_line(string.format('[%s](%s)', path, context.current_file))
   end
 
-  local end_line = M.output:get_line_count()
+     local end_line = M.output:get_line_count() ---@type number
 
   M._add_vertical_border(start_line, end_line, 'OpencodeMessageRoleUser', -3)
 end
@@ -345,6 +396,7 @@ function M._format_assistant_message(text)
 end
 
 ---@param type string Tool type (e.g., 'run', 'read', 'edit', etc.)
+---@param value string Value associated with the action (e.g., filename, command)
 function M._format_action(type, value)
   if not type or not value then
     return
@@ -473,9 +525,12 @@ function M._format_tool(part)
   end
 
   local start_line = M.output:get_line_count() + 1
-  local input = part.state and part.state.input or {}
-  local metadata = part.state.metadata or {}
-  local output = part.state and part.state.output or ''
+   ---@type TaskToolInput|BashToolInput|FileToolInput|TodoToolInput|GlobToolInput|GrepToolInput|WebFetchToolInput|ListToolInput
+   local input = (part.state and part.state.input) or {}
+   ---@type ToolMetadataBase|TaskToolMetadata|WebFetchToolMetadata|BashToolMetadata|FileToolMetadata|GlobToolMetadata|GrepToolMetadata|ListToolMetadata
+   local metadata = (part.state and part.state.metadata) or {}
+   ---@type string
+   local output = (part.state and part.state.output) or ''
 
   if tool == 'bash' then
     M._format_bash_tool(input --[[@as BashToolInput]], metadata --[[@as BashToolMetadata]])
@@ -503,7 +558,7 @@ function M._format_tool(part)
 
   M.output:add_empty_line()
 
-  local end_line = M.output:get_line_count()
+     local end_line = M.output:get_line_count() ---@type number
   if end_line - start_line > 1 then
     M._add_vertical_border(start_line, end_line - 1, 'OpencodeToolBorder', -1)
   end
@@ -532,7 +587,7 @@ function M._format_task_tool(input, metadata, output)
     end
   end
 
-  local end_line = M.output:get_line_count()
+     local end_line = M.output:get_line_count() ---@type number
   M.output:add_action({
     text = '[S]elect Child Session',
     type = 'select_child_session',
@@ -569,7 +624,7 @@ function M._format_diff(code, file_type)
         return {
           end_col = 0,
           end_row = line_idx,
-          virt_text = { { first_char, { hl_group } } },
+          virt_text = { { first_char, hl_group } },
           hl_group = hl_group,
           hl_eol = true,
           priority = 5000,