feat: Use per sessions user messages

Author: guillaumeboehm

lua/opencode/core.lua

@@ -147,16 +147,24 @@ function M.send_message(prompt, opts)
   params.parts = context.format_message(prompt, opts.context)
   M.before_run(opts)
 
-  state.user_message_count = state.user_message_count + 1
+  -- Capture the session ID to ensure we track the message count for the correct session
+  local session_id = state.active_session.id
+  local sent_message_count = vim.deepcopy(state.user_message_count)
+  sent_message_count[session_id] = (sent_message_count[session_id] or 0) + 1
+  state.user_message_count = sent_message_count
+
   state.api_client
-    :create_message(state.active_session.id, params)
+    :create_message(session_id, params)
     :and_then(function(response)
       if not response or not response.info or not response.parts then
         -- fall back to full render. incremental render is handled
         -- event manager
         ui.render_output()
       end
-      state.user_message_count = state.user_message_count - 1
+
+      local received_message_count = vim.deepcopy(state.user_message_count)
+      received_message_count[response.info.sessionID] = (received_message_count[response.info.sessionID] ~= nil) and (received_message_count[response.info.sessionID] - 1) or 0
+      state.user_message_count = received_message_count
 
       M.after_run(prompt)
     end)
@@ -369,17 +377,26 @@ function M.initialize_current_model()
   return state.current_model
 end
 
-local function on_user_message_count_change(_, new, old)
-  local done_thinking = new == 0 and old > 0
-  if config.hooks and config.hooks.on_done_thinking and done_thinking then
-    pcall(config.hooks.on_done_thinking, state.active_session)
+function M._on_user_message_count_change(_, new, old)
+  if config.hooks and config.hooks.on_done_thinking then
+      local all_sessions = session.get_all_workspace_sessions() or {}
+      local done_sessions = vim.tbl_filter(function(s)
+        local msg_count = new[s.id] or 0
+        local old_msg_count = (old and old[s.id]) or 0
+        return msg_count == 0 and old_msg_count > 0
+      end, all_sessions)
+
+      for _, done_session in ipairs(done_sessions) do
+        pcall(config.hooks.on_done_thinking, done_session)
+      end
   end
 end
 
-local function on_current_permission_change(_, new, old)
+function M._on_current_permission_change(_, new, old)
   local permission_requested = old == nil and new ~= nil
   if config.hooks and config.hooks.on_permission_requested and permission_requested then
-    pcall(config.hooks.on_permission_requested, state.active_session)
+    local local_session = session.get_by_id(state.active_session.id) or {}
+    pcall(config.hooks.on_permission_requested, local_session)
   end
 end
 
@@ -391,8 +408,8 @@ end
 
 function M.setup()
   state.subscribe('opencode_server', on_opencode_server)
-  state.subscribe('user_message_count', on_user_message_count_change)
-  state.subscribe('current_permission', on_current_permission_change)
+  state.subscribe('user_message_count', M._on_user_message_count_change)
+  state.subscribe('current_permission', M._on_current_permission_change)
 
   vim.schedule(function()
     M.opencode_ok()

lua/opencode/state.lua

@@ -32,7 +32,7 @@
 ---@field cost number
 ---@field tokens_count number
 ---@field job_count number
----@field user_message_count number
+---@field user_message_count table<string, number>
 ---@field opencode_server OpencodeServer|nil
 ---@field api_client OpencodeApiClient
 ---@field event_manager EventManager|nil
@@ -81,7 +81,7 @@ local _state = {
   tokens_count = 0,
   -- job
   job_count = 0,
-  user_message_count = 0,
+  user_message_count = {},
   opencode_server = nil,
   api_client = nil,
   event_manager = nil,

tests/unit/hooks_spec.lua

@@ -1,6 +1,7 @@
 local renderer = require('opencode.ui.renderer')
 local config = require('opencode.config')
 local state = require('opencode.state')
+local core = require('opencode.core')
 local helpers = require('tests.helpers')
 local ui = require('opencode.ui.ui')
 
@@ -115,23 +116,44 @@ describe('hooks', function()
   describe('on_done_thinking', function()
     it('should call hook when thinking is done', function()
       local called = false
+      local called_session = nil
 
-      config.hooks.on_done_thinking = function()
+      config.hooks.on_done_thinking = function(session)
         called = true
+        called_session = session
+      end
+
+      -- Mock session.get_all_workspace_sessions to return our test session
+      local session_module = require('opencode.session')
+      local original_get_all = session_module.get_all_workspace_sessions
+      session_module.get_all_workspace_sessions = function()
+        return { { id = 'test-session', title = 'Test' } }
       end
 
-      -- Simulate job count change from 1 to 0 (done thinking)
-      state.user_message_count = 1
-      state.user_message_count = 0
+      state.subscribe('user_message_count', core._on_user_message_count_change)
+
+      -- Simulate job count change from 1 to 0 (done thinking) for a specific session
+      state.active_session = { id = 'test-session', title = 'Test' }
+      state.user_message_count = { ['test-session'] = 1 }
+      state.user_message_count = { ['test-session'] = 0 }
+
+      -- Wait for async notification
+      vim.wait(100, function() return called end)
+
+      -- Restore original function
+      session_module.get_all_workspace_sessions = original_get_all
+      state.unsubscribe('user_message_count', core._on_user_message_count_change)
 
       assert.is_true(called)
+      assert.are.equal(called_session.id, 'test-session')
     end)
 
     it('should not error when hook is nil', function()
       config.hooks.on_done_thinking = nil
-      state.user_message_count = 1
+      state.active_session = { id = 'test-session', title = 'Test' }
+      state.user_message_count = { ['test-session'] = 1 }
       assert.has_no.errors(function()
-        state.user_message_count = 0
+        state.user_message_count = { ['test-session'] = 0 }
       end)
     end)
 
@@ -140,26 +162,48 @@ describe('hooks', function()
         error('test error')
       end
 
-      state.user_message_count = 1
+      state.active_session = { id = 'test-session', title = 'Test' }
+      state.user_message_count = { ['test-session'] = 1 }
       assert.has_no.errors(function()
-        state.user_message_count = 0
+        state.user_message_count = { ['test-session'] = 0 }
       end)
     end)
   end)
 
   describe('on_permission_requested', function()
     it('should call hook when permission is requested', function()
       local called = false
+      local called_session = nil
 
-      config.hooks.on_permission_requested = function()
+      config.hooks.on_permission_requested = function(session)
         called = true
+        called_session = session
+      end
+
+      -- Mock session.get_by_id to return our test session
+      local session_module = require('opencode.session')
+      local original_get_by_id = session_module.get_by_id
+      session_module.get_by_id = function(id)
+        return { id = id, title = 'Test' }
       end
 
+      -- Set up the subscription manually
+      state.subscribe('current_permission', core._on_current_permission_change)
+
       -- Simulate permission change from nil to a value
+      state.active_session = { id = 'test-session', title = 'Test' }
       state.current_permission = nil
       state.current_permission = { tool = 'test_tool', action = 'read' }
 
+      -- Wait for async notification
+      vim.wait(100, function() return called end)
+
+      -- Restore original function
+      session_module.get_by_id = original_get_by_id
+      state.unsubscribe('current_permission', core._on_current_permission_change)
+
       assert.is_true(called)
+      assert.are.equal(called_session.id, 'test-session')
     end)
 
     it('should not error when hook is nil', function()