fix(context): clear prompt attachments and selections reliably

Author: sudo-tee

lua/opencode/context/chat_context.lua

@@ -14,6 +14,23 @@ M.context = {
   linter_errors = nil,
 }
 
+local cleared_selections = {}
+local cleared_selections_context = nil
+
+---@param left OpencodeContextSelection|nil
+---@param right OpencodeContextSelection|nil
+---@return boolean
+local function is_same_selection(left, right)
+  return (
+    left
+    and right
+    and left.file
+    and right.file
+    and left.file.path == right.file.path
+    and left.lines == right.lines
+  ) == true
+end
+
 ---@param path string
 ---@param prompt? string
 ---@return OpencodeMessagePart
@@ -162,6 +179,9 @@ end
 -- Global context management functions
 
 function M.add_selection(selection)
+  cleared_selections = {}
+  cleared_selections_context = nil
+
   -- Ensure selections is always a table
   if not M.context.selections then
     M.context.selections = {}
@@ -266,7 +286,22 @@ function M.clear_subagents()
   state.context.set_context_updated_at(vim.uv.now())
 end
 
-function M.unload_attachments()
+---@param selections? OpencodeContextSelection[]
+function M.unload_attachments(selections)
+  cleared_selections = vim.deepcopy(M.context.selections or {})
+  for _, selection in ipairs(selections or {}) do
+    local already_cleared = false
+    for _, cleared_selection in ipairs(cleared_selections) do
+      if is_same_selection(selection, cleared_selection) then
+        already_cleared = true
+        break
+      end
+    end
+    if not already_cleared then
+      table.insert(cleared_selections, vim.deepcopy(selection))
+    end
+  end
+  cleared_selections_context = M.context
   M.context.mentioned_files = {}
   M.context.selections = {}
   state.context.set_context_updated_at(vim.uv.now())
@@ -367,6 +402,16 @@ end
 -- Load function that populates the global context state
 -- This is the core loading logic that was originally in the main context module
 function M.load()
+  local selections_cleared_by_send = {}
+  if cleared_selections_context == M.context then
+    selections_cleared_by_send = cleared_selections
+  else
+    cleared_selections = {}
+    cleared_selections_context = nil
+  end
+  cleared_selections = {}
+  cleared_selections_context = nil
+
   if not state.active_session and not state.is_opening then
     return
   end
@@ -412,7 +457,16 @@ function M.load()
       local selection_file = base_context.get_current_file_for_selection(buf)
       if selection_file then
         local selection = base_context.new_selection(selection_file, current_selection.text, current_selection.lines)
-        M.add_selection(selection)
+        local was_cleared = false
+        for _, cleared_selection in ipairs(selections_cleared_by_send) do
+          if is_same_selection(selection, cleared_selection) then
+            was_cleared = true
+            break
+          end
+        end
+        if not was_cleared then
+          M.add_selection(selection)
+        end
       end
     end
   end

lua/opencode/services/messaging.lua

