Improve LSP diagnostics handling for file operations

Kujtim Hoxha and termai created 7 months ago

- Split LSP file notification into separate functions
- Add waitForLspDiagnostics function to wait for diagnostics after file changes
- Move LSP diagnostics to after file operations in edit and write tools
- Fix string splitting in diff generation
- Reduce diagnostics timeout from 10 to 5 seconds

🤖 Generated with termai
Co-Authored-By: termai <[e-mail]>

Change summary

internal/llm/tools/diagnostics.go | 36 +++++++++++++++++++++++++++-----
internal/llm/tools/edit.go        | 11 +++++----
internal/llm/tools/write.go       |  3 +
3 files changed, 38 insertions(+), 12 deletions(-)

Detailed changes

internal/llm/tools/diagnostics.go

@@ -53,6 +53,7 @@ func (b *diagnosticsTool) Run(ctx context.Context, call ToolCall) (ToolResponse,
 
 	if params.FilePath != "" {
 		notifyLspOpenFile(ctx, params.FilePath, lsps)
+		waitForLspDiagnostics(ctx, params.FilePath, lsps)
 	}
 
 	output := appendDiagnostics(params.FilePath, lsps)
@@ -61,6 +62,22 @@ func (b *diagnosticsTool) Run(ctx context.Context, call ToolCall) (ToolResponse,
 }
 
 func notifyLspOpenFile(ctx context.Context, filePath string, lsps map[string]*lsp.Client) {
+	for _, client := range lsps {
+		// Open the file
+		err := client.OpenFile(ctx, filePath)
+		if err != nil {
+			// If there's an error opening the file, continue to the next client
+			continue
+		}
+	}
+}
+
+// waitForLspDiagnostics opens a file in LSP clients and waits for diagnostics to be published
+func waitForLspDiagnostics(ctx context.Context, filePath string, lsps map[string]*lsp.Client) {
+	if len(lsps) == 0 {
+		return
+	}
+
 	// Create a channel to receive diagnostic notifications
 	diagChan := make(chan struct{}, 1)
 
@@ -92,11 +109,18 @@ func notifyLspOpenFile(ctx context.Context, filePath string, lsps map[string]*ls
 		// Register our temporary handler
 		client.RegisterNotificationHandler("textDocument/publishDiagnostics", handler)
 
-		// Open the file
-		err := client.OpenFile(ctx, filePath)
-		if err != nil {
-			// If there's an error opening the file, continue to the next client
-			continue
+		// Notify change if the file is already open
+		if client.IsFileOpen(filePath) {
+			err := client.NotifyChange(ctx, filePath)
+			if err != nil {
+				continue
+			}
+		} else {
+			// Open the file if it's not already open
+			err := client.OpenFile(ctx, filePath)
+			if err != nil {
+				continue
+			}
 		}
 	}
 
@@ -104,7 +128,7 @@ func notifyLspOpenFile(ctx context.Context, filePath string, lsps map[string]*ls
 	select {
 	case <-diagChan:
 		// Diagnostics received
-	case <-time.After(10 * time.Second):
+	case <-time.After(5 * time.Second):
 		// Timeout after 5 seconds - this is a fallback in case no diagnostics are published
 	case <-ctx.Done():
 		// Context cancelled

internal/llm/tools/edit.go

@@ -74,7 +74,6 @@ func (e *editTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error)
 		params.FilePath = filepath.Join(wd, params.FilePath)
 	}
 
-	notifyLspOpenFile(ctx, params.FilePath, e.lspClients)
 	if params.OldString == "" {
 		result, err := createNewFile(params.FilePath, params.NewString)
 		if err != nil {
@@ -96,6 +95,8 @@ func (e *editTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error)
 		return NewTextErrorResponse(fmt.Sprintf("error replacing content: %s", err)), nil
 	}
 
+	// Wait for LSP diagnostics after editing the file
+	waitForLspDiagnostics(ctx, params.FilePath, e.lspClients)
 	result = fmt.Sprintf("<result>\n%s\n</result>\n", result)
 	result += appendDiagnostics(params.FilePath, e.lspClients)
 	return NewTextResponse(result), nil
@@ -303,23 +304,23 @@ func GenerateDiff(oldContent, newContent string) string {
 	diffs = dmp.DiffCharsToLines(diffs, dmpStrings)
 	diffs = dmp.DiffCleanupSemantic(diffs)
 	buff := strings.Builder{}
-	
+
 	// Add a header to make the diff more readable
 	buff.WriteString("Changes:\n")
-	
+
 	for _, diff := range diffs {
 		text := diff.Text
 
 		switch diff.Type {
 		case diffmatchpatch.DiffInsert:
-			for _, line := range strings.Split(text, "\n") {
+			for line := range strings.SplitSeq(text, "\n") {
 				if line == "" {
 					continue
 				}
 				_, _ = buff.WriteString("+ " + line + "\n")
 			}
 		case diffmatchpatch.DiffDelete:
-			for _, line := range strings.Split(text, "\n") {
+			for line := range strings.SplitSeq(text, "\n") {
 				if line == "" {
 					continue
 				}

internal/llm/tools/write.go

@@ -100,7 +100,6 @@ func (w *writeTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error
 		return NewTextErrorResponse(fmt.Sprintf("Failed to create parent directories: %s", err)), nil
 	}
 
-	notifyLspOpenFile(ctx, filePath, w.lspClients)
 	// Get old content for diff if file exists
 	oldContent := ""
 	if fileInfo != nil && !fileInfo.IsDir() {
@@ -135,6 +134,8 @@ func (w *writeTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error
 	// Record the file write
 	recordFileWrite(filePath)
 	recordFileRead(filePath)
+	// Wait for LSP diagnostics after writing the file
+	waitForLspDiagnostics(ctx, filePath, w.lspClients)
 
 	result := fmt.Sprintf("File successfully written: %s", filePath)
 	result = fmt.Sprintf("<result>\n%s\n</result>", result)