Document lock ordering, fix test data races, correct stateMu guard

- Add doc comments clarifying Server.mu and completionState.stateMu are
  independent (no ordering constraint).
- Fix test using wrong lock (s.mu instead of stateMu) for lastLLMCall.
- Replace time.Sleep polling in 7 tests with s.inflight.Wait() to
  eliminate data races under -race.

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

8 files changed, 26 insertions, 42 deletions

diff --git a/internal/lsp/chat_context_mode_test.go b/internal/lsp/chat_context_mode_test.go
index 895c2f3..876f092 100644
--- a/internal/lsp/chat_context_mode_test.go
+++ b/internal/lsp/chat_context_mode_test.go
@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"strings"
 	"testing"
-	"time"
 )
 
 // Ensure in-editor chat respects general.context_mode by adding window/full-file context.
@@ -25,10 +24,7 @@ func TestChat_RespectsContextModeWindow(t *testing.T) {
 	s.setDocument(uri, src)
 
 	s.detectAndHandleChat(uri)
-	// Wait briefly for async goroutine to call Chat
-	for i := 0; i < 40 && len(cap.msgs) == 0; i++ {
-		time.Sleep(10 * time.Millisecond)
-	}
+	s.inflight.Wait()
 	if len(cap.msgs) == 0 {
 		t.Fatalf("expected Chat to be called")
 	}
@@ -64,10 +60,8 @@ func TestChat_ContextModeMinimal_NoExtra(t *testing.T) {
 	uri := "file:///ctx2.go"
 	s.setDocument(uri, "package main\nhelp?>\n")
 	s.detectAndHandleChat(uri)
+	s.inflight.Wait()
 
-	for i := 0; i < 40 && len(cap.msgs) == 0; i++ {
-		time.Sleep(10 * time.Millisecond)
-	}
 	if len(cap.msgs) != 2 {
 		t.Fatalf("expected exactly 2 messages (system + user prompt), got %d", len(cap.msgs))
 	}
@@ -88,10 +82,8 @@ func TestChat_ContextModeAlwaysFull_AddsExtra(t *testing.T) {
 	uri := "file:///ctx3.go"
 	s.setDocument(uri, "package main\nline2\nhelp?>\n")
 	s.detectAndHandleChat(uri)
+	s.inflight.Wait()
 
-	for i := 0; i < 40 && len(cap.msgs) == 0; i++ {
-		time.Sleep(10 * time.Millisecond)
-	}
 	if len(cap.msgs) < 3 {
 		t.Fatalf("expected >=3 messages (system, extra, user prompt), got %d", len(cap.msgs))
 	}
@@ -118,10 +110,8 @@ func TestChat_ContextModeFileOnNewFunc_NoExtraWithoutSignature(t *testing.T) {
 	uri := "file:///ctx4.go"
 	s.setDocument(uri, "package main\nhelp?>\n")
 	s.detectAndHandleChat(uri)
+	s.inflight.Wait()
 
-	for i := 0; i < 40 && len(cap.msgs) == 0; i++ {
-		time.Sleep(10 * time.Millisecond)
-	}
 	if len(cap.msgs) != 2 {
 		t.Fatalf("expected exactly 2 messages (system + user prompt), got %d", len(cap.msgs))
 	}
@@ -141,10 +131,8 @@ func TestChat_ContextModeFileOnNewFunc_WithSignature_AddsExtra(t *testing.T) {
 	src := "package main\n\nfunc add(x int) int\nhelp?>\n"
 	s.setDocument(uri, src)
 	s.detectAndHandleChat(uri)
+	s.inflight.Wait()
 
-	for i := 0; i < 40 && len(cap.msgs) == 0; i++ {
-		time.Sleep(10 * time.Millisecond)
-	}
 	if len(cap.msgs) < 3 {
 		t.Fatalf("expected >=3 messages (system, extra, user prompt), got %d", len(cap.msgs))
 	}
diff --git a/internal/lsp/chat_prompt_test.go b/internal/lsp/chat_prompt_test.go
index 1f7b266..3a1146f 100644
--- a/internal/lsp/chat_prompt_test.go
+++ b/internal/lsp/chat_prompt_test.go
@@ -3,7 +3,6 @@ package lsp
 import (
 	"bytes"
 	"testing"
-	"time"
 )
 
 func TestDetectAndHandleChat_UsesConfiguredSystemPrompt(t *testing.T) {
@@ -20,10 +19,8 @@ func TestDetectAndHandleChat_UsesConfiguredSystemPrompt(t *testing.T) {
 	// Line that should trigger chat: ends with '>' and previous char in prefixes
 	s.setDocument(uri, "help?>\n")
 	s.detectAndHandleChat(uri)
-	// Wait briefly for async goroutine to call Chat
-	for i := 0; i < 20 && len(cap.msgs) == 0; i++ {
-		time.Sleep(10 * time.Millisecond)
-	}
+	// Wait for the background chat goroutine to finish.
+	s.inflight.Wait()
 	if len(cap.msgs) == 0 {
 		t.Fatalf("expected Chat to be called")
 	}
diff --git a/internal/lsp/completion_state.go b/internal/lsp/completion_state.go
index 9f034ab..692eafe 100644
--- a/internal/lsp/completion_state.go
+++ b/internal/lsp/completion_state.go
@@ -6,6 +6,9 @@ import (
 	"time"
 )
 
+// completionState manages the LRU completion cache, pending completions, and
+// throttle timing. Its stateMu is independent of Server.mu — the two locks
+// are never held simultaneously, so there is no ordering constraint.
 type completionState struct {
 	stateMu             sync.RWMutex
 	compCache           map[string]string
diff --git a/internal/lsp/debounce_throttle_more_test.go b/internal/lsp/debounce_throttle_more_test.go
index 7657cab..22d1888 100644
--- a/internal/lsp/debounce_throttle_more_test.go
+++ b/internal/lsp/debounce_throttle_more_test.go
@@ -26,9 +26,9 @@ func TestWaitForThrottle_WaitsRoughlyInterval(t *testing.T) {
 	cfg := s.cfg
 	cfg.CompletionThrottleMs = 20
 	s.cfg = cfg
-	s.mu.Lock()
+	s.stateMu.Lock()
 	s.lastLLMCall = time.Now()
-	s.mu.Unlock()
+	s.stateMu.Unlock()
 	start := time.Now()
 	if !s.waitForThrottle(context.Background()) {
 		t.Fatalf("waitForThrottle returned false")
diff --git a/internal/lsp/document_handlers_test.go b/internal/lsp/document_handlers_test.go
index 1fdd0da..efe6e30 100644
--- a/internal/lsp/document_handlers_test.go
+++ b/internal/lsp/document_handlers_test.go
@@ -6,7 +6,6 @@ import (
 	"io"
 	"log"
 	"testing"
-	"time"
 )
 
 func TestDidOpenChangeClose_UpdateDocs(t *testing.T) {
@@ -67,8 +66,8 @@ func TestDeferShowDocument_WritesLater(t *testing.T) {
 	uri := "file:///x.go"
 	out.Reset()
 	s.deferShowDocument(uri, Range{Start: Position{Line: 0}, End: Position{Line: 0}})
-	// wait >120ms per implementation
-	time.Sleep(160 * time.Millisecond)
+	// Wait for the background goroutine to finish.
+	s.inflight.Wait()
 	req := captureRequest(t, &out)
 	if req.Method != "window/showDocument" {
 		t.Fatalf("expected showDocument, got %s", req.Method)
diff --git a/internal/lsp/handlers_end_to_end_test.go b/internal/lsp/handlers_end_to_end_test.go
index 4528c1d..15b38df 100644
--- a/internal/lsp/handlers_end_to_end_test.go
+++ b/internal/lsp/handlers_end_to_end_test.go
@@ -8,7 +8,6 @@ import (
 	"log"
 	"strings"
 	"testing"
-	"time"
 
 	tut "codeberg.org/snonux/hexai/internal/testutil"
 )
@@ -229,10 +228,8 @@ func TestDetectAndHandleChat_InsertsReply(t *testing.T) {
 	s.setDocument(uri, "What time?>\n\n")
 	out.Reset()
 	s.detectAndHandleChat(uri)
-	// Allow async goroutine to write the request
-	for i := 0; i < 100 && out.Len() == 0; i++ {
-		time.Sleep(10 * time.Millisecond)
-	}
+	// Wait for the background chat goroutine to finish writing.
+	s.inflight.Wait()
 	if out.Len() == 0 {
 		t.Fatalf("no output written by detectAndHandleChat")
 	}
diff --git a/internal/lsp/server.go b/internal/lsp/server.go
index c039255..0cede98 100644
--- a/internal/lsp/server.go
+++ b/internal/lsp/server.go
@@ -32,11 +32,14 @@ type Server struct {
 	statusSink   StatusSink
 	exited       atomic.Bool
 	inflight     sync.WaitGroup // tracks background goroutines (inline prompt, chat, etc.)
-	mu           sync.RWMutex
-	docs         map[string]*document
-	logContext   bool
-	configStore  *runtimeconfig.Store
-	cfg          appconfig.App
+	// mu protects docs, cfg, logContext, configLoadOpts, nextID, and chatSubsystem.lastInput.
+	// It is never held while completionState.stateMu is held, and vice versa,
+	// so there is no lock ordering concern between them.
+	mu          sync.RWMutex
+	docs        map[string]*document
+	logContext  bool
+	configStore *runtimeconfig.Store
+	cfg         appconfig.App
 	codeActionSubsystem
 	chatSubsystem
 	// LLM request stats — atomic to avoid taking the server-wide mu lock.
diff --git a/internal/lsp/triggers_config_test.go b/internal/lsp/triggers_config_test.go
index dbcefd0..7ca3195 100644
--- a/internal/lsp/triggers_config_test.go
+++ b/internal/lsp/triggers_config_test.go
@@ -6,7 +6,6 @@ import (
 	"io"
 	"log"
 	"testing"
-	"time"
 
 	"codeberg.org/snonux/hexai/internal/appconfig"
 )
@@ -75,10 +74,8 @@ func TestDetectAndHandleChat_CustomConfig_InsertsReply(t *testing.T) {
 	s.setDocument(uri, "ok)#\n\n")
 	out.Reset()
 	s.detectAndHandleChat(uri)
-	// Give time for applyEdit request
-	for i := 0; i < 20 && out.Len() == 0; i++ {
-		time.Sleep(10 * time.Millisecond)
-	}
+	// Wait for the background chat goroutine to finish writing.
+	s.inflight.Wait()
 	if out.Len() == 0 {
 		t.Fatalf("no output written for custom chat config")
 	}