refactor(state): centralize protected state mutations into domain setters

Author: sudo-tee

lua/opencode/context.lua

@@ -22,11 +22,12 @@ local toggleable_context_keys = {
 ---@param context_key OpencodeToggleableContextKey
 ---@return table
 local function ensure_context_state(context_key)
-  state.current_context_config = state.current_context_config or {}
-  local current = state.current_context_config[context_key]
+  local current_config = state.current_context_config or {}
+  local current = current_config[context_key]
   local defaults = vim.tbl_get(config, 'context', context_key) or {}
-  state.current_context_config[context_key] = vim.tbl_deep_extend('force', {}, defaults, current or {})
-  return state.current_context_config[context_key]
+  current_config[context_key] = vim.tbl_deep_extend('force', {}, defaults, current or {})
+  state.context.set_current_context_config(current_config)
+  return current_config[context_key]
 end
 
 M.ChatContext = ChatContext
@@ -117,12 +118,12 @@ end
 -- Delegate global state management to ChatContext
 function M.add_selection(selection)
   ChatContext.add_selection(selection)
-  state.context_updated_at = vim.uv.now()
+  state.context.set_context_updated_at(vim.uv.now())
 end
 
 function M.remove_selection(selection)
   ChatContext.remove_selection(selection)
-  state.context_updated_at = vim.uv.now()
+  state.context.set_context_updated_at(vim.uv.now())
 end
 
 function M.clear_selections()
@@ -180,13 +181,13 @@ function M.add_file(file)
 
   file = vim.fn.fnamemodify(file, ':p')
   ChatContext.add_file(file)
-  state.context_updated_at = vim.uv.now()
+  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_updated_at = vim.uv.now()
+  state.context.set_context_updated_at(vim.uv.now())
 end
 
 function M.clear_files()
@@ -195,12 +196,12 @@ end
 
 function M.add_subagent(subagent)
   ChatContext.add_subagent(subagent)
-  state.context_updated_at = vim.uv.now()
+  state.context.set_context_updated_at(vim.uv.now())
 end
 
 function M.remove_subagent(subagent)
   ChatContext.remove_subagent(subagent)
-  state.context_updated_at = vim.uv.now()
+  state.context.set_context_updated_at(vim.uv.now())
 end
 
 function M.clear_subagents()
@@ -213,7 +214,7 @@ end
 
 function M.load()
   ChatContext.load()
-  state.context_updated_at = vim.uv.now()
+  state.context.set_context_updated_at(vim.uv.now())
 end
 
 -- Context creation with delta logic (delegates to ChatContext)

lua/opencode/context/chat_context.lua

@@ -175,7 +175,7 @@ function M.add_selection(selection)
   end
 
   table.insert(M.context.selections, selection)
-  state.context_updated_at = vim.uv.now()
+  state.context.set_context_updated_at(vim.uv.now())
 end
 
 function M.remove_selection(selection)
@@ -190,12 +190,12 @@ function M.remove_selection(selection)
       break
     end
   end
-  state.context_updated_at = vim.uv.now()
+  state.context.set_context_updated_at(vim.uv.now())
 end
 
 function M.clear_selections()
   M.context.selections = {}
-  state.context_updated_at = vim.uv.now()
+  state.context.set_context_updated_at(vim.uv.now())
 end
 
 function M.add_file(file)
@@ -210,7 +210,7 @@ function M.add_file(file)
   if not vim.tbl_contains(M.context.mentioned_files, file) then
     table.insert(M.context.mentioned_files, file)
   end
-  state.context_updated_at = vim.uv.now()
+  state.context.set_context_updated_at(vim.uv.now())
 end
 
 function M.remove_file(file)
@@ -226,12 +226,12 @@ function M.remove_file(file)
       break
     end
   end
-  state.context_updated_at = vim.uv.now()
+  state.context.set_context_updated_at(vim.uv.now())
 end
 
 function M.clear_files()
   M.context.mentioned_files = {}
-  state.context_updated_at = vim.uv.now()
+  state.context.set_context_updated_at(vim.uv.now())
 end
 
 function M.add_subagent(subagent)
@@ -243,7 +243,7 @@ function M.add_subagent(subagent)
   if not vim.tbl_contains(M.context.mentioned_subagents, subagent) then
     table.insert(M.context.mentioned_subagents, subagent)
   end
-  state.context_updated_at = vim.uv.now()
+  state.context.set_context_updated_at(vim.uv.now())
 end
 
 function M.remove_subagent(subagent)
@@ -258,18 +258,18 @@ function M.remove_subagent(subagent)
       break
     end
   end
-  state.context_updated_at = vim.uv.now()
+  state.context.set_context_updated_at(vim.uv.now())
 end
 
 function M.clear_subagents()
   M.context.mentioned_subagents = {}
-  state.context_updated_at = vim.uv.now()
+  state.context.set_context_updated_at(vim.uv.now())
 end
 
 function M.unload_attachments()
   M.context.mentioned_files = {}
   M.context.selections = {}
-  state.context_updated_at = vim.uv.now()
+  state.context.set_context_updated_at(vim.uv.now())
 end
 
 function M.get_mentioned_files()
@@ -402,7 +402,7 @@ function M.load()
     or not vim.deep_equal(prev_cursor_data, M.context.cursor_data)
     or not vim.deep_equal(prev_linter_errors, M.context.linter_errors)
   then
-    state.context_updated_at = vim.uv.now()
+    state.context.set_context_updated_at(vim.uv.now())
   end
 
   -- Handle current selection
@@ -471,7 +471,7 @@ function M.delta_context(opts)
     end
   end
 
-  state.context_updated_at = vim.uv.now()
+  state.context.set_context_updated_at(vim.uv.now())
   return ctx
 end
 

lua/opencode/core.lua

@@ -72,7 +72,7 @@ M.check_cwd = function()
       'CWD changed since last check, resetting session and context',
       { current_cwd = state.current_cwd, new_cwd = vim.fn.getcwd() }
     )
