From c8bdb0b659cd98bd4ef537d7a77b3ab66645ef19 Mon Sep 17 00:00:00 2001
From: Carlos Alexandro Becker <caarlos0@users.noreply.github.com>
Date: Fri, 26 Sep 2025 10:17:40 -0300
Subject: [PATCH] fix: improve shutdown (#1133)

Signed-off-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com>
---
 internal/llm/agent/mcp-tools.go |  4 +--
 internal/lsp/client.go          | 51 +++++----------------------------
 2 files changed, 9 insertions(+), 46 deletions(-)

diff --git a/internal/llm/agent/mcp-tools.go b/internal/llm/agent/mcp-tools.go
index d670a5797548cd52bbfd23c8cd16fea96b021e8a..a2e6b912ab503c61522501ad522a9f0a65fc37b0 100644
--- a/internal/llm/agent/mcp-tools.go
+++ b/internal/llm/agent/mcp-tools.go
@@ -259,9 +259,9 @@ func updateMCPState(name string, state MCPState, err error, client *client.Clien
 // CloseMCPClients closes all MCP clients. This should be called during application shutdown.
 func CloseMCPClients() error {
 	var errs []error
-	for c := range mcpClients.Seq() {
+	for name, c := range mcpClients.Seq2() {
 		if err := c.Close(); err != nil {
-			errs = append(errs, err)
+			errs = append(errs, fmt.Errorf("close mcp: %s: %w", name, err))
 		}
 	}
 	mcpBroker.Shutdown()
diff --git a/internal/lsp/client.go b/internal/lsp/client.go
index 259f6ba8c4876dcbeb441d48839685012c48ac32..ff9a3ac9b5249663c151fb2df04a4acb168e4de4 100644
--- a/internal/lsp/client.go
+++ b/internal/lsp/client.go
@@ -319,30 +319,6 @@ func (c *Client) NotifyChange(ctx context.Context, filepath string) error {
 	return c.client.NotifyDidChangeTextDocument(ctx, uri, int(fileInfo.Version), changes)
 }
 
-// CloseFile closes a file in the LSP server.
-//
-// NOTE: this is only ever called on LSP shutdown.
-func (c *Client) CloseFile(ctx context.Context, filepath string) error {
-	cfg := config.Get()
-	uri := string(protocol.URIFromPath(filepath))
-
-	if _, exists := c.openFiles.Get(uri); !exists {
-		return nil // Already closed
-	}
-
-	if cfg.Options.DebugLSP {
-		slog.Debug("Closing file", "file", filepath)
-	}
-
-	if err := c.client.NotifyDidCloseTextDocument(ctx, uri); err != nil {
-		return err
-	}
-
-	c.openFiles.Del(uri)
-
-	return nil
-}
-
 // IsFileOpen checks if a file is currently open.
 func (c *Client) IsFileOpen(filepath string) bool {
 	uri := string(protocol.URIFromPath(filepath))
@@ -353,29 +329,16 @@ func (c *Client) IsFileOpen(filepath string) bool {
 // CloseAllFiles closes all currently open files.
 func (c *Client) CloseAllFiles(ctx context.Context) {
 	cfg := config.Get()
-	filesToClose := make([]string, 0, c.openFiles.Len())
-
-	// First collect all URIs that need to be closed
+	debugLSP := cfg != nil && cfg.Options.DebugLSP
 	for uri := range c.openFiles.Seq2() {
-		// Convert URI back to file path using proper URI handling
-		filePath, err := protocol.DocumentURI(uri).Path()
-		if err != nil {
-			slog.Error("Failed to convert URI to path for file closing", "uri", uri, "error", err)
-			continue
+		if debugLSP {
+			slog.Debug("Closing file", "file", uri)
 		}
-		filesToClose = append(filesToClose, filePath)
-	}
-
-	// Then close them all
-	for _, filePath := range filesToClose {
-		err := c.CloseFile(ctx, filePath)
-		if err != nil && cfg != nil && cfg.Options.DebugLSP {
-			slog.Warn("Error closing file", "file", filePath, "error", err)
+		if err := c.client.NotifyDidCloseTextDocument(ctx, uri); err != nil {
+			slog.Warn("Error closing rile", "uri", uri, "error", err)
+			continue
 		}
-	}
-
-	if cfg != nil && cfg.Options.DebugLSP {
-		slog.Debug("Closed all files", "files", filesToClose)
+		c.openFiles.Del(uri)
 	}
 }