From 5d5edc7fe93dcee49d02c6e7a6bf408f277e5e89 Mon Sep 17 00:00:00 2001
From: Ned Palacios <npdoesmc@gmail.com>
Date: Fri, 22 Dec 2023 15:01:57 +0800
Subject: [PATCH] feat: connect to daemon  after lsp is connected

---
 server/daemon/client/client.go  |  4 ++++
 server/lsp_server/lsp_server.go | 19 ++++++++++++-------
 2 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/server/daemon/client/client.go b/server/daemon/client/client.go
index 194530b..a99f65e 100644
--- a/server/daemon/client/client.go
+++ b/server/daemon/client/client.go
@@ -41,6 +41,10 @@ func (c *Client) processIdField() jsonrpc2.CallOption {
 	return jsonrpc2.ExtraField("processId", c.processId)
 }
 
+func (c *Client) IsConnected() bool {
+	return c.rpcConn != nil && c.connState == ConnectedState
+}
+
 func (c *Client) EnsureConnection() error {
 	if c.rpcConn != nil || c.connState != NotConnectedState {
 		return nil
diff --git a/server/lsp_server/lsp_server.go b/server/lsp_server/lsp_server.go
index 57f6f88..383b15a 100644
--- a/server/lsp_server/lsp_server.go
+++ b/server/lsp_server/lsp_server.go
@@ -49,6 +49,18 @@ func (s *LspServer) Handle(ctx context.Context, c *jsonrpc2.Conn, r *jsonrpc2.Re
 
 	switch r.Method {
 	case lsp.MethodInitialize:
+		if !s.daemonClient.IsConnected() {
+			if err := s.daemonClient.Connect(); err != nil {
+				if err := s.daemonClient.EnsureConnection(); err != nil {
+					c.ReplyWithError(ctx, r.ID, &jsonrpc2.Error{
+						Code:    -32002,
+						Message: fmt.Sprintf("Unable to connect to daemon: %s", err.Error()),
+					})
+					return
+				}
+			}
+		}
+
 		c.Reply(ctx, r.ID, lsp.InitializeResult{
 			Capabilities: lsp.ServerCapabilities{
 				TextDocumentSync:   lsp.TextDocumentSyncKindFull,
@@ -199,13 +211,6 @@ func Start() error {
 	)
 
 	lspServer.daemonClient = daemonClient
-
-	if err := daemonClient.Connect(); err != nil {
-		if err := daemonClient.EnsureConnection(); err != nil {
-			return err
-		}
-	}
-
 	exitSignal := make(chan os.Signal, 1)
 	signal.Notify(exitSignal, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)