[feat!] Advanced stop logic (#184)

[breaking] Remove `PrtChatStop` for the general `PrtStop` command.

Author: frankroeder

README.md

@@ -204,7 +204,7 @@ Additional useful commands are implemented through hooks (see below).
 | `PrtChatFinder`           | Fuzzy search chat files using fzf             |
 | `PrtChatDelete`           | Delete the current chat file                  |
 | `PrtChatRespond`          | Trigger chat respond (in chat file)           |
-| `PrtStop`                 | Interrupt ongoing respond                     |
+| `PrtStop`                 | Interrupt any ongoing Parrot generation (works everywhere) |
 | `PrtProvider <provider>`  | Switch the provider (empty arg triggers fzf)  |
 | `PrtModel <model>`        | Switch the interactive command model (empty arg triggers fzf). Note: Chat model must be changed from within the chat buffer. |
 | `PrtStatus`               | Prints current provider and model selection   |
@@ -403,7 +403,7 @@ This plugin provides the following default key mappings:
 |--------------|-------------------------------------------------------------|
 | `<C-g>c`     | Opens a new chat via `PrtChatNew`                           |
 | `<C-g><C-g>` | Trigger the API to generate a response via `PrtChatRespond` |
-| `<C-g>s`     | Stop the current text generation via `PrtStop`              |
+| `<C-g>s`     | Stop any ongoing Parrot generation via `PrtStop`            |
 | `<C-g>d`     | Delete the current chat file via `PrtChatDelete`            |
 
 ### Provider Configuration Examples
@@ -942,6 +942,52 @@ or have suggestions for improving provider support.
   }
 ```
 
+## Cancellation
+
+You can stop any ongoing Parrot generation at any time using multiple methods:
+
+### Methods
+
+1. **Keybinding**: `<C-g>s` (configurable via `chat_shortcut_stop`)
+2. **Command**: `:PrtStop` (works everywhere)
+
+### Behavior
+
+When you cancel a generation:
+
+- **Immediate Termination**: The API request is stopped immediately
+- **Preserves Generated Text**: The text generated so far remains in the buffer
+- **Visual Feedback**: You receive a notification confirming the cancellation
+- **Preview Mode**: If cancelled during streaming, the preview won't be shown
+- **Multiple Jobs**: If multiple generations are running, all are stopped
+
+### Autocommand Event
+
+A `User PrtCancelled` event is fired when generation is cancelled, allowing you to create custom hooks:
+
+```lua
+vim.api.nvim_create_autocmd("User", {
+  pattern = "PrtCancelled",
+  callback = function()
+    -- Your custom logic here
+    print("Parrot generation was cancelled")
+  end,
+})
+```
+
+### Advanced Usage
+
+For buffer-specific cancellation in custom code:
+
+```lua
+-- Stop only jobs for current buffer
+local chat_handler = require("parrot").chat_handler
+chat_handler:stop({ buffer = vim.api.nvim_get_current_buf() })
+
+-- Stop without notification
+chat_handler:stop({ notify = false })
+```
+
 ## Bonus
 
 Access parrot.nvim directly from your terminal:

lua/parrot/chat_handler.lua

@@ -377,19 +377,81 @@ function ChatHandler:Cmd(params)
 end
 
 --- Stops all ongoing processes by killing associated jobs.
----@param signal number | nil Signal to send to the processes.
-function ChatHandler:stop(signal)
+---@param options table|number|nil Options table or signal number for backwards compatibility
+---   - options.signal: Signal to send to processes (default: 15)
+---   - options.buffer: Buffer number to stop jobs for (nil = stop all)
+---   - options.notify: Show notification after stopping (default: true)
+function ChatHandler:stop(options)
+  -- Backwards compatibility: if options is a number, treat it as signal
+  if type(options) == "number" then
+    options = { signal = options }
+  end
+  options = options or {}
+  local signal = options.signal or 15
+  local target_buf = options.buffer
+  local show_notification = options.notify ~= false
+
   if self.pool:is_empty() then
+    if show_notification then
+      logger.warning("No active Parrot processes to stop")
+    end
     return
   end
 
-  for _, process_info in self.pool:ipairs() do
-    if process_info.job.handle ~= nil and not process_info.job.handle:is_closing() then
-      vim.uv.kill(process_info.job.pid, signal or 15)
+  local stopped_count = 0
+  local cancelled_queries = {}
+
+  -- Collect jobs to stop (either for specific buffer or all)
+  for i = #self.pool._processes, 1, -1 do
+    local process_info = self.pool._processes[i]
+    local should_stop = target_buf == nil or process_info.buf == target_buf
+
+    if should_stop then
+      -- Mark associated query as cancelled
+      if process_info.qid then
+        self.queries:mark_cancelled(process_info.qid, "user")
+        table.insert(cancelled_queries, process_info.qid)
+      end
+
+      -- Kill the job
+      if process_info.job.handle ~= nil and not process_info.job.handle:is_closing() then
+        vim.uv.kill(process_info.job.pid, signal)
+        stopped_count = stopped_count + 1
+      end
+
+      -- Remove from pool
+      table.remove(self.pool._processes, i)
     end
   end
 
-  self.pool = Pool:new()
+  -- Clean up extmarks only (preserve generated text)
+  vim.schedule(function()
+    for _, qid in ipairs(cancelled_queries) do
+      local qt = self.queries:get(qid)
+      if qt then
+        -- Clear namespace/extmarks only
+        if qt.ns_id and qt.buf and vim.api.nvim_buf_is_valid(qt.buf) then
+          pcall(vim.api.nvim_buf_clear_namespace, qt.buf, qt.ns_id, 0, -1)
+        end
+      end
+    end
+
+    -- Show notification
+    if show_notification then
+      if stopped_count > 0 then
+        local msg = stopped_count == 1 and "Stopped 1 process" or string.format("Stopped %d processes", stopped_count)
+        if target_buf then
+          msg = msg .. " in current buffer"
+        end
+        logger.info(msg)
+      else
+        logger.warning("No active processes found" .. (target_buf and " in current buffer" or ""))
+      end
+    end
+
+    -- Fire autocmd event for user hooks
+    vim.cmd("doautocmd User PrtCancelled")
+  end)
 end
 
 --- Context command
@@ -1889,9 +1951,21 @@ function ChatHandler:query(buf, provider, payload, handler, on_exit)
     end,
   })
   job:start()
-  self.pool:add(job, buf)
+
+  -- Determine target type for better tracking
+  local target_type = "unknown"
+  if buf and vim.api.nvim_buf_is_valid(buf) then
+    local file_name = vim.api.nvim_buf_get_name(buf)
+    local utils_module = require("parrot.utils")
+    if utils_module.is_chat(buf, file_name, self.options.chat_dir) then
+      target_type = "chat"
+    end
+  end
+
+  self.pool:add(job, buf, qid, target_type)
   logger.debug("ChatHandler:query pool updated", {
     pool = self.pool,
+    qid = qid,
   })
 end
 

lua/parrot/config.lua

@@ -371,7 +371,7 @@ function M.setup(opts)
 
   M.cmd = {
     ChatFinder = "chat_finder",
-    ChatStop = "stop",
+    Stop = "stop",
     ChatNew = "chat_new",
     ChatToggle = "chat_toggle",
     ChatPaste = "chat_paste",

lua/parrot/pool.lua

@@ -10,8 +10,16 @@ end
 --- Adds a process to the pool.
 --- @param job table # A plenary job.
 --- @param buf number|nil # The buffer number (optional)
-function Pool:add(job, buf)
-  table.insert(self._processes, { job = job, buf = buf })
+--- @param qid string|nil # The query ID (optional)
+--- @param target_type string|nil # The target type (optional) - "chat", "rewrite", "popup", etc.
+function Pool:add(job, buf, qid, target_type)
+  table.insert(self._processes, {
+    job = job,
+    buf = buf,
+    qid = qid,
+    target_type = target_type,
+    timestamp = os.time(),
+  })
 end
 
 --- Checks if there is no other process running for the given buffer.
@@ -51,4 +59,61 @@ function Pool:ipairs()
   return ipairs(self._processes)
 end
 
+--- Gets processes for a specific buffer.
+--- @param buf number # The buffer number
+--- @return table # List of process info tables for the buffer
+function Pool:get_for_buffer(buf)
+  local result = {}
+  for _, process_info in ipairs(self._processes) do
+    if process_info.buf == buf then
+      table.insert(result, process_info)
+    end
+  end
+  return result
+end
+
+--- Gets the most recent active job (by timestamp).
+--- @return table|nil # The most recent process info, or nil if pool is empty
+function Pool:get_active_job()
+  if self:is_empty() then
+    return nil
+  end
+
+  local most_recent = nil
+  for _, process_info in ipairs(self._processes) do
+    if most_recent == nil or process_info.timestamp > most_recent.timestamp then
+      most_recent = process_info
+    end
+  end
+  return most_recent
+end
+
+--- Stops jobs for a specific buffer.
+--- @param buf number # The buffer number
+--- @param signal number|nil # Signal to send (default 15)
+--- @return number # Number of jobs stopped
+function Pool:stop_buffer(buf, signal)
+  signal = signal or 15
+  local stopped_count = 0
+
+  for i = #self._processes, 1, -1 do
+    local process_info = self._processes[i]
+    if process_info.buf == buf then
+      if process_info.job.handle ~= nil and not process_info.job.handle:is_closing() then
+        vim.uv.kill(process_info.job.pid, signal)
+        stopped_count = stopped_count + 1
+      end
+      table.remove(self._processes, i)
+    end
+  end
+
+  return stopped_count
+end
+
+--- Checks if there are any active jobs in the pool.
+--- @return boolean # True if there are active jobs, false otherwise
+function Pool:has_active_jobs()
+  return not self:is_empty()
+end
+
 return Pool

lua/parrot/preview_response_handler.lua

@@ -99,6 +99,13 @@ function PreviewResponseHandler:handle_chunk(qid, chunk)
     return
   end
 
+  -- Check if query was cancelled
+  if qt.cancelled then
+    -- For preview, we don't show the preview if cancelled
+    logger.debug("PreviewResponseHandler: Query cancelled, aborting preview")
+    return
+  end
+
   if chunk and chunk ~= "" then
     self.response = self.response .. chunk
     qt.response = self.response
@@ -269,7 +276,14 @@ end
 --- Creates a completion handler that shows the preview
 ---@return function
 function PreviewResponseHandler:create_completion_handler()
-  return vim.schedule_wrap(function(_)
+  return vim.schedule_wrap(function(qid)
+    -- Check if cancelled before showing preview
+    local qt = self.queries:get(qid)
+    if qt and qt.cancelled then
+      logger.debug("PreviewResponseHandler: Completion cancelled, not showing preview")
+      return
+    end
+
     -- Clean up the response (remove code fences, etc.)
     self.response = self
       .response

lua/parrot/queries.lua

@@ -13,6 +13,10 @@ end
 --- @param qid number # Query ID.
 --- @param data table # Query data.
 function Queries:add(qid, data)
+  -- Initialize cancellation state
+  data.cancelled = false
+  data.cancellation_reason = nil
+  data.cancellation_time = nil
   self._queries[qid] = data
 end
 
@@ -38,6 +42,39 @@ function Queries:get(qid)
   return self._queries[qid]
 end
 
+--- Gets queries for a specific buffer.
+--- @param buf number # The buffer number
+--- @return table # List of query IDs for the buffer
+function Queries:get_for_buffer(buf)
+  local result = {}
+  for qid, query_data in pairs(self._queries) do
+    if query_data.buf == buf then
+      table.insert(result, qid)
+    end
+  end
+  return result
+end
+
+--- Marks a query as cancelled.
+--- @param qid string # Query ID.
+--- @param reason string|nil # Cancellation reason (optional)
+function Queries:mark_cancelled(qid, reason)
+  local query = self._queries[qid]
+  if query then
+    query.cancelled = true
+    query.cancellation_reason = reason or "user"
+    query.cancellation_time = os.time()
+  end
+end
+
+--- Checks if a query was cancelled.
+--- @param qid string # Query ID.
+--- @return boolean # True if cancelled, false otherwise.
+function Queries:is_cancelled(qid)
+  local query = self._queries[qid]
+  return query and query.cancelled or false
+end
+
 --- Cleans up old queries from the collection based on the specified criteria.
 --- @param N number # Number of queries to keep.
 --- @param age number # Age of queries to keep in seconds.

lua/parrot/response_handler.lua

@@ -60,6 +60,11 @@ function ResponseHandler:handle_chunk(qid, chunk)
     return
   end
 
+  -- Check if query was cancelled - just stop processing, preserve existing text
+  if qt.cancelled then
+    return
+  end
+
   if not self.skip_first_undojoin then
     utils.undojoin(self.buffer)
   end

tests/parrot/pool_spec.lua

@@ -18,7 +18,10 @@ describe("Pool", function()
         local job = { pid = 1 }
         local buf = 10
         pool:add(job, buf)
-        assert.are.same({ { job = job, buf = buf } }, pool._processes)
+        assert.equals(1, #pool._processes)
+        assert.are.same(job, pool._processes[1].job)
+        assert.equals(buf, pool._processes[1].buf)
+        assert.is_not_nil(pool._processes[1].timestamp)
       end)
     end)
   end)
@@ -57,7 +60,9 @@ describe("Pool", function()
         pool:add(job1, 10)
         pool:add(job2, 20)
         pool:remove(1)
-        assert.are.same({ { job = job2, buf = 20 } }, pool._processes)
+        assert.equals(1, #pool._processes)
+        assert.are.same(job2, pool._processes[1].job)
+        assert.equals(20, pool._processes[1].buf)
       end)
     end)
   end)