codeium.nvim/lua/codeium/api.lua

539 lines
13 KiB
Lua

local versions = require("codeium.versions")
local config = require("codeium.config")
local io = require("codeium.io")
local log = require("codeium.log")
local update = require("codeium.update")
local notify = require("codeium.notify")
local util = require("codeium.util")
local enums = require("codeium.enums")
local api_key = nil
local status = {
api_key_error = nil,
}
local function find_port(manager_dir, start_time)
local files = io.readdir(manager_dir)
for _, file in ipairs(files) do
local number = tonumber(file.name, 10)
if file.type == "file" and number and io.stat_mtime(manager_dir .. "/" .. file.name) >= start_time then
return number
end
end
return nil
end
local cookie_generator = 1
local function next_cookie()
cookie_generator = cookie_generator + 1
return cookie_generator
end
local function get_request_metadata(request_id)
return {
api_key = api_key,
ide_name = "neovim",
ide_version = versions.nvim,
extension_name = "neovim",
extension_version = versions.extension,
request_id = request_id or next_cookie(),
}
end
local Server = {}
Server.__index = Server
function Server.check_status()
return status
end
function Server.load_api_key()
local json, err = io.read_json(config.options.config_path)
if err or type(json) ~= "table" then
if err == "ENOENT" then
-- Allow any UI plugins to load
local message = "please log in with :Codeium Auth"
status.api_key_error = message
vim.defer_fn(function()
notify.info(message)
end, 100)
else
local message = "failed to load the api key"
status.api_key_error = message
notify.info(message)
end
api_key = nil
return
end
status.api_key_error = nil
api_key = json.api_key
end
function Server.save_api_key()
local _, result = io.write_json(config.options.config_path, {
api_key = api_key,
})
status.api_key_error = nil
if result then
local message = "failed to save the api key"
status.api_key_error = message
notify.error(message, result)
end
end
function Server.authenticate()
local attempts = 0
local uuid = io.generate_uuid()
local url = "https://"
.. config.options.api.portal_url
.. "/profile?response_type=token&redirect_uri=vim-show-auth-token&state="
.. uuid
.. "&scope=openid%20profile%20email&redirect_parameters_type=query"
local prompt
local function on_submit(value)
if not value then
return
end
local endpoint = "https://api.codeium.com/register_user/"
if config.options.enterprise_mode then
endpoint = "https://" .. config.options.api.host .. ":" .. config.options.api.port
if config.options.api.path then
endpoint = endpoint .. "/" .. config.options.api.path:gsub("^/", "")
end
endpoint = endpoint .. "/exa.seat_management_pb.SeatManagementService/RegisterUser"
end
io.post(endpoint, {
headers = {
accept = "application/json",
},
body = {
firebase_id_token = value,
},
callback = function(body, err)
if err and not err.response then
notify.error("failed to validate token", err)
return
end
local ok, json = pcall(vim.fn.json_decode, body)
if not ok then
notify.error("failed to decode json", json)
return
end
if json and json.api_key and json.api_key ~= "" then
api_key = json.api_key
Server.save_api_key()
notify.info("api key saved")
return
end
attempts = attempts + 1
if attempts == 3 then
notify.error("too many failed attempts")
return
end
notify.error("api key is incorrect")
prompt()
end,
})
end
prompt = function()
require("codeium.views.auth-menu")(url, on_submit)
end
prompt()
end
function Server:new()
local m = {}
setmetatable(m, self)
local o = {}
setmetatable(o, m)
local port = nil
local job = nil
local current_cookie = nil
local workspaces = {}
local healthy = false
local last_heartbeat = nil
local last_heartbeat_error = nil
local function request(fn, payload, callback)
local url = "http://127.0.0.1:" .. port .. "/exa.language_server_pb.LanguageServerService/" .. fn
io.post(url, {
body = payload,
callback = callback,
})
end
local function do_heartbeat()
request("Heartbeat", {
metadata = get_request_metadata(),
}, function(_, err)
last_heartbeat = os.time()
last_heartbeat_error = nil
if err then
notify.warn("heartbeat failed", err)
last_heartbeat_error = err
else
healthy = true
end
end)
end
function m.is_healthy()
return healthy
end
function m.checkhealth(logger)
logger.info("Checking server status")
if m.is_healthy() then
logger.ok("Server is healthy on port: " .. port)
else
logger.warn("Server is unhealthy")
end
logger.info("Language Server binary: " .. update.get_bin_info().bin)
if last_heartbeat == nil then
logger.warn("No heartbeat executed")
else
logger.info("Last heartbeat: " .. os.date("%D %H:%M:%S", last_heartbeat))
if last_heartbeat_error ~= nil then
logger.error(last_heartbeat_error)
else
logger.ok("Heartbeat ok")
end
end
end
function m.start()
m.shutdown()
current_cookie = next_cookie()
if not api_key then
io.timer(1000, 0, m.start)
return
end
local manager_dir = config.manager_path
if not manager_dir then
manager_dir = io.tempdir("codeium/manager")
vim.fn.mkdir(manager_dir, "p")
end
local start_time = io.touch(manager_dir .. "/start")
local function on_exit(_, err)
if not current_cookie then
return
end
healthy = false
if err then
job = nil
current_cookie = nil
notify.error("codeium server crashed", err)
io.timer(1000, 0, function()
log.debug("restarting server after crash")
m.start()
end)
end
end
local function on_output(_, v, j)
log.debug(j.pid .. ": " .. v)
end
local api_server_url = "https://"
.. config.options.api.host
.. ":"
.. config.options.api.port
.. (config.options.api.path and "/" .. config.options.api.path:gsub("^/", "") or "")
local job_args = {
update.get_bin_info().bin,
"--api_server_url",
api_server_url,
"--manager_dir",
manager_dir,
"--record_cortex_telemetry",
"false",
"--file_watch_max_dir_count",
config.options.file_watch_max_dir_count,
enable_handlers = true,
enable_recording = false,
on_exit = on_exit,
on_stdout = on_output,
on_stderr = on_output,
}
if config.options.enable_chat then
table.insert(job_args, "--enable_chat_web_server")
table.insert(job_args, "--enable_chat_client")
end
if config.options.enable_local_search then
table.insert(job_args, "--enable_local_search")
end
if config.options.enable_index_service then
table.insert(job_args, "--enable_index_service")
table.insert(job_args, "--search_max_workspace_file_count")
table.insert(job_args, config.options.search_max_workspace_file_count)
end
if config.options.api.portal_url then
table.insert(job_args, "--portal_url")
table.insert(job_args, "https://" .. config.options.api.portal_url)
end
if config.options.enterprise_mode then
table.insert(job_args, "--enterprise_mode")
end
if config.options.detect_proxy ~= nil then
table.insert(job_args, "--detect_proxy=" .. tostring(config.options.detect_proxy))
end
local job = io.job(job_args)
job:start()
local function start_heartbeat()
io.timer(100, 5000, function(cancel_heartbeat)
if not current_cookie then
cancel_heartbeat()
else
do_heartbeat()
end
end)
end
io.timer(100, 500, function(cancel)
if not current_cookie then
cancel()
return
end
port = find_port(manager_dir, start_time)
if port then
cancel()
start_heartbeat()
end
end)
end
local function noop(...) end
local pending_request = { 0, noop }
function m.request_completion(document, editor_options, other_documents, callback)
pending_request[2](true)
local metadata = get_request_metadata()
local this_pending_request
local complete
complete = function(...)
complete = noop
this_pending_request(false)
callback(...)
end
this_pending_request = function(is_complete)
if pending_request[1] == metadata.request_id then
pending_request = { 0, noop }
end
this_pending_request = noop
request("CancelRequest", {
metadata = get_request_metadata(),
request_id = metadata.request_id,
}, function(_, err)
if err then
log.warn("failed to cancel in-flight request", err)
end
end)
if is_complete then
complete(false, nil)
end
end
pending_request = { metadata.request_id, this_pending_request }
request("GetCompletions", {
metadata = metadata,
editor_options = editor_options,
document = document,
other_documents = other_documents,
}, function(body, err)
if err then
if err.status == 503 or err.status == 408 then
-- Service Unavailable or Timeout error
return complete(false, nil)
end
local ok, json = pcall(vim.fn.json_decode, err.response.body)
if ok and json then
if json.state and json.state.state == "CODEIUM_STATE_INACTIVE" then
if json.state.message then
log.debug("completion request failed", json.state.message)
end
return complete(false, nil)
end
if json.code == "canceled" then
log.debug("completion request cancelled at the server", json.message)
return complete(false, nil)
end
end
notify.error("completion request failed", err)
complete(false, nil)
return
end
local ok, json = pcall(vim.fn.json_decode, body)
if not ok then
notify.error("completion request failed", "invalid JSON:", json)
return
end
log.trace("completion: ", json)
complete(true, json)
end)
return function()
this_pending_request(true)
end
end
function m.accept_completion(completion_id)
request("AcceptCompletion", {
metadata = get_request_metadata(),
completion_id = completion_id,
}, noop)
end
function m.refresh_context()
-- bufnr for current buffer is 0
local bufnr = 0
local line_ending = util.get_newline(bufnr)
local lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, true)
-- Ensure that there is always a newline at the end of the file
table.insert(lines, "")
local text = table.concat(lines, line_ending)
local filetype = vim.bo.filetype
local language = enums.languages[filetype] or enums.languages.unspecified
local doc = {
editor_language = filetype,
language = language,
cursor_offset = 0,
text = text,
line_ending = line_ending,
absolute_uri = util.get_uri(vim.api.nvim_buf_get_name(bufnr)),
workspace_uri = util.get_uri(util.get_project_root()),
}
request("RefreshContextForIdeAction", {
active_document = doc,
}, function(_, err)
if err then
notify.error("failed refresh context: " .. err.out)
return
end
end)
end
function m.add_workspace()
local project_root = util.get_project_root()
-- workspace already tracked by server
if workspaces[project_root] then
return
end
-- unable to track hidden path
for entry in project_root:gmatch("[^/]+") do
if entry:sub(1, 1) == "." then
return
end
end
request("AddTrackedWorkspace", { workspace = project_root }, function(_, err)
if err then
notify.error("failed to add workspace: " .. err.out)
return
end
workspaces[project_root] = true
end)
end
function m.get_chat_ports()
request("GetProcesses", {
metadata = get_request_metadata(),
}, function(body, err)
if err then
notify.error("failed to get chat ports", err)
return
end
local ports = vim.fn.json_decode(body)
local url = "http://127.0.0.1:"
.. ports.chatClientPort
.. "?api_key="
.. api_key
.. "&has_enterprise_extension="
.. (config.options.enterprise_mode and "true" or "false")
.. "&web_server_url=ws://127.0.0.1:"
.. ports.chatWebServerPort
.. "&ide_name=neovim"
.. "&ide_version="
.. versions.nvim
.. "&app_name=codeium.nvim"
.. "&extension_name=codeium.nvim"
.. "&extension_version="
.. versions.extension
.. "&ide_telemetry_enabled=true"
.. "&has_index_service="
.. (config.options.enable_index_service and "true" or "false")
.. "&locale=en_US"
-- cross-platform solution to open the web app
local os_info = io.get_system_info()
if os_info.os == "linux" then
if vim.fn.executable("xdg-open") == 1 then
os.execute("xdg-open '" .. url .. "'")
end
elseif os_info.os == "macos" then
os.execute("open '" .. url .. "'")
elseif os_info.os == "windows" then
os.execute(string.format('start "" "%s"', url))
end
require("codeium.views.chat").open(url)
end)
end
function m.shutdown()
current_cookie = nil
if job then
job.on_exit = nil
job:shutdown()
end
end
m.__index = m
return o
end
return Server