refactor(state): unify model clearing and improve renderer/session reset

Author: sudo-tee

lua/opencode/core.lua

@@ -37,12 +37,10 @@ end)
 M.switch_session = Promise.async(function(session_id)
   local selected_session = session.get_by_id(session_id):await()
 
-  state.model.clear_model()
-  state.model.clear_mode()
+  state.model.clear()
   M.ensure_current_mode():await()
 
   state.session.set_active(selected_session)
-  state.session.reset_restore_points()
   if state.ui.is_visible() then
     ui.focus_input()
   else

lua/opencode/state/model.lua

@@ -1,5 +1,17 @@
 local store = require('opencode.state.store')
 
+---@class OpencodeModelStateMutations
+---@field set_mode fun(mode: string|nil)
+---@field clear_mode fun()
+---@field set_model fun(model: string|nil)
+---@field clear_model fun()
+---@field clear fun()
+---@field set_model_info fun(info: table|nil)
+---@field set_variant fun(variant: string|nil)
+---@field clear_variant fun()
+---@field set_mode_model_map fun(mode_map: table<string, string>)
+---@field set_mode_model_override fun(mode: string, model: string)
+
 local M = {}
 
 ---@param mode string|nil
@@ -20,6 +32,12 @@ function M.clear_model()
   return store.set('current_model', nil)
 end
 
+function M.clear()
+  store.set('current_model', nil)
+  store.set('current_mode', nil)
+  store.set('current_variant', nil)
+end
+
 ---@param info table|nil
 function M.set_model_info(info)
   return store.set('current_model_info', info)
@@ -42,10 +60,11 @@ end
 ---@param mode string
 ---@param model string
 function M.set_mode_model_override(mode, model)
-  local state = store.state()
-  local mode_map = vim.deepcopy(state.user_mode_model_map)
-  mode_map[mode] = model
-  return store.set('user_mode_model_map', mode_map)
+  return store.update('user_mode_model_map', function(current)
+    local updated = vim.deepcopy(current)
+    updated[mode] = model
+    return updated
+  end)
 end
 
 return M

lua/opencode/state/renderer.lua

@@ -8,6 +8,8 @@ local store = require('opencode.state.store')
 ---@field set_pending_permissions fun(permissions: OpencodePermission[])
 ---@field set_cost fun(cost: number)
 ---@field set_tokens_count fun(count: number)
+---@field set_stats fun(tokens_count: number, cost: number)
+---@field reset fun()
 
 local M = {}
 
@@ -41,4 +43,19 @@ function M.set_tokens_count(count)
   return store.set('tokens_count', count)
 end
 
+---@param tokens_count number
+---@param cost number
+function M.set_stats(tokens_count, cost)
+  store.set('tokens_count', tokens_count)
+  store.set('cost', cost)
+end
+
+function M.reset()
+  store.set('messages', {})
+  store.set('last_user_message', nil)
+  store.set('tokens_count', 0)
+  store.set('cost', 0)
+  store.set('pending_permissions', {})
+end
+
 return M

lua/opencode/state/session.lua

@@ -12,7 +12,9 @@ local M = {}
 
 ---@param session Session|nil
 function M.set_active(session)
-  M.clear_active()
+  store.set('restore_points', {})
+  store.set('last_sent_context', nil)
+  store.set('user_message_count', {})
   return store.set('active_session', session)
 end
 

lua/opencode/ui/renderer.lua

@@ -43,18 +43,14 @@ function M.reset()
 
   output_window.clear()
 
-  state.renderer.set_messages({})
-  state.renderer.set_last_user_message(nil)
-  state.renderer.set_tokens_count(0)
-
   local permissions = state.pending_permissions or {}
   if #permissions > 0 and state.api_client then
     for _, permission in ipairs(permissions) do
       require('opencode.api').permission_deny(permission)
     end
   end
   permission_window.clear_all()
-  state.renderer.set_pending_permissions({})
+  state.renderer.reset()
 
   trigger_on_data_rendered()
 end
@@ -1198,11 +1194,14 @@ function M._update_stats_from_message(message)
   end
 
   local tokens = message.info.tokens
-  if tokens and tokens.input > 0 then
+  if tokens and tokens.input > 0 and message.info.cost and type(message.info.cost) == 'number' then
+    state.renderer.set_stats(
+      tokens.input + tokens.output + tokens.cache.read + tokens.cache.write,
+      message.info.cost
+    )
+  elseif tokens and tokens.input > 0 then
     state.renderer.set_tokens_count(tokens.input + tokens.output + tokens.cache.read + tokens.cache.write)
-  end
-
-  if message.info.cost and type(message.info.cost) == 'number' then
+  elseif message.info.cost and type(message.info.cost) == 'number' then
     state.renderer.set_cost(message.info.cost)
   end
 end