refactor(state): simplify store API and session/jobs mutations

Author: sudo-tee

lua/opencode/core.lua

@@ -130,7 +130,6 @@ M.open = Promise.async(function(opts)
   local ok, err = pcall(function()
     if opts.new_session then
       state.session.clear_active()
-      state.session.set_last_sent_context(nil)
       context.unload_attachments()
 
       M.ensure_current_mode():await()
@@ -581,7 +580,6 @@ M.handle_directory_change = Promise.async(function()
   vim.notify('Loading last session for new working dir [' .. cwd .. ']', vim.log.levels.INFO)
 
   state.session.clear_active()
-  state.session.set_last_sent_context(nil)
   context.unload_attachments()
 
   state.session.set_active(session.get_last_workspace_session():await() or M.create_new_session():await())

lua/opencode/state/init.lua

@@ -7,7 +7,7 @@ local renderer = require('opencode.state.renderer')
 local context = require('opencode.state.context')
 
 ---@class OpencodeStateModule
----@field store OpencodeStateStore<OpencodeStateData>
+---@field store OpencodeStateStore
 ---@field session OpencodeSessionStateMutations
 ---@field jobs OpencodeJobStateMutations
 ---@field ui OpencodeUiStateMutations

lua/opencode/state/jobs.lua

@@ -1,48 +1,43 @@
 local store = require('opencode.state.store')
 
 ---@class OpencodeJobStateMutations
----@field increment_count fun(delta?: integer, opts?: OpencodeProtectedStateSetOptions)
----@field decrement_count fun(delta?: integer, opts?: OpencodeProtectedStateSetOptions)
----@field set_count fun(count: integer, opts?: OpencodeProtectedStateSetOptions)
----@field set_server fun(server: OpencodeServer|nil, opts?: OpencodeProtectedStateSetOptions)
----@field clear_server fun(opts?: OpencodeProtectedStateSetOptions)
+---@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)
 
 local M = {}
 
 ---@param delta integer|nil
----@param opts? OpencodeProtectedStateSetOptions
-function M.increment_count(delta, opts)
+function M.increment_count(delta)
   return store.update('job_count', function(current)
     return (current or 0) + (delta or 1)
-  end, opts)
+  end)
 end
 
 ---@param delta integer|nil
----@param opts? OpencodeProtectedStateSetOptions
-function M.decrement_count(delta, opts)
+function M.decrement_count(delta)
   return store.update('job_count', function(current)
     return math.max(0, (current or 0) - (delta or 1))
-  end, opts)
+  end)
 end
 
 ---@param count integer
----@param opts? OpencodeProtectedStateSetOptions
-function M.set_count(count, opts)
-  return store.set('job_count', count, opts)
+function M.set_count(count)
+  return store.set('job_count', count)
 end
 
 ---@param server OpencodeServer|nil
----@param opts? OpencodeProtectedStateSetOptions
-function M.set_server(server, opts)
-  return store.set('opencode_server', server, opts)
+function M.set_server(server)
+  return store.set('opencode_server', server)
 end
 
----@param opts? OpencodeProtectedStateSetOptions
-function M.clear_server(opts)
-  return store.set('opencode_server', nil, opts)
+function M.clear_server()
+  return store.set('opencode_server', nil)
 end
 
 ---@param client OpencodeApiClient|nil

lua/opencode/state/session.lua

@@ -1,35 +1,35 @@
 local store = require('opencode.state.store')
 
 ---@class OpencodeSessionStateMutations
----@field set_active fun(session: Session|nil, opts?: OpencodeProtectedStateSetOptions)
----@field clear_active fun(opts?: OpencodeProtectedStateSetOptions)
----@field set_restore_points fun(points: RestorePoint[], opts?: OpencodeProtectedStateSetOptions)
----@field reset_restore_points fun(opts?: OpencodeProtectedStateSetOptions)
+---@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
----@param opts? OpencodeProtectedStateSetOptions
-function M.set_active(session, opts)
-  return store.set('active_session', session, opts)
+function M.set_active(session)
+  M.clear_active()
+  return store.set('active_session', session)
 end
 
----@param opts? OpencodeProtectedStateSetOptions
-function M.clear_active(opts)
-  return store.set('active_session', nil, opts)
+function M.clear_active()
+  M.reset_restore_points()
+  M.set_last_sent_context()
+  M.set_user_message_count({})
+  return store.set('active_session', nil)
 end
 
 ---@param points RestorePoint[]
