From c103810a0dd25a454a95a9b8ca94353e4478ac4c Mon Sep 17 00:00:00 2001
From: monkoose <pythonproof@gmail.com>
Date: Thu, 4 Apr 2024 11:23:06 +0300
Subject: [PATCH] Add restart command

---
 README.md                   |  1 +
 lua/neocodeium/commands.lua |  5 +++++
 lua/neocodeium/server.lua   | 35 +++++++++++++++++++++++++++++------
 3 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index f1c2e76..c7b9e40 100644
--- a/README.md
+++ b/README.md
@@ -130,6 +130,7 @@ NeoCodeium provides `:NeoCodeium` user command, which has some useful actions:
 - `:NeoCodeium toggle` - toggles NeoCodeium completion.
 - `:NeoCodeium disable_buffer` - disables NeoCodeium completion in the current buffer.
 - `:NeoCodeium enable_buffer` - enables NeoCodeium completion in the current buffer.
+- `:NeoCodeium restart` - restarts Codeium server (useful when server stop responding for any reason).
 
 #### 🎨 Color groups
 
diff --git a/lua/neocodeium/commands.lua b/lua/neocodeium/commands.lua
index 111872e..cfb0ea7 100644
--- a/lua/neocodeium/commands.lua
+++ b/lua/neocodeium/commands.lua
@@ -6,6 +6,7 @@ local conf = require("neocodeium.utils.conf")
 local log = require("neocodeium.log")
 local options = require("neocodeium.options").options
 local api_key = require("neocodeium.api_key")
+local server = require("neocodeium.server")
 
 local vf = vim.fn
 local json = vim.json
@@ -160,6 +161,10 @@ function commands.enable_buffer()
   vim.b.neocodeium_enabled = true
 end
 
+function commands.restart()
+  server:restart()
+end
+
 function commands.toggle()
   options.enabled = not options.enabled
 end
diff --git a/lua/neocodeium/server.lua b/lua/neocodeium/server.lua
index 3c9f3cd..744ca36 100644
--- a/lua/neocodeium/server.lua
+++ b/lua/neocodeium/server.lua
@@ -4,6 +4,7 @@ local log = require("neocodeium.log")
 local api_key = require("neocodeium.api_key")
 local options = require("neocodeium.options").options
 local stdio = require("neocodeium.utils.stdio")
+local echo = require("neocodeium.utils.echo")
 local utils = require("neocodeium.utils")
 local Bin = require("neocodeium.binary")
 
@@ -17,9 +18,12 @@ local json = vim.json
 ---@class Server
 ---@field bin Binary
 ---@field port? string
----@field handle? uv.uv_handle_t
+---@field handle? uv.uv_process_t
+---@field pid? integer
+---@field is_restart boolean
 local Server = {
   bin = Bin.new(),
+  is_restart = false,
 }
 
 -- Auxiliary functions ------------------------------------- {{{1
@@ -66,7 +70,7 @@ function Server:start()
   local stderr = assert(uv.new_pipe())
 
   ---@diagnostic disable-next-line: missing-fields
-  self.handle = uv.spawn(self.bin.path, {
+  self.handle, self.pid = uv.spawn(self.bin.path, {
     args = args,
     stdio = { stdin, stdout, stderr },
   }, function(_, _)
@@ -77,6 +81,14 @@ function Server:start()
     uv.close(stderr)
     if self.handle then
       uv.close(self.handle)
+      self.handle = nil
+    end
+    self.port = nil
+    if self.is_restart then
+      self.is_restart = false
+      vim.schedule(function()
+        self:run()
+      end)
     end
   end)
 
@@ -110,12 +122,23 @@ function Server:run()
   end
 end
 
----Properly stops the server and closes the handle
+---Properly stops the server, executing self.handle on_exit callback.
 function Server:stop()
-  if self.handle then
-    uv.close(self.handle)
-    self.handle = nil
+  if self.pid then
+    uv.kill(self.pid, "sigint")
+  end
+end
+
+---Restarts the server
+function Server:restart()
+  if self.handle and not uv.is_closing(self.handle) then
+    self.is_restart = true
+    self:stop()
+  else
+    self:run()
   end
+  log.info("Server restarted")
+  echo.info("server restarted")
 end
 
 ---Sends request to the server