-    state.current_cwd = vim.fn.getcwd()
+    state.context.set_current_cwd(vim.fn.getcwd())
     state.session.clear_active()
     context.unload_attachments()
   end
@@ -130,7 +130,7 @@ M.open = Promise.async(function(opts)
   local ok, err = pcall(function()
     if opts.new_session then
       state.session.clear_active()
-      state.last_sent_context = nil
+      state.session.set_last_sent_context(nil)
       context.unload_attachments()
 
       M.ensure_current_mode():await()
@@ -187,7 +187,7 @@ M.send_message = Promise.async(function(prompt, opts)
   opts = opts or {}
 
   opts.context = vim.tbl_deep_extend('force', state.current_context_config or {}, opts.context or {})
-  state.current_context_config = opts.context
+  state.context.set_current_context_config(opts.context)
   context.load()
   opts.model = opts.model or M.initialize_current_model():await()
   opts.agent = opts.agent or state.current_mode or config.default_mode
@@ -223,7 +223,7 @@ M.send_message = Promise.async(function(prompt, opts)
     local sent_message_count = vim.deepcopy(state.user_message_count)
     local new_value = (sent_message_count[session_id] or 0) + num
     sent_message_count[session_id] = new_value >= 0 and new_value or 0
-    state.user_message_count = sent_message_count
+    state.session.set_user_message_count(sent_message_count)
   end
 
   update_sent_message_count(1)
@@ -267,7 +267,7 @@ end)
 ---@param prompt string
 function M.after_run(prompt)
   context.unload_attachments()
-  state.last_sent_context = vim.deepcopy(context.get_context())
+  state.session.set_last_sent_context(vim.deepcopy(context.get_context()))
   context.delta_context()
   require('opencode.history').write(prompt)
   M._abort_count = 0
@@ -437,7 +437,7 @@ M.opencode_ok = Promise.async(function()
   if not state.opencode_cli_version or state.opencode_cli_version == '' then
     local result = Promise.system({ config.opencode_executable, '--version' }):await()
     local out = (result and result.stdout or ''):gsub('%s+$', '')
-    state.opencode_cli_version = out:match('(%d+%%.%d+%%.%d+)') or out
+    state.jobs.set_opencode_cli_version(out:match('(%d+%%.%d+%%.%d+)') or out)
   end
 
   local required = state.required_version
@@ -581,7 +581,7 @@ 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.last_sent_context = nil
+  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())
@@ -615,7 +615,7 @@ function M.setup()
     M.opencode_ok()
   end)
   local OpencodeApiClient = require('opencode.api_client')
-  state.api_client = OpencodeApiClient.create()
+  state.jobs.set_api_client(OpencodeApiClient.create())
 end
 
 return M

lua/opencode/event_manager.lua

@@ -593,8 +593,9 @@ function EventManager:get_subscriber_count(event_name)
 end
 
 function EventManager.setup()
-  state.event_manager = EventManager.new()
-  state.event_manager:start()
+  local manager = EventManager.new()
+  state.jobs.set_event_manager(manager)
+  manager:start()
 end
 
 return EventManager

lua/opencode/image_handler.lua

@@ -150,7 +150,7 @@ function M.paste_image_from_clipboard()
 
   if success then
     context.add_file(image_path)
-    state.context_updated_at = os.time()
+    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/state/context.lua

@@ -0,0 +1,20 @@
+local store = require('opencode.state.store')
+
+local M = {}
+
+---@param config OpencodeContextConfig|nil
+function M.set_current_context_config(config)
+  return store.set('current_context_config', config)
+end
+
+---@param timestamp number|nil
+function M.set_context_updated_at(timestamp)
+  return store.set('context_updated_at', timestamp)
+end
+
+---@param cwd string|nil
+function M.set_current_cwd(cwd)
+  return store.set('current_cwd', cwd)
+end
+
+return M

lua/opencode/state/init.lua

@@ -28,13 +28,18 @@
 ---@field clear_active fun(opts?: OpencodeProtectedStateSetOptions)
 ---@field set_restore_points fun(points: RestorePoint[], opts?: OpencodeProtectedStateSetOptions)
 ---@field reset_restore_points fun(opts?: OpencodeProtectedStateSetOptions)
+---@field set_last_sent_context fun(context: OpencodeContext|nil)
+---@field set_user_message_count fun(count: table<string, number>)
 
 ---@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 set_api_client fun(client: OpencodeApiClient|nil)
+---@field set_event_manager fun(manager: EventManager|nil)
+---@field set_opencode_cli_version fun(version: string|nil)
 
 ---@class OpencodeUiStateMutations
 ---@field set_windows fun(windows: OpencodeWindowState|nil)
@@ -48,6 +53,9 @@
 ---@field set_current_code_buf fun(bufnr: integer|nil)
 ---@field set_last_window_width_ratio fun(ratio: number|nil)
 ---@field clear_last_window_width_ratio fun()
+---@field set_input_content fun(lines: table)
+---@field set_saved_window_options fun(opts: table|nil)
+---@field set_pre_zoom_width fun(width: integer|nil)
 
 ---@class OpencodeModelStateMutations
 ---@field set_mode fun(mode: string|nil)
@@ -60,6 +68,19 @@
 ---@field set_mode_model_map fun(mode_map: table<string, string>)
 ---@field set_mode_model_override fun(mode: string, model: string)
 
+---@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)
+
+---@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)
+
 ---@class OpencodeState
 ---@field windows OpencodeWindowState|nil
 ---@field is_opening boolean