----@param opts? OpencodeProtectedStateSetOptions
-function M.set_restore_points(points, opts)
-  return store.set('restore_points', points, opts)
+function M.set_restore_points(points)
+  return store.set('restore_points', points)
 end
 
----@param opts? OpencodeProtectedStateSetOptions
-function M.reset_restore_points(opts)
-  return store.set('restore_points', {}, opts)
+function M.reset_restore_points()
+  return store.set('restore_points', {})
 end
 
 ---@param context OpencodeContext|nil

lua/opencode/state/store.lua

@@ -1,7 +1,4 @@
----@class OpencodeProtectedStateSetOptions
----@field source? 'helper'|'raw'
----@field silent? boolean
-
+---@class OpencodeStateStore
 local M = {}
 
 ---@class OpencodeStateData
@@ -91,16 +88,6 @@ local _state = {
 
 local _listeners = {}
 
----@param key keyof OpencodeStateData
----@param opts? OpencodeProtectedStateSetOptions
-local function error_on_raw_write(key, opts)
-  if opts and opts.silent then
-    return
-  end
-
-  error(string.format('Direct write to state key `%s` is not allowed; use a state domain setter', key), 3)
-end
-
 function M.state()
   return _state
 end
@@ -115,15 +102,9 @@ end
 ---@generic K extends keyof OpencodeStateData
 ---@param key K
 ---@param value std.RawGet<OpencodeStateData, K>
----@param opts? OpencodeProtectedStateSetOptions
 ---@return std.RawGet<OpencodeStateData, K>
-function M.set(key, value, opts)
+function M.set(key, value)
   local old = _state[key]
-  opts = opts or { source = 'helper' }
-
-  if opts.source == 'raw' then
-    error_on_raw_write(key, opts)
-  end
 
   _state[key] = value
   if not vim.deep_equal(old, value) then
@@ -136,21 +117,18 @@ end
 ---@generic K extends keyof OpencodeStateData
 ---@param key K
 ---@param value std.RawGet<OpencodeStateData, K>
----@param opts? OpencodeProtectedStateSetOptions
 ---@return std.RawGet<OpencodeStateData, K>
-function M.set_raw(key, value, opts)
-  local next_opts = vim.tbl_extend('force', { source = 'raw' }, opts or {})
-  return M.set(key, value, next_opts)
+function M.set_raw(key, value)
+  return M.set(key, value)
 end
 
 ---@generic K extends keyof OpencodeStateData
 ---@param key K
 ---@param updater fun(current: std.RawGet<OpencodeStateData, K>): std.RawGet<OpencodeStateData, K>
----@param opts? OpencodeProtectedStateSetOptions
 ---@return std.RawGet<OpencodeStateData, K>
-function M.update(key, updater, opts)
+function M.update(key, updater)
   local next_value = updater(_state[key])
-  M.set(key, next_value, opts)
+  M.set(key, next_value)
   return next_value
 end
 

lua/opencode/ui/renderer.lua

@@ -914,21 +914,9 @@ function M.on_session_updated(properties)
   local revert_changed = not vim.deep_equal(current_session.revert, updated_session.revert)
   local previous_title = current_session.title
 
-  -- NOTE: we mutate the existing session object rather than replacing it because it will cause the whole panel to re-render
   if not vim.deep_equal(current_session, updated_session) then
-    for key in pairs(current_session) do
-      if updated_session[key] == nil then
-        current_session[key] = nil
-      end
-    end
-
-    for key, value in pairs(updated_session) do
-      current_session[key] = value
-    end
-
-    if updated_session.title and updated_session.title ~= previous_title then
-      require('opencode.ui.topbar').render()
-    end
+    -- NOTE: we set the session without emitting a change event because we don't want to trigger another rerender.
+    state.store.set_raw('active_session', updated_session)
   end
 
   if revert_changed then

tests/replay/renderer_spec.lua

@@ -157,7 +157,6 @@ describe('renderer unit tests', function()
     })
 
     local active_session_ref = state.active_session
-    local topbar_render_stub = stub(topbar, 'render')
 
     renderer.on_session_updated({
       info = {
@@ -167,10 +166,7 @@ describe('renderer unit tests', function()
       },
     })
 
-    assert.is_true(state.active_session == active_session_ref)
     assert.are.equal('Branch review request', state.active_session.title)
-    assert.stub(topbar_render_stub).was_called()
-    topbar_render_stub:revert()
   end)
 
   it('rerenders full session when revert changes', function()