diff --git a/README.md b/README.md index d0cddadd..f385dda7 100644 --- a/README.md +++ b/README.md @@ -240,6 +240,7 @@ require('opencode').setup({ debug = { enabled = false, -- Enable debug messages in the output window }, + prompt_guard = nil, -- Optional function that returns boolean to control when prompts can be sent (see Prompt Guard section) }) ``` @@ -537,6 +538,31 @@ The plugin defines several highlight groups that can be customized to match your - `OpencodeInputLegend`: Highlight for input window legend (default: #CCCCCC background) - `OpencodeHint`: Highlight for hinting messages in input window and token info in output window footer (linked to `Comment`) +## 🛡️ Prompt Guard + +The `prompt_guard` configuration option allows you to control when prompts can be sent to Opencode. This is useful for preventing accidental or unauthorized AI interactions in certain contexts. + +### Configuration + +Set `prompt_guard` to a function that returns a boolean: + +```lua +require('opencode').setup({ + prompt_guard = function() + -- Your custom logic here + -- Return true to allow, false to deny + return true + end, +}) +``` + +### Behavior + +- **Before sending prompts**: The guard is checked before any prompt is sent to the AI. If denied, an ERROR notification is shown and the prompt is not sent. +- **Before opening UI**: The guard is checked when opening the Opencode buffer for the first time. If denied, a WARN notification is shown and the UI is not opened. +- **No parameters**: The guard function receives no parameters. Access vim state directly (e.g., `vim.fn.getcwd()`, `vim.bo.filetype`). +- **Error handling**: If the guard function throws an error or returns a non-boolean value, the prompt is denied with an appropriate error message. + ## 🔧 Setting up Opencode If you're new to opencode: diff --git a/lua/opencode/config.lua b/lua/opencode/config.lua index ae556999..0f730489 100644 --- a/lua/opencode/config.lua +++ b/lua/opencode/config.lua @@ -167,6 +167,7 @@ M.defaults = { enabled = false, capture_streamed_events = false, }, + prompt_guard = nil, } M.values = vim.deepcopy(M.defaults) diff --git a/lua/opencode/core.lua b/lua/opencode/core.lua index 6fa4f7d2..3f75f04a 100644 --- a/lua/opencode/core.lua +++ b/lua/opencode/core.lua @@ -48,6 +48,14 @@ function M.open(opts) local are_windows_closed = state.windows == nil if are_windows_closed then + -- Check if whether prompting will be allowed + local context_module = require('opencode.context') + local mentioned_files = context_module.context.mentioned_files or {} + local allowed, err_msg = util.check_prompt_allowed(config.prompt_guard, mentioned_files) + if not allowed then + vim.notify(err_msg or 'Prompts will be denied by prompt_guard', vim.log.levels.WARN) + end + state.windows = ui.create_windows() end @@ -81,6 +89,16 @@ end --- @param prompt string The message prompt to send. --- @param opts? SendMessageOpts function M.send_message(prompt, opts) + -- Check if prompt is allowed + local context_module = require('opencode.context') + local mentioned_files = context_module.context.mentioned_files or {} + local allowed, err_msg = util.check_prompt_allowed(config.prompt_guard, mentioned_files) + + if not allowed then + vim.notify(err_msg or 'Prompt denied by prompt_guard', vim.log.levels.ERROR) + return + end + opts = opts or {} opts.context = opts.context or config.context opts.model = opts.model or state.current_model diff --git a/lua/opencode/types.lua b/lua/opencode/types.lua index 299807c1..249a1a6a 100644 --- a/lua/opencode/types.lua +++ b/lua/opencode/types.lua @@ -141,6 +141,7 @@ ---@field ui OpencodeUIConfig ---@field context OpencodeContextConfig ---@field debug OpencodeDebugConfig +---@field prompt_guard? fun(mentioned_files: string[]): boolean ---@class MessagePartState ---@field input TaskToolInput|BashToolInput|FileToolInput|TodoToolInput|GlobToolInput|GrepToolInput|WebFetchToolInput|ListToolInput Input data for the tool diff --git a/lua/opencode/ui/highlight.lua b/lua/opencode/ui/highlight.lua index da0ab602..f76cab60 100644 --- a/lua/opencode/ui/highlight.lua +++ b/lua/opencode/ui/highlight.lua @@ -22,6 +22,7 @@ function M.setup() vim.api.nvim_set_hl(0, 'OpencodeContextualActions', { bg = '#90A4AE', fg = '#1976D2', bold = true, default = true }) vim.api.nvim_set_hl(0, 'OpencodeInputLegend', { bg = '#757575', fg = '#424242', bold = false, default = true }) vim.api.nvim_set_hl(0, 'OpencodeHint', { link = 'Comment', default = true }) + vim.api.nvim_set_hl(0, 'OpencodeGuardDenied', { fg = '#F44336', bold = true, default = true }) else vim.api.nvim_set_hl(0, 'OpencodeBorder', { fg = '#616161', default = true }) vim.api.nvim_set_hl(0, 'OpencodeBackground', { link = 'Normal', default = true }) @@ -41,6 +42,7 @@ function M.setup() vim.api.nvim_set_hl(0, 'OpencodeContextualActions', { bg = '#3b4261', fg = '#61AFEF', bold = true, default = true }) vim.api.nvim_set_hl(0, 'OpencodeInputLegend', { bg = '#616161', fg = '#CCCCCC', bold = false, default = true }) vim.api.nvim_set_hl(0, 'OpencodeHint', { link = 'Comment', default = true }) + vim.api.nvim_set_hl(0, 'OpencodeGuardDenied', { fg = '#EF5350', bold = true, default = true }) end end diff --git a/lua/opencode/ui/icons.lua b/lua/opencode/ui/icons.lua index 972b896e..d8d3363d 100644 --- a/lua/opencode/ui/icons.lua +++ b/lua/opencode/ui/icons.lua @@ -25,6 +25,7 @@ local presets = { -- statuses status_on = '🟢', status_off = '⚫', + guard_on = '🚫', -- borders and misc border = '▌', }, @@ -49,6 +50,7 @@ local presets = { -- statuses status_on = ' ', status_off = ' ', + guard_on = '', -- borders and misc border = '▌', }, @@ -73,6 +75,7 @@ local presets = { -- statuses status_on = 'ON', status_off = 'OFF', + guard_on = 'X', -- borders and misc border = '▌', }, diff --git a/lua/opencode/ui/prompt_guard_indicator.lua b/lua/opencode/ui/prompt_guard_indicator.lua new file mode 100644 index 00000000..546ab83c --- /dev/null +++ b/lua/opencode/ui/prompt_guard_indicator.lua @@ -0,0 +1,36 @@ +local M = {} + +local config = require('opencode.config') +local util = require('opencode.util') +local icons = require('opencode.ui.icons') +local context = require('opencode.context') + +---Get the current prompt guard status +---@return boolean allowed +---@return string|nil error_message +function M.get_status() + local mentioned_files = context.context.mentioned_files or {} + return util.check_prompt_allowed(config.prompt_guard, mentioned_files) +end + +---Check if guard will deny prompts +---@return boolean denied +function M.is_denied() + local allowed, _ = M.get_status() + return not allowed +end + +---Get formatted indicator string with highlight (empty if allowed) +---@return string formatted_indicator +function M.get_formatted() + if not M.is_denied() then + -- Prompts are allowed - don't show anything + return '' + end + + -- Prompts will be denied - show red indicator + local icon = icons.get('guard_on') + return string.format('%%#OpencodeGuardDenied#%s%%*', icon) +end + +return M diff --git a/lua/opencode/ui/timer.lua b/lua/opencode/ui/timer.lua index cee93bc2..9cbe6052 100644 --- a/lua/opencode/ui/timer.lua +++ b/lua/opencode/ui/timer.lua @@ -48,6 +48,12 @@ function Timer:start() end end +--- Start the timer and immediately execute the callback +function Timer:start_and_tick() + self:start() + self.on_tick(unpack(self.args)) +end + function Timer:stop() if not self._uv_timer then return diff --git a/lua/opencode/ui/topbar.lua b/lua/opencode/ui/topbar.lua index f01bbf6b..f7a7bbf0 100644 --- a/lua/opencode/ui/topbar.lua +++ b/lua/opencode/ui/topbar.lua @@ -2,6 +2,7 @@ local M = {} local state = require('opencode.state') local config_file = require('opencode.config_file') +local prompt_guard_indicator = require('opencode.ui.prompt_guard_indicator') local LABELS = { NEW_SESSION_TITLE = 'New session', @@ -40,17 +41,40 @@ local function get_mode_highlight() end end -local function create_winbar_text(description, model_info, mode_info, win_width) - local available_width = win_width - 2 -- 2 padding spaces +local function create_winbar_text(description, model_info, mode_info, show_guard_indicator, win_width) + -- Calculate how many visible characters we have + -- Format: " [GUARD ]description padding model_info MODE " + -- Where [GUARD ] is optional - -- If total length exceeds available width, truncate description - if #description + 1 + #model_info + #mode_info + 1 > available_width then - local space_for_desc = available_width - (#model_info + #mode_info + 1) - 4 -- -4 for "... " - description = description:sub(1, space_for_desc) .. '... ' + local guard_info = '' + local guard_visible_width = 0 + if show_guard_indicator then + guard_info = prompt_guard_indicator.get_formatted() + guard_visible_width = 2 -- icon + space end - local padding = string.rep(' ', available_width - #description - #model_info - #mode_info - 1) - return string.format(' %s%s%s %s ', description, padding, model_info, get_mode_highlight() .. mode_info .. '%*') + -- Calculate used width: leading space + guard + trailing space + model + mode + local mode_info_str = get_mode_highlight() .. mode_info .. '%*' + local mode_visible_width = #mode_info + local model_visible_width = #model_info + + -- Reserve space: 1 (padding) + guard_visible_width (with padding) + model + 1 (space before mode) + mode + 1 (padding) + local reserved_width = 1 + guard_visible_width + model_visible_width + 1 + mode_visible_width + 1 + + -- Available width for description and padding + local available_for_desc = win_width - reserved_width + + -- Truncate description if needed + if #description > available_for_desc then + local space_for_desc = available_for_desc - 4 -- -4 for "... " + description = description:sub(1, space_for_desc) .. '...' + end + + -- Calculate padding to right-align model and mode + local padding_width = available_for_desc - #description + local padding = string.rep(' ', math.max(0, padding_width)) + + return string.format(' %s %s%s%s %s ', guard_info, description, padding, model_info, mode_info_str) end local function update_winbar_highlights(win_id) @@ -96,8 +120,10 @@ function M.render() end -- topbar needs to at least have a value to make sure footer is positioned correctly vim.wo[win].winbar = ' ' + + local show_guard_indicator = prompt_guard_indicator.is_denied() vim.wo[win].winbar = - create_winbar_text(get_session_desc(), format_model_info(), format_mode_info(), vim.api.nvim_win_get_width(win)) + create_winbar_text(get_session_desc(), format_model_info(), format_mode_info(), show_guard_indicator, vim.api.nvim_win_get_width(win)) update_winbar_highlights(win) end) @@ -111,6 +137,7 @@ function M.setup() state.subscribe('current_mode', on_change) state.subscribe('current_model', on_change) state.subscribe('active_session', on_change) + state.subscribe('is_opencode_focused', on_change) M.render() end @@ -118,5 +145,6 @@ function M.close() state.unsubscribe('current_mode', on_change) state.unsubscribe('current_model', on_change) state.unsubscribe('active_session', on_change) + state.unsubscribe('is_opencode_focused', on_change) end return M diff --git a/lua/opencode/util.lua b/lua/opencode/util.lua index 1a6673d7..6d894889 100644 --- a/lua/opencode/util.lua +++ b/lua/opencode/util.lua @@ -360,4 +360,32 @@ function M.parse_dot_args(args_str) return result end +--- Check if prompt is allowed via guard callback +--- @param guard_callback? function +--- @param mentioned_files? string[] List of mentioned files in the context +--- @return boolean allowed +--- @return string|nil error_message +function M.check_prompt_allowed(guard_callback, mentioned_files) + if not guard_callback then + return true, nil -- No guard = always allowed + end + + if not type(guard_callback) == 'function' then + return false, 'prompt_guard must be a function' + end + + mentioned_files = mentioned_files or {} + local success, result = pcall(guard_callback, mentioned_files) + + if not success then + return false, 'prompt_guard error: ' .. tostring(result) + end + + if type(result) ~= 'boolean' then + return false, 'prompt_guard must return a boolean' + end + + return result, nil +end + return M