@@ -109,29 +130,30 @@
 ---@field jobs OpencodeJobStateMutations
 ---@field ui OpencodeUiStateMutations
 ---@field model OpencodeModelStateMutations
+---@field renderer OpencodeRendererStateMutations
+---@field context OpencodeContextStateMutations
 
 local store = require('opencode.state.store')
 local session = require('opencode.state.session')
 local jobs = require('opencode.state.jobs')
 local ui = require('opencode.state.ui')
 local model = require('opencode.state.model')
-local test_helpers = require('opencode.state.test_helpers')
+local renderer = require('opencode.state.renderer')
+local context = require('opencode.state.context')
 
 local M = {
   store = store,
   session = session,
   jobs = jobs,
   ui = ui,
   model = model,
-  test_helpers = test_helpers,
+  renderer = renderer,
+  context = context,
   subscribe = store.subscribe,
   unsubscribe = store.unsubscribe,
   notify = store.notify,
   append = store.append,
   remove = store.remove,
-  set_raw = store.set_raw,
-  allow_raw_writes_for_tests = test_helpers.allow_raw_writes_for_tests,
-  silence_protected_writes = test_helpers.silence_protected_writes,
 }
 
 function M.is_running()
@@ -142,8 +164,8 @@ return setmetatable(M, {
   __index = function(_, key)
     return store.get(key)
   end,
-  __newindex = function(_, key, value)
-    store.set(key, value, { source = 'raw' })
+  __newindex = function(_, key, _value)
+    error(string.format('Direct write to state key `%s` is not allowed; use a state domain setter', key), 2)
   end,
   __pairs = function()
     return pairs(store.state())

lua/opencode/state/jobs.lua

@@ -35,4 +35,19 @@ function M.clear_server(opts)
   return store.set('opencode_server', nil, opts)
 end
 
+---@param client OpencodeApiClient|nil
+function M.set_api_client(client)
+  return store.set('api_client', client)
+end
+
+---@param manager EventManager|nil
+function M.set_event_manager(manager)
+  return store.set('event_manager', manager)
+end
+
+---@param version string|nil
+function M.set_opencode_cli_version(version)
+  return store.set('opencode_cli_version', version)
+end
+
 return M