local session = require 'bugabinga.pivi.session' local help = require 'bugabinga.pivi.help' local transcript = require 'bugabinga.pivi.transcript' local context = require 'bugabinga.pivi.context' local follow = require 'bugabinga.pivi.follow' local publish = require 'bugabinga.pivi.publish' local runtime = require 'bugabinga.pivi.runtime' local comparison = require 'bugabinga.pivi.diff' local guide = require 'bugabinga.pivi.guide' local layout = require 'bugabinga.pivi.layout' local compose = require 'bugabinga.pivi.compose' local INPUT_FILETYPE = 'pivi-input' local notify = function( message, level ) vim.notify( 'pivi: ' .. message, level or vim.log.levels.INFO ) end local fetch_commands = function( current, callback ) current.request( { type = 'get_commands', }, function( response ) local commands = {} if response.success and response.data then for _, entry in ipairs( response.data.commands or {} ) do table.insert( commands, entry ) end end current.commands = commands if callback then callback( commands ) end end ) end --- Returns the active session, starting one when none exists. --- A missing or unusable Pi is reported, never raised. local require_session = function() local current = session.active() if current then return current end local available, failure = runtime.executable() if not available then notify( failure, vim.log.levels.ERROR ) return nil end local ok, created = pcall( session.create, {} ) if not ok then notify( 'cannot start pi: ' .. tostring( created ), vim.log.levels.ERROR ) return nil end local started = created.client.failure() if started then notify( 'cannot start pi: ' .. started, vim.log.levels.ERROR ) created.stop() return nil end -- Warm the command list so completion has something true to offer. fetch_commands( created ) return created end --- Decides whether `/` opens slash-command entry. --- Only a pivi input buffer with the cursor on the first character qualifies. --- @param bufnr number --- @param col number zero based cursor column local slash_opens_commands = function( bufnr, col ) if not vim.api.nvim_buf_is_valid( bufnr ) then return false end if vim.bo[bufnr].filetype ~= INPUT_FILETYPE then return false end return col == 0 end local command_names = function( current ) local found = {} for _, entry in ipairs( current and current.commands or {} ) do table.insert( found, entry.name ) end table.sort( found ) return found end local knows = function( current, name ) for _, entry in ipairs( current and current.commands or {} ) do if entry.name == name then return true end end return false end --- Runs a Pi command. --- Pi dispatches a command it owns and treats anything else as something the --- user said, so an unknown name would quietly become a message to the model. local run_slash_command = function( current, name, arguments ) if not knows( current, name ) then notify( ( 'pi has no command "%s" here; terminal-only commands are not offered over this channel' ):format( name ), vim.log.levels.WARN ) return false end local message = '/' .. name if arguments and arguments ~= '' then message = message .. ' ' .. arguments end local accepted, rejection = current.submit( message, { capture = false, } ) if not accepted then notify( rejection, vim.log.levels.WARN ) end return accepted end local pick_slash_command = function( current, on_pick ) fetch_commands( current, function( commands ) if #commands == 0 then notify( 'no pi commands available', vim.log.levels.WARN ) return end -- A picker filters the text it is shown, so showing a description here -- means typing a command name competes with every description that -- happens to contain the same letters. vim.ui.select( commands, { prompt = 'pi command', format_item = function( entry ) return '/' .. entry.name end, }, function( choice ) if not choice then return end on_pick( choice ) end ) end ) end --- Opens a modifiable prompt buffer bound to a session. local prompts = {} --- The window showing a session's transcript, when there is one. --- A session is not always a live one: a caller may hold a stand-in, so the --- buffer is only trusted once it looks like a buffer. local transcript_window = function( current ) if type( current.buffer ) ~= 'number' then return nil end if not vim.api.nvim_buf_is_valid( current.buffer ) then return nil end local win = vim.fn.bufwinid( current.buffer ) if win == -1 then return nil end return win end --- Shows a session's prompt, reusing the one it already has. --- A draft survives, because reopening a prompt is asking to continue writing, --- not to start over. --- What `@` can refer to: what is open first, then what is nearby. --- A buffer is what the user is actually looking at, so it outranks a path --- that merely exists. local references = function( cwd ) local found, seen = {}, {} for _, bufnr in ipairs( vim.api.nvim_list_bufs() ) do if vim.api.nvim_buf_is_loaded( bufnr ) and vim.bo[bufnr].buftype == '' then local name = vim.api.nvim_buf_get_name( bufnr ) if name ~= '' then local shown = vim.fn.fnamemodify( name, ':.' ) if not seen[shown] then seen[shown] = true table.insert( found, { path = shown, open = true, modified = vim.bo[bufnr].modified, } ) end end end end local ok, listed = pcall( vim.fn.systemlist, { 'git', '-C', cwd, 'ls-files', } ) if ok and vim.v.shell_error == 0 then for _, path in ipairs( listed ) do if not seen[path] then seen[path] = true table.insert( found, { path = path, open = false, } ) end end end return found end local pick_reference = function( current, on_pick ) local choices = references( current.cwd or vim.uv.cwd() ) if #choices == 0 then notify( 'nothing to refer to', vim.log.levels.WARN ) return end vim.ui.select( choices, { prompt = 'refer to', format_item = function( choice ) return choice.path end, }, function( choice ) if choice then on_pick( choice.path ) end end ) end local open_prompt = function( current ) local existing = prompts[current.name] if existing and vim.api.nvim_buf_is_valid( existing ) then local win = vim.fn.bufwinid( existing ) if win == -1 then win = layout.open_prompt( existing, transcript_window( current ) ) or -1 end if win ~= -1 then vim.api.nvim_set_current_win( win ) vim.cmd.startinsert() end return existing end local bufnr = vim.api.nvim_create_buf( false, true ) vim.bo[bufnr].buftype = 'nofile' vim.bo[bufnr].bufhidden = 'hide' vim.bo[bufnr].filetype = INPUT_FILETYPE vim.b[bufnr].pivi_session = current.name prompts[current.name] = bufnr pcall( vim.api.nvim_buf_set_name, bufnr, 'pivi://' .. current.name .. '/prompt' ) vim.keymap.set( { 'n', 'i', }, '/', function() local col = vim.api.nvim_win_get_cursor( 0 )[2] if not slash_opens_commands( bufnr, col ) then return '/' end vim.schedule( function() pick_slash_command( current, function( choice ) run_slash_command( current, choice.name, nil ) transcript.show( current.buffer ) end ) end ) return '' end, { buffer = bufnr, expr = true, desc = 'pivi slash command entry', } ) -- `@` names a file, as it does in pi itself. It only takes over when it -- starts a word, so an address or a decorator can still be typed. vim.keymap.set( 'i', '@', function() local line = vim.api.nvim_get_current_line() local column = vim.api.nvim_win_get_cursor( 0 )[2] local before = column > 0 and line:sub( column, column ) or '' if before ~= '' and not before:match '%s' then return '@' end vim.schedule( function() pick_reference( current, function( path ) vim.api.nvim_put( { '@' .. path .. ' ', }, 'c', true, true ) vim.cmd.startinsert { bang = true, } end ) end ) return '' end, { buffer = bufnr, expr = true, desc = 'pivi refer to a file', } ) guide.bind( bufnr ) local send = function() local lines = vim.api.nvim_buf_get_lines( bufnr, 0, -1, false ) local text = vim.trim( table.concat( lines, '\n' ) ) if text == '' then return end local accepted, rejection = current.submit( text, { capture = false, } ) if not accepted then notify( rejection, vim.log.levels.WARN ) return end vim.api.nvim_buf_set_lines( bufnr, 0, -1, false, { '', } ) -- A submission the user cannot read the answer to is a dead end. transcript.show( current.buffer ) end vim.keymap.set( { 'n', 'i', }, '(pivi-submit)', send, { buffer = bufnr, desc = 'pivi submit the prompt', } ) -- Accepting a completion and accepting a prompt are the same intent, so a -- key bound to this yields to the popup whenever one is open. vim.keymap.set( 'i', '(pivi-accept)', function() if vim.fn.pumvisible() == 1 then return '' end vim.schedule( send ) return '' end, { buffer = bufnr, expr = true, desc = 'pivi accept the completion, else the prompt', } ) local win = layout.open_prompt( bufnr, transcript_window( current ) ) if win then transcript.decorate( win, current.name ) vim.wo[win].winbar = '%#Title#π ' .. current.name .. '%* %#Comment#' .. ( vim.g.pivi_submit_hint or 'send' ) .. ' / commands%*' if vim.g.pivi_start_insert ~= false then vim.cmd.startinsert() end end return bufnr end local submit_text = function( text, options ) local current = require_session() if not current then return false end local accepted, rejection = current.submit( text, options ) if not accepted then notify( rejection, vim.log.levels.WARN ) return false end transcript.show( current.buffer ) return true end local subcommands = {} subcommands.ask = function( arguments, modifiers ) if arguments == '' then notify( 'ask needs a question', vim.log.levels.WARN ) return end local options = {} if modifiers and modifiers.range and modifiers.range > 0 then options.snapshot = context.snapshot { range = { kind = 'range', first = { modifiers.line1, 0, }, last = { modifiers.line2, 0, }, }, } end submit_text( arguments, options ) end subcommands.steer = function( arguments ) submit_text( arguments, { delivery = 'steer', } ) end subcommands['follow-up'] = function( arguments ) submit_text( arguments, { delivery = 'follow_up', } ) end subcommands.help = function( arguments ) if arguments == '' then notify( 'help needs a question', vim.log.levels.WARN ) return end help.ask( arguments ) end subcommands.command = function( arguments ) local current = require_session() if not current then return end local name, rest = arguments:match '^(%S+)%s*(.*)$' if name then -- The list may not have arrived yet, so it is fetched before judging a name. fetch_commands( current, function() if run_slash_command( current, name, rest ) then transcript.show( current.buffer ) end end ) return end pick_slash_command( current, function( choice ) run_slash_command( current, choice.name, nil ) transcript.show( current.buffer ) end ) end subcommands.start = function( arguments ) local available, failure = runtime.executable() if not available then notify( failure, vim.log.levels.ERROR ) return end local created = session.create { name = arguments ~= '' and arguments or nil, } fetch_commands( created ) publish.set_active( created.name ) transcript.show( created.buffer ) notify( 'session ' .. created.name ) end subcommands.stop = function( arguments ) local target = arguments ~= '' and session.get( arguments ) or session.active() if not target then notify( 'no session', vim.log.levels.WARN ) return end target.stop() notify( 'stopped ' .. target.name ) end subcommands.sessions = function() local names = session.list() if #names == 0 then notify( 'no sessions', vim.log.levels.WARN ) return end vim.ui.select( names, { prompt = 'pi session', }, function( choice ) if not choice then return end session.set_active( choice ) local chosen = session.get( choice ) if chosen then transcript.show( chosen.buffer ) end end ) end subcommands.show = function() local current = session.active() if not current then notify( 'no session', vim.log.levels.WARN ) return end transcript.show( current.buffer ) end subcommands.prompt = function() local current = require_session() if not current then return end open_prompt( current ) end subcommands.abort = function() local current = session.active() if not current then notify( 'no session', vim.log.levels.WARN ) return end current.abort() end subcommands.queue = function() local current = session.active() if not current then notify( 'no session', vim.log.levels.WARN ) return end local queue = current.status.queue or {} local entries = {} for _, item in ipairs( queue.steering or {} ) do table.insert( entries, 'steer: ' .. item ) end for _, item in ipairs( queue.follow_up or {} ) do table.insert( entries, 'follow-up: ' .. item ) end if #entries == 0 then notify 'queue is empty' return end notify( table.concat( entries, '\n' ) ) end subcommands.clear = function() local current = session.active() if not current then notify( 'no session', vim.log.levels.WARN ) return end current.clear_queue( function() notify 'queue cleared' end ) end subcommands.diff = function() local bufnr = vim.api.nvim_get_current_buf() if vim.bo[bufnr].filetype ~= transcript.filetype then notify( 'move the cursor into a transcript first', vim.log.levels.WARN ) return end local changed, _, faithful = comparison.open( bufnr ) if not changed then notify( 'this entry reports no change to compare', vim.log.levels.WARN ) return end if faithful == false then notify( 'the file has moved on since this change', vim.log.levels.WARN ) end end --- Switches model. --- `/model` is a terminal-only command, but the capability itself is on this --- channel, so pivi offers it natively instead of pretending the command works. subcommands.model = function( arguments ) local current = require_session() if not current then return end current.request( { type = 'get_available_models', }, function( response ) if not response.success or not response.data then notify( 'pi did not offer a model list', vim.log.levels.WARN ) return end local models = response.data.models or {} if #models == 0 then notify( 'pi has no models configured', vim.log.levels.WARN ) return end local choose = function( model ) current.request( { type = 'set_model', provider = model.provider, modelId = model.id, }, function( answer ) if answer.success then notify( model.provider .. '/' .. model.id ) else notify( answer.error or 'pi refused the model', vim.log.levels.ERROR ) end end ) end if arguments ~= '' then for _, model in ipairs( models ) do if ( model.provider .. '/' .. model.id ) == arguments or model.id == arguments then choose( model ) return end end notify( 'no model matches ' .. arguments, vim.log.levels.WARN ) return end vim.ui.select( models, { prompt = 'pi model', format_item = function( model ) return model.provider .. '/' .. model.id end, }, function( model ) if model then choose( model ) end end ) end ) end --- Runs a shell command through pi, so its output joins the conversation. subcommands.bang = function( arguments ) local current = require_session() if not current then return end if arguments == '' then notify( 'bang needs a command', vim.log.levels.WARN ) return end local accepted, rejection = current.bash( arguments, false ) if not accepted then notify( rejection, vim.log.levels.WARN ) return end transcript.show( current.buffer ) end --- The same, kept out of the conversation. subcommands['bang!'] = function( arguments ) local current = require_session() if not current then return end if arguments == '' then notify( 'bang needs a command', vim.log.levels.WARN ) return end local accepted, rejection = current.bash( arguments, true ) if not accepted then notify( rejection, vim.log.levels.WARN ) return end transcript.show( current.buffer ) end --- Switches thinking level, the other capability `/thinking` cannot reach here. subcommands.thinking = function( arguments ) local current = require_session() if not current then return end current.request( { type = 'get_available_thinking_levels', }, function( response ) local levels = response.success and response.data and response.data.levels or {} if #levels == 0 then notify( 'this model offers no thinking levels', vim.log.levels.WARN ) return end local choose = function( level ) current.request( { type = 'set_thinking_level', level = level, }, function( answer ) if answer.success then notify( 'thinking ' .. level ) else notify( answer.error or 'pi refused the level', vim.log.levels.ERROR ) end end ) end if arguments ~= '' then if vim.list_contains( levels, arguments ) then choose( arguments ) else notify( 'no such level: ' .. arguments, vim.log.levels.WARN ) end return end vim.ui.select( levels, { prompt = 'thinking', }, function( level ) if level then choose( level ) end end ) end ) end --- Compacts the conversation, optionally saying what to keep. subcommands.compact = function( arguments ) local current = require_session() if not current then return end local request = { type = 'compact', } if arguments ~= '' then request.customInstructions = arguments end notify 'compacting' current.request( request, function( response ) if not response.success then notify( response.error or 'compaction failed', vim.log.levels.ERROR ) return end local data = response.data or {} notify( ( 'compacted %s to about %s tokens' ):format( tostring( data.tokensBefore or '?' ), tostring( data.estimatedTokensAfter or '?' ) ) ) end ) end --- Turns automatic compaction on or off. subcommands.autocompact = function( arguments ) local current = require_session() if not current then return end local wanted = arguments ~= 'off' current.request( { type = 'set_auto_compaction', enabled = wanted, }, function( response ) if response.success then notify( 'automatic compaction ' .. ( wanted and 'on' or 'off' ) ) else notify( response.error or 'pi refused', vim.log.levels.ERROR ) end end ) end --- What this session has cost so far. subcommands.session = function() local current = require_session() if not current then return end current.request( { type = 'get_session_stats', }, function( response ) if not response.success or not response.data then notify( 'pi did not report session statistics', vim.log.levels.WARN ) return end local data = response.data local tokens = data.tokens or {} local usage = data.contextUsage or {} local told = { 'session ' .. tostring( data.sessionId or current.name ), ( 'messages %s, tool calls %s' ):format( tostring( data.totalMessages or '?' ), tostring( data.toolCalls or '?' ) ), ( 'tokens %s in, %s out, %s total' ):format( tostring( tokens.input or 0 ), tostring( tokens.output or 0 ), tostring( tokens.total or 0 ) ), ( 'cost %s' ):format( tostring( data.cost or 0 ) ), } if usage.percent then table.insert( told, ( 'context %s%% of %s' ):format( tostring( usage.percent ), tostring( usage.contextWindow ) ) ) end notify( table.concat( told, '\n' ) ) end ) end --- Names the session, so it can be recognised when resumed. subcommands.name = function( arguments ) local current = require_session() if not current then return end if arguments == '' then notify( 'name needs a name', vim.log.levels.WARN ) return end current.request( { type = 'set_session_name', name = arguments, }, function( response ) if response.success then notify( 'named ' .. arguments ) else notify( response.error or 'pi refused the name', vim.log.levels.ERROR ) end end ) end --- Starts a fresh conversation in the process that is already running. --- `start` spawns another pi; this one does not. subcommands.new = function() local current = require_session() if not current then return end current.request( { type = 'new_session', }, function( response ) if not response.success then notify( response.error or 'pi refused a new session', vim.log.levels.ERROR ) return end if response.data and response.data.cancelled then notify( 'an extension cancelled the new session', vim.log.levels.WARN ) return end transcript.clear( current.buffer ) notify 'new conversation' end ) end --- Writes the conversation out as it is. subcommands.export = function( arguments ) local current = require_session() if not current then return end local request = { type = 'export_html', } if arguments ~= '' then request.outputPath = vim.fn.expand( arguments ) end current.request( request, function( response ) if response.success and response.data then notify( 'exported to ' .. tostring( response.data.path ) ) else notify( response.error or 'export failed', vim.log.levels.ERROR ) end end ) end subcommands.guide = function() guide.show( vim.api.nvim_get_current_buf() ) end subcommands.follow = function( arguments ) follow.toggle( arguments ~= '' and arguments or nil ) notify( 'follow ' .. ( follow.mode() or 'off' ) ) end local names = function() local list = {} for key in pairs( subcommands ) do table.insert( list, key ) end table.sort( list ) return list end local dispatch = function( options ) local input = options.args or '' local name, rest = input:match '^(%S+)%s*(.*)$' -- Bare `:Pi` is the way in: it starts a session when there is none and puts -- the cursor where the user can type. if not name then subcommands.prompt() return end local handler = subcommands[name] if not handler then notify( 'unknown subcommand: ' .. name, vim.log.levels.ERROR ) return end handler( rest or '', options ) end local setup = function() vim.api.nvim_create_user_command( 'Pi', dispatch, { nargs = '*', range = true, desc = 'pivi native Pi integration', complete = function( lead, line ) local matching = function( candidates ) return vim.tbl_filter( function( candidate ) return candidate:sub( 1, #lead ) == lead end, candidates ) end -- Completing an argument needs the session's own knowledge, so the list -- is kept warm and offered as it is, rather than blocking on a request. local subcommand = line:match '^%s*Pi%s+(%S+)%s+%S*$' if subcommand == 'command' then return matching( command_names( session.active() ) ) end if subcommand == 'stop' or subcommand == 'sessions' then return matching( session.list() ) end if subcommand == 'follow' then return matching { 'window', 'tab', 'current', } end if subcommand == 'autocompact' then return matching { 'on', 'off', } end if line:match '^%s*Pi%s+%S*$' == nil then return {} end return matching( names() ) end, } ) vim.keymap.set( 'n', '(pivi-show)', function() subcommands.show() end, { desc = 'pivi show transcript', } ) vim.keymap.set( 'n', '(pivi-prompt)', function() subcommands.prompt() end, { desc = 'pivi open prompt', } ) vim.keymap.set( 'n', '(pivi-abort)', function() subcommands.abort() end, { desc = 'pivi abort run', } ) vim.keymap.set( 'n', '(pivi-follow)', function() subcommands.follow '' end, { desc = 'pivi toggle follow', } ) vim.keymap.set( 'n', '(pivi-guide)', function() subcommands.guide() end, { desc = 'pivi show the guide', } ) end return { input_filetype = INPUT_FILETYPE, setup = setup, dispatch = dispatch, subcommands = subcommands, names = names, slash_opens_commands = slash_opens_commands, open_prompt = open_prompt, }