feat: Add optional Neovim diagnostics to context (#15)

Author: ibeex

README.md

@@ -12,7 +12,7 @@ It's funny that all AI plugins for Neovim are quite complex to interact with, li
 
 ## Features
 
-- **Context aware**: Sends your current buffer, cwd, and selection as context.
+- **Context aware**: Sends your current buffer, cwd, selection, and optional diagnostics as context.
 - **Unsaved-buffer aware**: Tells pi to treat the sent Neovim buffer content as the source of truth, even if the on-disk file is stale.
 - **Simple configuration**: Just set your preferred AI model.
 - **Gets out of your way**: You ask it. It does it. Done.
@@ -69,6 +69,9 @@ require("pi").setup({
     selection = {
       surrounding_lines = 40,
     },
+    diagnostics = {
+      enabled = false,
+    },
   },
   skills = true,
   extensions = true,
@@ -86,6 +89,7 @@ require("pi").setup({
 | `context.max_bytes` | `24000` | Maximum size in bytes for sent context before trimming. |
 | `context.ask.surrounding_lines` | `80` | Number of lines before and after the current cursor line to include for `:PiAsk`. |
 | `context.selection.surrounding_lines` | `40` | Number of lines before and after the current visual selection to include for `:PiAskSelection`. |
+| `context.diagnostics.enabled` | `false` | Includes Neovim diagnostics in the sent context. `:PiAsk` sends all buffer diagnostics; `:PiAskSelection` sends only diagnostics overlapping the selected lines. |
 | `skills` | `true` | Whether pi discovers and loads skills. Set to `false` to pass `--no-skills`. |
 | `extensions` | `true` | Whether pi discovers and loads extensions. Set to `false` to pass `--no-extensions`. |
 
@@ -142,6 +146,7 @@ vim.keymap.set("v", "<leader>ai", ":PiAskSelection<CR>", { desc = "Ask pi (selec
 - Uses `nvim-notify` for status updates when available; otherwise falls back to a small floating status window.
 - Reloads changed loaded buffers on success so pi's on-disk edits are reflected in Neovim.
 - Treats sent buffer/selection context as newer than disk, so unsaved Neovim changes are the source of truth for the agent.
+- Optionally includes Neovim diagnostics from LSPs/linters via `vim.diagnostic`.
 - Trims oversized context for speed instead of always sending the full file.
 
 

lua/pi/config.lua

@@ -1,3 +1,4 @@
+--- Configuration helpers for pi.nvim.
 local M = {}
 
 local VALID_THINKING_LEVELS = {
@@ -24,19 +25,27 @@ M.defaults = {
     selection = {
       surrounding_lines = 40,
     },
+    diagnostics = {
+      enabled = false,
+    },
   },
   skills = true,
   extensions = true,
 }
 
 local values = vim.deepcopy(M.defaults)
 
+--- Validates a positive numeric configuration value.
+--- @param name string Human-readable config key name.
+--- @param value number Value to validate.
 local function validate_number(name, value)
   if type(value) ~= "number" or value < 1 then
     error(string.format("pi.nvim: %s must be a positive number", name))
   end
 end
 
+--- Validates user-provided configuration overrides.
+--- @param opts table User configuration passed to `setup`.
 function M.validate(opts)
   if opts.binary ~= nil and not (type(opts.binary) == "string" or type(opts.binary) == "table") then
     error("pi.nvim: binary must be a string or list of strings")
@@ -74,6 +83,14 @@ function M.validate(opts)
         validate_number("context.selection.surrounding_lines", context.selection.surrounding_lines)
       end
     end
+    if context.diagnostics ~= nil then
+      if type(context.diagnostics) ~= "table" then
+        error("pi.nvim: context.diagnostics must be a table")
+      end
+      if context.diagnostics.enabled ~= nil and type(context.diagnostics.enabled) ~= "boolean" then
+        error("pi.nvim: context.diagnostics.enabled must be a boolean")
+      end
+    end
   end
   if opts.skills ~= nil and type(opts.skills) ~= "boolean" then
     error("pi.nvim: skills must be a boolean")
@@ -97,13 +114,18 @@ function M.validate(opts)
   end
 end
 
+--- Merges user options with defaults and stores the effective config.
+--- @param opts? table User configuration overrides.
+--- @return table values Effective configuration.
 function M.setup(opts)
   opts = opts or {}
   M.validate(opts)
   values = vim.tbl_deep_extend("force", vim.deepcopy(M.defaults), opts)
   return values
 end
 
+--- Returns the currently active configuration.
+--- @return table values Effective configuration.
 function M.get()
   return values
 end

lua/pi/context.lua

@@ -1,3 +1,4 @@
+--- Context builders for pi.nvim prompts.
 local M = {}
 
 local SYSTEM_PROMPT = [[You are running inside the pi.nvim Neovim plugin. The user has sent a request and will not be able to reply back. You must complete the task immediately without asking any questions or requesting clarification. Take action now and do what was asked.
@@ -8,6 +9,9 @@ local BUFFER_SOURCE_OF_TRUTH_NOTE = [[NOTE: The context below comes from the cur
 
 local EMPTY_FILE_NOTE = [[NOTE: This file is currently empty. Please create or populate it directly by applying the necessary edits so pi.nvim can write the file.]]
 
+--- Returns whether a buffer contains only empty or whitespace-only lines.
+--- @param bufnr integer Buffer handle.
+--- @return boolean is_empty Whether the buffer has meaningful content.
 local function buffer_is_empty(bufnr)
   local line_count = vim.api.nvim_buf_line_count(bufnr)
   if line_count == 0 then
@@ -22,6 +26,9 @@ local function buffer_is_empty(bufnr)
   return true
 end
 
+--- Returns whether a buffer maps to a real file on disk.
+--- @param bufnr integer Buffer handle.
+--- @return boolean is_file_backed Whether the buffer is file-backed.
 function M.buffer_is_file_backed(bufnr)
   if vim.bo[bufnr].buftype ~= "" then
     return false
@@ -30,6 +37,8 @@ function M.buffer_is_file_backed(bufnr)
   return filename ~= nil and filename ~= ""
 end
 
+--- Returns the current visual selection as a 1-based inclusive line range.
+--- @return table|nil range Selection range with `start` and `end` keys.
 function M.get_visual_selection_range()
   local start_pos = vim.fn.getpos("'<")
   local end_pos = vim.fn.getpos("'>")
@@ -44,6 +53,10 @@ function M.get_visual_selection_range()
   return { start = start_line, ["end"] = end_line }
 end
 
+--- Builds the prompt label shown in `vim.ui.input`.
+--- @param bufnr integer Buffer handle.
+--- @param selection_range table|nil Optional selected line range.
+--- @return string label Input prompt label.
 function M.format_prompt_label(bufnr, selection_range)
   local components = {}
   local filename = vim.api.nvim_buf_get_name(bufnr)
@@ -59,31 +72,144 @@ function M.format_prompt_label(bufnr, selection_range)
   return string.format("ask pi (%s): ", table.concat(components, ":"))
 end
 
+--- Returns a buffer's filetype, defaulting to plain text.
+--- @param bufnr integer Buffer handle.
+--- @return string filetype Buffer filetype or `text`.
 local function filetype_for(bufnr)
   return vim.bo[bufnr].filetype ~= "" and vim.bo[bufnr].filetype or "text"
 end
 
+--- Extracts nearby lines around a 1-based center line.
+--- @param lines string[] All buffer lines.
+--- @param center_line integer 1-based center line.
+--- @param surrounding_lines integer Number of surrounding lines to include on each side.
+--- @return string[] slice Selected lines.
+--- @return integer start_line First included line number.
+--- @return integer end_line Last included line number.
 local function slice_lines_around(lines, center_line, surrounding_lines)
   local start_line = math.max(1, center_line - surrounding_lines)
   local end_line = math.min(#lines, center_line + surrounding_lines)
   return vim.list_slice(lines, start_line, end_line), start_line, end_line
 end
 
+--- Truncates text to a byte limit.
+--- @param text string Text to trim.
+--- @param max_bytes integer Maximum byte length.
+--- @return string content Trimmed or original text.
+--- @return boolean did_trim Whether truncation happened.
 local function truncate_to_bytes(text, max_bytes)
   if #text <= max_bytes then
     return text, false
   end
   return text:sub(1, max_bytes), true
 end
 
+--- Wraps text inside a labeled fenced code block.
+--- @param label string Section label.
+--- @param text string Section content.
+--- @return string block Formatted block.
 local function content_block(label, text)
   return string.format("%s:\n```\n%s\n```", label, text)
 end
 
+--- Converts a diagnostic severity enum into a stable uppercase label.
+--- @param severity integer|nil Diagnostic severity.
+--- @return string label Severity label.
+local function diagnostic_severity_label(severity)
+  local labels = {
+    [vim.diagnostic.severity.ERROR] = "ERROR",
+    [vim.diagnostic.severity.WARN] = "WARN",
+    [vim.diagnostic.severity.INFO] = "INFO",
+    [vim.diagnostic.severity.HINT] = "HINT",
+  }
+  return labels[severity] or "UNKNOWN"
+end
+
+--- Returns whether a diagnostic overlaps an inclusive line range.
+--- @param diagnostic table Diagnostic item from `vim.diagnostic.get`.
+--- @param start_line integer Inclusive start line.
+--- @param end_line integer Inclusive end line.
+--- @return boolean overlaps Whether the diagnostic intersects the range.
+local function diagnostic_overlaps_range(diagnostic, start_line, end_line)
+  local diagnostic_start = (diagnostic.lnum or 0) + 1
+  local diagnostic_end = (diagnostic.end_lnum or diagnostic.lnum or 0) + 1
+  return diagnostic_start <= end_line and diagnostic_end >= start_line
+end
+
+--- Formats a diagnostic as a human-readable bullet list item.
+--- @param diagnostic table Diagnostic item from `vim.diagnostic.get`.
+--- @return string line Rendered diagnostic line.
+local function format_diagnostic(diagnostic)
+  local line = (diagnostic.lnum or 0) + 1
+  local col = (diagnostic.col or 0) + 1
+  local details = { diagnostic_severity_label(diagnostic.severity) }
+
+  if diagnostic.source and diagnostic.source ~= "" then
+    details[#details + 1] = diagnostic.source
+  end
+
+  return string.format("- line %d:%d: %s [%s]", line, col, diagnostic.message, table.concat(details, ", "))
+end
+
+--- Builds an optional diagnostics block for the current buffer or selection.
+--- @param bufnr integer Buffer handle.
+--- @param config table Active pi.nvim configuration.
+--- @param opts? table Optional range filter.
+--- @return string|nil block Formatted diagnostics block.
+--- @return string|nil note Optional trimming note.
+local function diagnostics_block(bufnr, config, opts)
+  if not config.context.diagnostics or not config.context.diagnostics.enabled then
+    return nil
+  end
+
+  opts = opts or {}
+  local diagnostics = vim.diagnostic.get(bufnr)
+  if not diagnostics or vim.tbl_isempty(diagnostics) then
+    return nil
+  end
+
+  if opts.range then
+    diagnostics = vim.tbl_filter(function(diagnostic)
+      return diagnostic_overlaps_range(diagnostic, opts.range.start, opts.range["end"])
+    end, diagnostics)
+    if vim.tbl_isempty(diagnostics) then
+      return nil
+    end
+  end
+
+  table.sort(diagnostics, function(a, b)
+    if a.lnum ~= b.lnum then
+      return a.lnum < b.lnum
+    end
+    return (a.col or 0) < (b.col or 0)
+  end)
+
+  local formatted = {}
+  for _, diagnostic in ipairs(diagnostics) do
+    formatted[#formatted + 1] = format_diagnostic(diagnostic)
+  end
+
+  local content, did_trim_bytes = truncate_to_bytes(table.concat(formatted, "\n"), config.context.max_bytes)
+  local label = opts.range and "Diagnostics in selection" or "Diagnostics"
+  local note = nil
+
+  if did_trim_bytes then
+    note = string.format("NOTE: Diagnostics were trimmed for speed (max_bytes=%d).", config.context.max_bytes)
+  end
+
+  return content_block(label, content), note
+end
+
+--- Returns the system prompt appended to pi invocations.
+--- @return string prompt Internal system prompt.
 function M.get_system_prompt()
   return SYSTEM_PROMPT
 end
 
+--- Builds prompt context for `:PiAsk` around the current cursor line.
+--- @param bufnr integer Buffer handle.
+--- @param config table Active pi.nvim configuration.
+--- @return string context Prompt context payload.
 function M.get_buffer_context(bufnr, config)
   local lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, false)
   local cursor_line = vim.api.nvim_win_get_cursor(0)[1]
@@ -108,13 +234,25 @@ function M.get_buffer_context(bufnr, config)
     )
   end
 
+  local diagnostics, diagnostics_note = diagnostics_block(bufnr, config)
+  if diagnostics then
+    parts[#parts + 1] = diagnostics
+  end
+  if diagnostics_note then
+    parts[#parts + 1] = diagnostics_note
+  end
+
   if buffer_is_empty(bufnr) then
     parts[#parts + 1] = EMPTY_FILE_NOTE
   end
 
   return table.concat(parts, "\n\n")
 end
 
+--- Builds prompt context for `:PiAskSelection`.
+--- @param bufnr integer Buffer handle.
+--- @param config table Active pi.nvim configuration.
+--- @return string context Prompt context payload.
 function M.get_visual_context(bufnr, config)
   local filename = vim.api.nvim_buf_get_name(bufnr)
   local all_lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, false)
@@ -145,6 +283,14 @@ function M.get_visual_context(bufnr, config)
     )
   end
 
+  local diagnostics, diagnostics_note = diagnostics_block(bufnr, config, { range = selection_range })
+  if diagnostics then
+    parts[#parts + 1] = diagnostics
+  end
+  if diagnostics_note then
+    parts[#parts + 1] = diagnostics_note
+  end
+
   if buffer_is_empty(bufnr) then
     parts[#parts + 1] = EMPTY_FILE_NOTE
   end