2 files changed, 39 insertions, 14 deletions
diff --git a/internal/lsp/context.go b/internal/lsp/context.go
index 8b584fb..4395abf 100644
--- a/internal/lsp/context.go
+++ b/internal/lsp/context.go
@@ -7,6 +7,15 @@ import (
 	"codeberg.org/snonux/hexai/internal/logging"
 )
 
+type contextModeBuilder func(*Server, bool, string, Position) (string, bool)
+
+var contextModeRegistry = map[string]contextModeBuilder{
+	"minimal":          minimalModeContext,
+	"window":           windowModeContext,
+	"file-on-new-func": fileOnNewFuncModeContext,
+	"always-full":      alwaysFullModeContext,
+}
+
 // buildAdditionalContext builds extra context messages based on the configured mode.
 // Modes:
 // - minimal: no extra context
@@ -14,23 +23,31 @@ import (
 // - file-on-new-func: include full file only when defining a new function
 // - always-full: always include the full file
 func (s *Server) buildAdditionalContext(newFunc bool, uri string, pos Position) (string, bool) {
-	mode := s.contextMode()
-	switch mode {
-	case "minimal":
-		return "", false
-	case "window":
-		return s.windowContext(uri, pos), true
-	case "file-on-new-func":
-		if newFunc {
-			return s.fullFileContext(uri), true
-		}
-		return "", false
-	case "always-full":
-		return s.fullFileContext(uri), true
-	default:
+	builder, ok := contextModeRegistry[s.contextMode()]
+	if !ok {
 		// fallback to minimal if unknown
 		return "", false
 	}
+	return builder(s, newFunc, uri, pos)
+}
+
+func minimalModeContext(_ *Server, _ bool, _ string, _ Position) (string, bool) {
+	return "", false
+}
+
+func windowModeContext(s *Server, _ bool, uri string, pos Position) (string, bool) {
+	return s.windowContext(uri, pos), true
+}
+
+func fileOnNewFuncModeContext(s *Server, newFunc bool, uri string, _ Position) (string, bool) {
+	if !newFunc {
+		return "", false
+	}
+	return s.fullFileContext(uri), true
+}
+
+func alwaysFullModeContext(s *Server, _ bool, uri string, _ Position) (string, bool) {
+	return s.fullFileContext(uri), true
 }
 
 func (s *Server) windowContext(uri string, pos Position) string {
diff --git a/internal/lsp/context_test.go b/internal/lsp/context_test.go
index 875eec9..fd555cf 100644
--- a/internal/lsp/context_test.go
+++ b/internal/lsp/context_test.go
@@ -61,6 +61,14 @@ func TestBuildAdditionalContext_AlwaysFull(t *testing.T) {
 	}
 }
 
+func TestBuildAdditionalContext_UnknownModeFallsBackToMinimal(t *testing.T) {
+	s := newTestServer()
+	s.cfg.ContextMode = "unknown-mode"
+	if ctx, ok := s.buildAdditionalContext(false, "file:///x.go", Position{}); ok || ctx != "" {
+		t.Fatalf("expected no context for unknown mode; got ok=%v ctx=%q", ok, ctx)
+	}
+}
+
 func TestTruncateToApproxTokens(t *testing.T) {
 	text := strings.Repeat("abcd", 10)     // 40 chars
 	got := truncateToApproxTokens(text, 5) // ~20 chars