feat(store): add batched state updates and mutate API

sudo-tee · sudo-tee · commit 5429155c5735 · 2026-03-18T12:54:42.000-04:00
diff --git a/lua/opencode/context.lua b/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
@@ -200,12 +196,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
diff --git a/lua/opencode/core.lua b/lua/opencode/core.lua
@@ -92,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()
diff --git a/lua/opencode/server_job.lua b/lua/opencode/server_job.lua
@@ -293,6 +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.store.set_raw('opencode_server', state.opencode_server)
         state.opencode_server.port = port_num
         local server_pid = job and job.pid
         port_mapping.register(port_num, vim.fn.getcwd(), true, 'serve', nil, server_pid)
diff --git a/lua/opencode/state.lua b/lua/opencode/state.lua
@@ -1 +1 @@
-return require('opencode.state.init') --[[@as OpencodeStateModule]]
+return require('opencode.state.init') --[[@as OpencodeState]]
diff --git a/lua/opencode/state/context.lua b/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)
diff --git a/lua/opencode/state/init.lua b/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,
diff --git a/lua/opencode/state/jobs.lua b/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
diff --git a/lua/opencode/state/model.lua b/lua/opencode/state/model.lua
@@ -1,17 +1,6 @@
 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
@@ -33,9 +22,11 @@ function M.clear_model()
 end
 
 function M.clear()
-  store.set('current_model', nil)
-  store.set('current_mode', nil)
-  store.set('current_variant', nil)
+  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
diff --git a/lua/opencode/state/renderer.lua b/lua/opencode/state/renderer.lua
@@ -1,16 +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)
----@field set_stats fun(tokens_count: number, cost: number)
----@field reset fun()
-
 local M = {}
 
 ---@param messages OpencodeMessage[]|nil
@@ -33,6 +23,10 @@ function M.set_pending_permissions(permissions)
   return store.set('pending_permissions', permissions)
 end
 
+---@param mutator fun(current_permissions: OpencodePermission[]): OpencodePermission[]
+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)
@@ -46,16 +40,21 @@ 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)
+  return store.batch(function()
+    store.set('tokens_count', tokens_count)
+    store.set('cost', cost)
+  end)
 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', {})
+  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
diff --git a/lua/opencode/state/session.lua b/lua/opencode/state/session.lua
@@ -1,28 +1,25 @@
 local store = require('opencode.state.store')
 
 ---@class OpencodeSessionStateMutations
----@field set_active fun(session: Session|nil)
----@field clear_active fun()
----@field set_restore_points fun(points: RestorePoint[])
----@field reset_restore_points fun()
----@field set_last_sent_context fun(context: OpencodeContext|nil)
----@field set_user_message_count fun(count: table<string, number>)
-
 local M = {}
 
 ---@param session Session|nil
 function M.set_active(session)
-  store.set('restore_points', {})
-  store.set('last_sent_context', nil)
-  store.set('user_message_count', {})
-  return store.set('active_session', session)
+  return store.batch(function()
+    store.set('restore_points', {})
+    store.set('last_sent_context', nil)
+    store.set('user_message_count', {})
+    return store.set('active_session', session)
+  end)
 end
 
 function M.clear_active()
-  M.reset_restore_points()
-  M.set_last_sent_context()
-  M.set_user_message_count({})
-  return store.set('active_session', nil)
+  return store.batch(function()
+    store.set('restore_points', {})
+    store.set('last_sent_context', nil)
+    store.set('user_message_count', {})
+    return store.set('active_session', nil)
+  end)
 end
 
 ---@param points RestorePoint[]
diff --git a/lua/opencode/state/store.lua b/lua/opencode/state/store.lua
@@ -87,6 +87,54 @@ local _state = {
 }
 
 local _listeners = {}
+local _batch_depth = 0
+local _batched_changes = {}
+local _batched_order = {}
+
+---@generic K extends keyof OpencodeStateData
+---@param key K
+---@param new_val std.RawGet<OpencodeStateData, K>
+---@param old_val std.RawGet<OpencodeStateData, K>
+local function queue_emit(key, new_val, old_val)
+  if vim.deep_equal(old_val, new_val) then
+    return false
+  end
+
+  if _batch_depth == 0 then
+    M.emit(key, new_val, old_val)
+    return true
+  end
+
+  if not _batched_changes[key] then
+    _batched_changes[key] = {
+      old_val = old_val,
+      new_val = new_val,
+    }
+    table.insert(_batched_order, key)
+    return true
+  end
+
+  _batched_changes[key].new_val = new_val
+  return true
+end
+
+local function flush_batched_emits()
+  if _batch_depth > 0 or #_batched_order == 0 then
+    return
+  end
+
+  local pending_changes = _batched_changes
+  local pending_order = _batched_order
+  _batched_changes = {}
+  _batched_order = {}
+
+  for _, key in ipairs(pending_order) do
+    local change = pending_changes[key]
+    if change then
+      M.emit(key, change.new_val, change.old_val)
+    end
+  end
+end
 
 function M.state()
   return _state