@@ -65,6 +65,8 @@ M.send_message = Promise.async(function(prompt, opts)
   params.system = opts.system or config.default_system_prompt or nil
 
   local session_id = state.active_session.id
+  local sent_context = vim.deepcopy(context.get_context())
+  context.unload_attachments()
 
   local function update_sent_message_count(num)
     local sent_message_count = vim.deepcopy(state.user_message_count)
@@ -86,7 +88,7 @@ M.send_message = Promise.async(function(prompt, opts)
         return
       end
 
-      M.after_run(prompt)
+      M.after_run(prompt, sent_context)
     end)
     :catch(function(err)
       log.notify('Error sending message to session: ' .. vim.inspect(err), vim.log.levels.ERROR)
@@ -97,9 +99,13 @@ M.send_message = Promise.async(function(prompt, opts)
 end)
 
 ---@param prompt string
-function M.after_run(prompt)
-  context.unload_attachments()
-  state.session.set_last_sent_context(vim.deepcopy(context.get_context()))
+---@param sent_context? OpencodeContext
+function M.after_run(prompt, sent_context)
+  local context_sent = vim.deepcopy(sent_context or context.get_context())
+  if not sent_context then
+    context.unload_attachments()
+  end
+  state.session.set_last_sent_context(context_sent)
   context.delta_context()
   require('opencode.history').write(prompt)
   vim.g.opencode_abort_count = 0

tests/unit/context_spec.lua

@@ -899,6 +899,58 @@ describe('ChatContext.load() preserves selections on file switch', function()
     ChatContext.get_diagnostics = original_get_diagnostics
   end)
 
+  it('should not restore a selection cleared after sending', function()
+    local selection = {
+      file = { path = '/tmp/bar.lua', name = 'bar.lua', extension = 'lua' },
+      content = 'sent selection',
+      lines = '5, 7',
+    }
+    ChatContext.add_selection(selection)
+
+    local original_get_current_buf = BaseContext.get_current_buf
+    local original_get_current_file = BaseContext.get_current_file
+    local original_get_current_file_for_selection = BaseContext.get_current_file_for_selection
+    local original_get_current_cursor_data = BaseContext.get_current_cursor_data
+    local original_is_context_enabled = BaseContext.is_context_enabled
+    local original_get_current_selection = BaseContext.get_current_selection
+    local original_get_diagnostics = ChatContext.get_diagnostics
+
+    BaseContext.get_current_buf = function()
+      return 2, 2
+    end
+    BaseContext.get_current_file = function()
+      return selection.file
+    end
+    BaseContext.get_current_file_for_selection = function()
+      return selection.file
+    end
+    BaseContext.get_current_cursor_data = function()
+      return nil
+    end
+    BaseContext.is_context_enabled = function(context_type)
+      return context_type == 'selection'
+    end
+    BaseContext.get_current_selection = function()
+      return { text = selection.content, lines = selection.lines }
+    end
+    ChatContext.get_diagnostics = function()
+      return {}
+    end
+
+    ChatContext.unload_attachments()
+    ChatContext.load()
+
+    assert.same({}, ChatContext.context.selections)
+
+    BaseContext.get_current_buf = original_get_current_buf
+    BaseContext.get_current_file = original_get_current_file
+    BaseContext.get_current_file_for_selection = original_get_current_file_for_selection
+    BaseContext.get_current_cursor_data = original_get_current_cursor_data
+    BaseContext.is_context_enabled = original_is_context_enabled
+    BaseContext.get_current_selection = original_get_current_selection
+    ChatContext.get_diagnostics = original_get_diagnostics
+  end)
+
   it('should clear stale current_file when current_file context is disabled', function()
     ChatContext.context.current_file = {
       path = '/tmp/stale.lua',

tests/unit/services_messaging_spec.lua

@@ -8,6 +8,7 @@ loaded.services_messaging_spec = true
 local messaging = require('opencode.services.messaging')
 local session_runtime = require('opencode.services.session_runtime')
 local config_file = require('opencode.config_file')
+local context = require('opencode.context')
 local state = require('opencode.state')
 local Promise = require('opencode.promise')
 local stub = require('luassert.stub')
@@ -280,6 +281,16 @@ describe('opencode.services.messaging', function()
     state.session.set_active({ id = 'sess1' })
     state.session.set_user_message_count({})
 
+    local original_context = vim.deepcopy(context.get_context())
+    context.get_context().mentioned_files = { '/tmp/attached.lua' }
+    context.get_context().selections = {
+      {
+        file = { path = '/tmp/attached.lua', name = 'attached.lua', extension = 'lua' },
+        content = 'selected',
+        lines = '1, 2',
+      },
+    }
+
     local count_before = state.user_message_count['sess1'] or 0
     local count_during = nil
 
@@ -299,8 +310,79 @@ describe('opencode.services.messaging', function()
     assert.equal(0, count_before)
     assert.equal(1, count_during)
     assert.equal(0, count_after)
+    assert.same({}, context.get_context().mentioned_files)
+    assert.same({}, context.get_context().selections)
 
     state.api_client.create_message = orig
     session_runtime.cancel = orig_cancel
+    for key, value in pairs(original_context) do
+      context.get_context()[key] = value
+    end
+  end)
+
+  it('clears attachments before the request is sent', function()
+    state.ui.set_windows({ mock = 'windows' })
+    state.session.set_active({ id = 'sess1' })
+
+    local original_context = vim.deepcopy(context.get_context())
+    context.get_context().mentioned_files = { '/tmp/attached.lua' }
+    context.get_context().selections = {
+      {
+        file = { path = '/tmp/attached.lua', name = 'attached.lua', extension = 'lua' },
+        content = 'selected',
+        lines = '1, 2',
+      },
+    }
+
+    local observed_context
+    local original_create_message = state.api_client.create_message
+    state.api_client.create_message = function(_, _session_id, _params)
+      observed_context = vim.deepcopy(context.get_context())
+      return Promise.new():resolve({ info = { id = 'm1' }, parts = {} })
+    end
+
+    messaging.send_message('hello world'):wait()
+
+    assert.same({}, observed_context.mentioned_files)
+    assert.same({}, observed_context.selections)
+
+    state.api_client.create_message = original_create_message
+    for key, value in pairs(original_context) do
+      context.get_context()[key] = value
+    end
+  end)
+
+  it('clears sent attachments from the active context', function()
+    state.session.set_active({ id = 'sess1' })
+
+    local original_context = vim.deepcopy(context.get_context())
+    local sent_context = {
+      current_file = nil,
+      cursor_data = nil,
+      linter_errors = nil,
+      mentioned_files = { '/tmp/attached.lua' },
+      mentioned_subagents = {},
+      selections = {
+        {
+          file = { path = '/tmp/attached.lua', name = 'attached.lua', extension = 'lua' },
+          content = 'selected',
+          lines = '1, 2',
+        },
+      },
+    }
+    for key, value in pairs(sent_context) do
+      context.get_context()[key] = value
+    end
+
+    local delta_stub = stub(context, 'delta_context')
+    messaging.after_run('hello')
+
+    assert.same({}, context.get_context().mentioned_files)
+    assert.same({}, context.get_context().selections)
+
+    delta_stub:revert()
+    for key, value in pairs(original_context) do
+      context.get_context()[key] = value
+    end
   end)
 end)