refactor(state): improve use of state helpers (#330)

Author: sudo-tee

lua/opencode/context.lua

@@ -22,14 +22,10 @@ local toggleable_context_keys = {
 ---@param context_key OpencodeToggleableContextKey
 ---@return table
 local function ensure_context_state(context_key)
-  local current_config = state.current_context_config or {}
-  local current = current_config[context_key]
-  local new_config = vim.deepcopy(current_config)
   local defaults = vim.tbl_get(config, 'context', context_key) or {}
-
-  new_config[context_key] = vim.tbl_deep_extend('force', {}, defaults, current or {})
-  state.context.set_current_context_config(new_config)
-  return new_config[context_key]
+  return state.context.update_current_context_config(function(current_config)
+    current_config[context_key] = vim.tbl_deep_extend('force', {}, defaults, current_config[context_key] or {})
+  end)[context_key]
 end
 
 M.ChatContext = ChatContext
@@ -120,12 +116,10 @@ end
 -- Delegate global state management to ChatContext
 function M.add_selection(selection)
   ChatContext.add_selection(selection)
-  state.context.set_context_updated_at(vim.uv.now())
 end
 
 function M.remove_selection(selection)
   ChatContext.remove_selection(selection)
-  state.context.set_context_updated_at(vim.uv.now())
 end
 
 function M.clear_selections()
@@ -200,12 +194,7 @@ function M.build_inline_selection_text(range)
   end
 
   local filetype = vim.bo[buf].filetype or ''
-  local text = string.format(
-    '**`%s`**\n\n```%s\n%s\n```',
-    file.path,
-    filetype,
-    current_selection.text
-  )
+  local text = string.format('**`%s`**\n\n```%s\n%s\n```', file.path, filetype, current_selection.text)
 
   return text
 end
@@ -225,13 +214,11 @@ function M.add_file(file)
 
   file = vim.fn.fnamemodify(file, ':p')
   ChatContext.add_file(file)
-  state.context.set_context_updated_at(vim.uv.now())
 end
 
 function M.remove_file(file)
   file = vim.fn.fnamemodify(file, ':p')
   ChatContext.remove_file(file)
-  state.context.set_context_updated_at(vim.uv.now())
 end
 
 function M.clear_files()
@@ -240,12 +227,10 @@ end
 
 function M.add_subagent(subagent)
   ChatContext.add_subagent(subagent)
-  state.context.set_context_updated_at(vim.uv.now())
 end
 
 function M.remove_subagent(subagent)
   ChatContext.remove_subagent(subagent)
-  state.context.set_context_updated_at(vim.uv.now())
 end
 
 function M.clear_subagents()
@@ -258,7 +243,6 @@ end
 
 function M.load()
   ChatContext.load()
-  state.context.set_context_updated_at(vim.uv.now())
 end
 
 -- Context creation with delta logic (delegates to ChatContext)

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
@@ -94,8 +92,7 @@ M.open = Promise.async(function(opts)
 
   if are_windows_closed then
     if not ui.is_opencode_focused() then
-      state.ui.set_last_code_window(vim.api.nvim_get_current_win())
-      state.ui.set_current_code_buf(vim.api.nvim_get_current_buf())
+      state.ui.set_code_context(vim.api.nvim_get_current_win(), vim.api.nvim_get_current_buf())
     end
 
     M.is_prompting_allowed()

lua/opencode/image_handler.lua

@@ -150,7 +150,6 @@ function M.paste_image_from_clipboard()
 
   if success then
     context.add_file(image_path)
-    state.context.set_context_updated_at(os.time())
     vim.notify('Image saved and added to context: ' .. vim.fn.fnamemodify(image_path, ':t'), vim.log.levels.INFO)
     return true
   end

lua/opencode/server_job.lua

@@ -293,7 +293,7 @@ function M.spawn_local_server(promise, port, hostname)
       log.notify(string.format('Started local server at %s', base_url), vim.log.levels.INFO)
       if url_port then
         local port_num = tonumber(url_port)
-        state.opencode_server.port = port_num
+        state.jobs.set_server_port(port_num)
         local server_pid = job and job.pid
         port_mapping.register(port_num, vim.fn.getcwd(), true, 'serve', nil, server_pid)
         log.debug(

lua/opencode/state.lua

@@ -1 +1 @@
-return require('opencode.state.init') --[[@as OpencodeStateModule]]
+return require('opencode.state.init') --[[@as OpencodeState]]

lua/opencode/state/context.lua

@@ -1,18 +1,17 @@
-
 local store = require('opencode.state.store')
 
 ---@class OpencodeContextStateMutations
----@field set_current_context_config fun(config: OpencodeContextConfig|nil)
----@field set_context_updated_at fun(timestamp: number|nil)
----@field set_current_cwd fun(cwd: string|nil)
-
 local M = {}
 
 ---@param config OpencodeContextConfig|nil
 function M.set_current_context_config(config)
   return store.set('current_context_config', config)
 end
 
+function M.update_current_context_config(mutator)
+  return store.mutate('current_context_config', mutator)
+end
+
 ---@param timestamp number|nil
 function M.set_context_updated_at(timestamp)
   return store.set('context_updated_at', timestamp)

lua/opencode/state/init.lua

@@ -16,6 +16,7 @@ local context = require('opencode.state.context')
 ---@field context OpencodeContextStateMutations
 
 ---@alias OpencodeState OpencodeStateModule & OpencodeStateData
+
 ---@type OpencodeState
 local M = {
   store = store,

lua/opencode/state/jobs.lua

@@ -1,16 +1,6 @@
 local store = require('opencode.state.store')
 
 ---@class OpencodeJobStateMutations
----@field increment_count fun(delta?: integer)
----@field decrement_count fun(delta?: integer)
----@field set_count fun(count: integer)
----@field set_server fun(server: OpencodeServer|nil)
----@field clear_server fun()
----@field set_api_client fun(client: OpencodeApiClient|nil)
----@field set_event_manager fun(manager: EventManager|nil)
----@field set_opencode_cli_version fun(version: string|nil)
----@field is_running fun():boolean
-
 local M = {}
 
 ---@param delta integer|nil
@@ -41,6 +31,18 @@ function M.clear_server()
   return store.set('opencode_server', nil)
 end
 
+---@param port integer
+function M.set_server_port(port)
+  local server = store.get('opencode_server')
+  if not server then
+    error('Opencode server is not set; cannot set port')
+  end
+
+  store.mutate('opencode_server', function(s)
+    s.port = port
+  end)
+end
+
 ---@param client OpencodeApiClient|nil
 function M.set_api_client(client)
   return store.set('api_client', client)

lua/opencode/state/model.lua

@@ -1,5 +1,6 @@
 local store = require('opencode.state.store')
 
+---@class OpencodeModelStateMutations
 local M = {}
 
 ---@param mode string|nil
@@ -20,6 +21,14 @@ function M.clear_model()
   return store.set('current_model', nil)
 end
 
+function M.clear()
+  return store.batch(function()
+    store.set('current_model', nil)
+    store.set('current_mode', nil)
+    store.set('current_variant', nil)
+  end)
+end
+
 ---@param info table|nil
 function M.set_model_info(info)
   return store.set('current_model_info', info)
@@ -42,10 +51,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

@@ -1,14 +1,6 @@
-
 local store = require('opencode.state.store')
 
 ---@class OpencodeRendererStateMutations
----@field set_messages fun(messages: OpencodeMessage[]|nil)
----@field set_current_message fun(message: OpencodeMessage|nil)
----@field set_last_user_message fun(message: OpencodeMessage|nil)
----@field set_pending_permissions fun(permissions: OpencodePermission[])
----@field set_cost fun(cost: number)
----@field set_tokens_count fun(count: number)
-
 local M = {}
 
 ---@param messages OpencodeMessage[]|nil
@@ -31,6 +23,10 @@ function M.set_pending_permissions(permissions)
   return store.set('pending_permissions', permissions)
 end
 
+---@param mutator fun(current_permissions: OpencodePermission[]): nil
+function M.update_pending_permissions(mutator)
+  return store.mutate('pending_permissions', mutator)
+end
 ---@param cost number
 function M.set_cost(cost)
   return store.set('cost', cost)
@@ -41,4 +37,24 @@ 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)
+  return store.batch(function()
+    store.set('tokens_count', tokens_count)
+    store.set('cost', cost)
+  end)
+end
+
+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', {})
+  end)
+end
+
 return M