feat(cache): only start on init

Author: meeehdi-dev

cache/cache.go

@@ -2,9 +2,7 @@ package cache
 
 import (
 	"crypto/sha256"
-	"fmt"
 	"llm-language-server/lsp"
-	"strings"
 	"time"
 )
 
@@ -13,8 +11,10 @@ type CacheValue struct {
 	Ex   time.Time
 }
 
-var cacheMap = make(map[string]CacheValue)
-var EX = 60 * time.Second
+var cacheMap map[string]CacheValue
+var started = false
+
+var ex = 60 * time.Second
 
 func checkExpiredKeys() {
 	for key, value := range cacheMap {
@@ -25,41 +25,45 @@ func checkExpiredKeys() {
 }
 
 func getKey(prompt string, suffix string) string {
-	cacheKey := fmt.Sprintf("%s<FIM_MIDDLE>%s", prompt, suffix)
 	cacheKeyHash := sha256.New()
-	cacheKeyHash.Write([]byte(cacheKey))
+	cacheKeyHash.Write([]byte(prompt + "_" + suffix))
 	return string(cacheKeyHash.Sum(nil))
 }
 
 func Set(prompt string, suffix string, value []lsp.CompletionItem) {
-	key := getKey(prompt, suffix)
-	cacheMap[key] = CacheValue{Data: value, Ex: time.Now().Add(EX)}
-	for _, item := range value {
-		chars := strings.Split(item.InsertText, "")
-		pendingLine := ""
-		for _, char := range chars {
-			pendingLine += char
-			item.InsertText = item.InsertText[1:]
-			pendingKey := getKey(prompt+pendingLine, suffix)
-			cacheMap[pendingKey] = CacheValue{Data: []lsp.CompletionItem{item}, Ex: time.Now().Add(EX)}
-		}
+	if !started {
+		return
 	}
+
+	key := getKey(prompt, suffix)
+	cacheMap[key] = CacheValue{Data: value, Ex: time.Now().Add(ex)}
 }
 
 func Get(prompt string, suffix string) ([]lsp.CompletionItem, bool) {
+	if !started {
+		return nil, false
+	}
+
 	key := getKey(prompt, suffix)
 	value, exists := cacheMap[key]
 
 	if exists {
-		return value.Data, exists
+		return value.Data, true
 	} else {
-		return nil, exists
+		return nil, false
 	}
 }
 
+func Reset() {
+	cacheMap = make(map[string]CacheValue)
+}
+
 var ticker *time.Ticker
 
 func Init() {
+	Reset()
+	started = true
+
 	ticker = time.NewTicker(time.Second)
 	defer ticker.Stop()
 
@@ -76,7 +80,3 @@ func Shutdown() {
 		ticker.Stop()
 	}
 }
-
-func Reset() {
-	cacheMap = make(map[string]CacheValue)
-}

handler/handler.go

@@ -53,6 +53,9 @@ func HandleRequestMessage(writer *os.File, message []byte) error {
 		if err != nil {
 			return err
 		}
+		if params.InitializationOptions.Cache == nil || *params.InitializationOptions.Cache != false {
+			go cache.Init()
+		}
 		output := lsp.NewInitializeResponse(request.ID)
 		writer.Write(lsp.NewResponseMessage(output))
 		return nil

lsp/initialize.go

@@ -8,6 +8,7 @@ type ClientInfo struct {
 type InitializationOptions struct {
 	Provider string `json:"provider"`
 	Params   any    `json:"params"`
+	Cache    *bool  `json:"cache,omitempty"`
 }
 
 type InitializeParams struct {

main.go

@@ -4,7 +4,6 @@ import (
 	"bufio"
 	"flag"
 	"fmt"
-	"llm-language-server/cache"
 	"llm-language-server/handler"
 	"llm-language-server/jsonrpc"
 	"llm-language-server/lsp"
@@ -28,8 +27,6 @@ func main() {
 	writer := os.Stdout
 	writer.Write(lsp.NewNotificationMessage(lsp.NewLogMesssage(fmt.Sprintf("Successfully started! (debug: %t)", *debug), lsp.Info)))
 
-	go cache.Init()
-
 	for scanner.Scan() {
 		bytes := scanner.Bytes()
 		message, err := jsonrpc.DecodeMessage(bytes)