Luigit
repositories / dotfiles

dotfiles

bugabingas dorkfiles

owned by admin

neovim/lua/bugabinga/pivi/rpc.lua

Raw
local jsonl = require 'bugabinga.pivi.jsonl'

local next_request_id = 0

local new_request_id = function()
  next_request_id = next_request_id + 1
  return 'pivi-' .. tostring( next_request_id )
end

local default_spawn = function( command, options )
  return vim.system( command, {
    cwd = options.cwd,
    env = options.env,
    stdin = true,
    text = true,
    stdout = options.stdout,
    stderr = options.stderr,
  }, options.on_exit )
end

--- Starts a Pi RPC child process.
--- Output arrives in fast contexts, so every handler call is deferred through the main loop.
--- @param options table command, cwd, spawn, on_event, on_response, on_stderr, on_exit, on_decode_error
local start = function( options )
  vim.validate( 'options', options, 'table' )
  vim.validate( 'options.command', options.command, 'table' )

  local spawn = options.spawn or default_spawn
  local splitter = jsonl.splitter()
  local pending = {}
  local state = { running = true, exit_code = nil, }

  local failures = {}

  --- Handler failures never escape into the main loop, where a late error
  --- would raise a message prompt instead of failing a request.
  local guard = function( handler, ... )
    if not handler then return end
    local ok, failure = pcall( handler, ... )
    if not ok then
      local reported = tostring( failure )
      table.insert( failures, reported )
      -- A swallowed handler error is a bug that looks like nothing happening,
      -- so it is reported rather than only recorded.
      if options.on_failure then pcall( options.on_failure, reported ) end
    end
  end

  local dispatch = function( records )
    if vim.in_fast_event() then
      table.insert( failures, 'dispatch ran in a fast event context' )
      return
    end
    for _, record in ipairs( records ) do
      local frame, decode_error = jsonl.parse( record )
      if not frame then
        guard( options.on_decode_error, record, decode_error )
      elseif frame.type == 'response' then
        local callback = frame.id and pending[frame.id]
        if callback then
          pending[frame.id] = nil
          guard( callback, frame )
        end
        guard( options.on_response, frame )
      else
        guard( options.on_event, frame )
      end
    end
  end

  local on_stdout = function( _, data )
    local records = splitter.feed( data )
    if #records == 0 then return end
    vim.schedule( function() dispatch( records ) end )
  end

  local on_stderr = function( _, data )
    if data == nil or data == '' then return end
    if not options.on_stderr then return end
    vim.schedule( function() guard( options.on_stderr, data ) end )
  end

  local on_exit = function( result )
    local trailing = splitter.finish()
    vim.schedule( function()
      state.running = false
      state.exit_code = result and result.code
      if #trailing > 0 then dispatch( trailing ) end
      for id, callback in pairs( pending ) do
        pending[id] = nil
        guard( callback, {
          id = id,
          type = 'response',
          success = false,
          error = 'pi process exited before responding',
        } )
      end
      guard( options.on_exit, state.exit_code )
    end )
  end

  -- A missing or unusable executable is an ordinary reported failure, never a
  -- traceback in the user's editor.
  local started, spawned = pcall( spawn, options.command, {
    cwd = options.cwd,
    env = options.env,
    stdout = on_stdout,
    stderr = on_stderr,
    on_exit = on_exit,
  } )

  local handle = started and spawned or nil

  if not handle then
    state.running = false
    state.failure = tostring( spawned )
    table.insert( failures, state.failure )
    vim.schedule( function() guard( options.on_exit, nil ) end )
  end

  local send = function( command )
    if not state.running or not handle then
      return false, state.failure or 'pi process is not running'
    end
    local ok, encoded = pcall( vim.json.encode, command )
    if not ok then return false, tostring( encoded ) end
    handle:write( encoded .. '\n' )
    return true, nil
  end

  --- Sends a command and correlates its response by id.
  local request = function( command, callback )
    local payload = vim.deepcopy( command )
    payload.id = payload.id or new_request_id()
    if callback then pending[payload.id] = callback end
    local ok, failure = send( payload )
    if not ok then
      pending[payload.id] = nil
      if callback then
        callback( { id = payload.id, type = 'response', success = false, error = failure, } )
      end
    end
    return payload.id
  end

  local notify = function( command ) return send( command ) end

  local GRACE_MS = 2000

  --- Ends the child by closing its input first.
  --- Pi exits on its own when stdin closes; killing immediately would orphan
  --- the runtime behind a launcher shim, which keeps processes alive forever.
  local stop = function()
    if not state.running or not handle then return end

    pcall( function() handle:write( nil ) end )

    local timer = vim.uv.new_timer()
    if not timer then
      pcall( function() handle:kill( 'sigterm' ) end )
      return
    end

    timer:start( GRACE_MS, 0, vim.schedule_wrap( function()
      timer:stop()
      timer:close()
      if state.running then pcall( function() handle:kill( 'sigterm' ) end ) end
    end ) )
  end

  local is_running = function() return state.running end

  return {
    request = request,
    notify = notify,
    stop = stop,
    is_running = is_running,
    failure = function() return state.failure end,
    handle = handle,
    failures = function() return vim.deepcopy( failures ) end,
  }
end

return {
  start = start,
}