@@ -107,9 +155,7 @@ function M.set(key, value)
   local old = _state[key]
 
   _state[key] = value
-  if not vim.deep_equal(old, value) then
-    M.emit(key, value, old)
-  end
+  queue_emit(key, value, old)
 
   return value
 end
@@ -119,7 +165,8 @@ end
 ---@param value std.RawGet<OpencodeStateData, K>
 ---@return std.RawGet<OpencodeStateData, K>
 function M.set_raw(key, value)
-  return M.set(key, value)
+  _state[key] = value
+  return value
 end
 
 ---@generic K extends keyof OpencodeStateData
@@ -132,6 +179,43 @@ function M.update(key, updater)
   return next_value
 end
 
+---@param callback fun(store: OpencodeStateStore)
+function M.batch(callback)
+  _batch_depth = _batch_depth + 1
+  local ok, result = pcall(callback, M)
+  _batch_depth = _batch_depth - 1
+
+  if _batch_depth == 0 then
+    flush_batched_emits()
+  end
+
+  if not ok then
+    error(result, 0)
+  end
+
+  return result
+end
+
+---@generic K extends keyof OpencodeStateData
+---@param key K
+---@param mutator fun(current: std.RawGet<OpencodeStateData, K>):nil
+---@return std.RawGet<OpencodeStateData, K>
+function M.mutate(key, mutator)
+  if _state[key] == nil then
+    _state[key] = {} --[[@as std.RawGet<OpencodeStateData, K>]]
+  end
+
+  if type(_state[key]) ~= 'table' then
+    error('State key is not a table: ' .. key)
+  end
+
+  local current = _state[key]
+  local old = vim.tbl_extend('force', {}, current)
+  mutator(current)
+  queue_emit(key, current, old)
+  return current
+end
+
 ---@generic K extends keyof OpencodeStateData
 ---@param key K|K[]|nil
 ---@param cb fun(key:K, new_val:std.RawGet<OpencodeStateData, K>, old_val:std.RawGet<OpencodeStateData, K>)
@@ -204,16 +288,10 @@ function M.append(key, value)
   if type(value) ~= 'table' then
     error('Value must be a table to append')
   end
-  if not _state[key] then
-    _state[key] = {}
-  end
-  if type(_state[key]) ~= 'table' then
-    error('State key is not a table: ' .. key)
-  end
 
-  local old = vim.deepcopy(_state[key] --[[@as table]])
-  table.insert(_state[key] --[[@as table]], value)
-  M.emit(key, _state[key], old)
+  M.mutate(key, function(current)
+    table.insert(current --[[@as table]], value)
+  end)
 end
 
 ---@generic K extends keyof OpencodeStateData
@@ -223,13 +301,10 @@ function M.remove(key, idx)
   if not _state[key] then
     return
   end
-  if type(_state[key]) ~= 'table' then
-    error('State key is not a table: ' .. key)
-  end
 
-  local old = vim.deepcopy(_state[key] --[[@as table]])
-  table.remove(_state[key] --[[@as table]], idx)
-  M.emit(key, _state[key], old)
+  M.mutate(key, function(current)
+    table.remove(current --[[@as table]], idx)
+  end)
 end
 
 return M
diff --git a/lua/opencode/state/ui.lua b/lua/opencode/state/ui.lua
diff --git a/lua/opencode/ui/autocmds.lua b/lua/opencode/ui/autocmds.lua
diff --git a/lua/opencode/ui/renderer.lua b/lua/opencode/ui/renderer.lua
diff --git a/lua/opencode/ui/ui.lua b/lua/opencode/ui/ui.lua
diff --git a/tests/unit/state_spec.lua b/tests/unit/state_spec.lua

Original file line number	Diff line number	Diff line change
`@@ -1 +1 @@`
`1`		`-return require('opencode.state.init') --[[@as OpencodeStateModule]]`
	`1`	`+return require('opencode.state.init') --[[@as OpencodeState]]`