Track fire-and-forget goroutines with sync.WaitGroup for clean shutdown

Adds inflight WaitGroup to Server and wraps inline-prompt,
chat-response, and deferShowDocument goroutines. Run() waits
for all in-flight work before returning.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

2 files changed, 18 insertions, 3 deletions

diff --git a/internal/lsp/handlers_document.go b/internal/lsp/handlers_document.go
index a411f4e..f69c626 100644
--- a/internal/lsp/handlers_document.go
+++ b/internal/lsp/handlers_document.go
@@ -124,7 +124,11 @@ func (s *Server) maybeRunInlinePrompt(uri string, lineIdx int, raw string, openS
 	}
 	if s.currentLLMClient() != nil {
 		pos := Position{Line: lineIdx, Character: len(raw)}
-		go s.runInlinePrompt(uri, pos)
+		s.inflight.Add(1)
+		go func() {
+			defer s.inflight.Done()
+			s.runInlinePrompt(uri, pos)
+		}()
 	}
 	return true
 }
@@ -193,7 +197,11 @@ func (s *Server) handleChatPrompt(uri string, lineIdx int, match chatPromptLine)
 		}
 		return
 	}
-	go s.requestChatResponse(uri, lineIdx, match)
+	s.inflight.Add(1)
+	go func() {
+		defer s.inflight.Done()
+		s.requestChatResponse(uri, lineIdx, match)
+	}()
 }
 
 func (s *Server) requestChatResponse(uri string, lineIdx int, match chatPromptLine) {
@@ -432,7 +440,9 @@ func (s *Server) clientShowDocument(uri string, sel *Range) {
 // The goroutine respects s.serverCtx so it won't write after shutdown.
 func (s *Server) deferShowDocument(uri string, sel Range) {
 	ctx := s.serverCtx
+	s.inflight.Add(1)
 	go func() {
+		defer s.inflight.Done()
 		if ctx == nil {
 			time.Sleep(120 * time.Millisecond)
 			s.clientShowDocument(uri, &sel)
diff --git a/internal/lsp/server.go b/internal/lsp/server.go
index 6442342..c039255 100644
--- a/internal/lsp/server.go
+++ b/internal/lsp/server.go
@@ -31,6 +31,7 @@ type Server struct {
 	serverCancel context.CancelFunc
 	statusSink   StatusSink
 	exited       atomic.Bool
+	inflight     sync.WaitGroup // tracks background goroutines (inline prompt, chat, etc.)
 	mu           sync.RWMutex
 	docs         map[string]*document
 	logContext   bool
@@ -332,8 +333,12 @@ func (s *Server) emitGlobalStatus(reqs int64, rpm float64, sent int64, recv int6
 }
 
 // Run starts the server's main loop, reading and dispatching LSP messages until EOF or exit.
+// On shutdown it cancels the server context and waits for in-flight goroutines.
 func (s *Server) Run() error {
-	defer s.cancelRequests()
+	defer func() {
+		s.cancelRequests()
+		s.inflight.Wait()
+	}()
 	for {
 		body, err := s.readMessage()
 		if errors.Is(err, io.EOF) {