From 74224f8008ee5f109d0861476591707dc73d9865 Mon Sep 17 00:00:00 2001
From: Davidyz <hzjlyz@gmail.com>
Date: Sat, 12 Apr 2025 13:56:41 +0800
Subject: [PATCH 1/2] feat(nvim): avoid retrieving files that is already
 present in the context.

---
 lua/vectorcode/integrations/codecompanion.lua | 26 +++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/lua/vectorcode/integrations/codecompanion.lua b/lua/vectorcode/integrations/codecompanion.lua
index 9033ae1b..21c36907 100644
--- a/lua/vectorcode/integrations/codecompanion.lua
+++ b/lua/vectorcode/integrations/codecompanion.lua
@@ -7,6 +7,7 @@
 ---@field use_lsp boolean?
 ---@field auto_submit table<string, boolean>?
 ---@field ls_on_start boolean?
+---@field no_duplicate boolean?
 
 local vc_config = require("vectorcode.config")
 local check_cli_wrap = vc_config.check_cli_wrap
@@ -21,6 +22,10 @@ local function flatten_table_to_string(t)
   return table.concat(vim.iter(t):flatten(math.huge):totable(), "\n")
 end
 
+---@alias path string
+---@type table<CodeCompanion.Chat, table<path, boolean>> For each chat, keep a record of what files has been sent.
+local added_context = {}
+
 local job_runner = nil
 ---@param use_lsp boolean
 local function initialise_runner(use_lsp)
@@ -58,6 +63,7 @@ local make_tool = check_cli_wrap(function(opts)
     use_lsp = false,
     auto_submit = { ls = false, query = false },
     ls_on_start = false,
+    no_duplicate = false,
   }, opts or {})
   local capping_message = ""
   if opts.max_num > 0 then
@@ -105,6 +111,20 @@ local make_tool = check_cli_wrap(function(opts)
               )
             end
           end
+
+          if opts.no_duplicate then
+            -- exclude files that has been added to the context
+            local existing_files = { "--exclude" }
+            for path, ok in pairs(added_context[agent.chat] or {}) do
+              if ok then
+                table.insert(existing_files, path)
+              end
+            end
+            if #existing_files > 1 then
+              vim.list_extend(args, existing_files)
+            end
+          end
+
           job_runner.run_async(args, function(result, error)
             vim.schedule(function()
               if opts.auto_submit[action.command] then
@@ -315,6 +335,12 @@ Remember:
                   file.document
                 ),
               }, { visible = false })
+              if opts.no_duplicate then
+                if added_context[agent.chat] == nil then
+                  added_context[agent.chat] = {}
+                end
+                added_context[agent.chat][file.path] = true
+              end
             end
           end
         elseif cmd.command == "ls" then

From d782e783d7b37d71c09f022c7d9595c22aca0ce7 Mon Sep 17 00:00:00 2001
From: Davidyz <hzjlyz@gmail.com>
Date: Sun, 13 Apr 2025 17:37:36 +0800
Subject: [PATCH 2/2] refactor(nvim): use `CodeCompanion.Chat.Ref` for
 reference tracking

---
 lua/vectorcode/integrations/codecompanion.lua | 30 +++++++++----------
 1 file changed, 14 insertions(+), 16 deletions(-)

diff --git a/lua/vectorcode/integrations/codecompanion.lua b/lua/vectorcode/integrations/codecompanion.lua
index 21c36907..ffed80ca 100644
--- a/lua/vectorcode/integrations/codecompanion.lua
+++ b/lua/vectorcode/integrations/codecompanion.lua
@@ -13,6 +13,8 @@ local vc_config = require("vectorcode.config")
 local check_cli_wrap = vc_config.check_cli_wrap
 local notify_opts = vc_config.notify_opts
 
+local tool_result_source = "VectorCodeToolResult"
+
 ---@param t table|string
 ---@return string
 local function flatten_table_to_string(t)
@@ -22,10 +24,6 @@ local function flatten_table_to_string(t)
   return table.concat(vim.iter(t):flatten(math.huge):totable(), "\n")
 end
 
----@alias path string
----@type table<CodeCompanion.Chat, table<path, boolean>> For each chat, keep a record of what files has been sent.
-local added_context = {}
-
 local job_runner = nil
 ---@param use_lsp boolean
 local function initialise_runner(use_lsp)
@@ -63,7 +61,7 @@ local make_tool = check_cli_wrap(function(opts)
     use_lsp = false,
     auto_submit = { ls = false, query = false },
     ls_on_start = false,
-    no_duplicate = false,
+    no_duplicate = true,
   }, opts or {})
   local capping_message = ""
   if opts.max_num > 0 then
@@ -112,12 +110,12 @@ local make_tool = check_cli_wrap(function(opts)
             end
           end
 
-          if opts.no_duplicate then
+          if opts.no_duplicate and agent.chat.refs ~= nil then
             -- exclude files that has been added to the context
             local existing_files = { "--exclude" }
-            for path, ok in pairs(added_context[agent.chat] or {}) do
-              if ok then
-                table.insert(existing_files, path)
+            for _, ref in pairs(agent.chat.refs) do
+              if ref.source == tool_result_source then
+                table.insert(existing_files, ref.id)
               end
             end
             if #existing_files > 1 then
@@ -318,6 +316,7 @@ Remember:
       success = function(agent, cmd, stdout)
         stdout = stdout[1]
         if cmd.command == "query" then
+          agent.chat.ui:unlock_buf()
           for i, file in pairs(stdout) do
             if opts.max_num < 0 or i <= opts.max_num then
               agent.chat:add_message({
@@ -334,13 +333,12 @@ Remember:
                   file.path,
                   file.document
                 ),
-              }, { visible = false })
-              if opts.no_duplicate then
-                if added_context[agent.chat] == nil then
-                  added_context[agent.chat] = {}
-                end
-                added_context[agent.chat][file.path] = true
-              end
+              }, { visible = false, id = file.path })
+              agent.chat.references:add({
+                source = tool_result_source,
+                id = file.path,
+                opts = { visible = false },
+              })
             end
           end
         elseif cmd.command == "ls" then