lsp: add 'Resolve diagnostics' code action scoped to selection

- Parse diagnostics from CodeAction context; filter to overlap with selection
- Build LLM prompt from selection-only diagnostics; replace only selected range
- Keep existing 'Rewrite selection' action; return both when applicable
- Add Diagnostic and CodeActionContext types; make CodeActionParams.Context raw JSON
- Add helpers for range overlap; unit tests for filtering/overlap
- Update README to document resolve-diagnostics action

6 files changed, 320 insertions, 197 deletions

diff --git a/README.md b/README.md
index 6cd0827..9aeebca 100644
--- a/README.md
+++ b/README.md
@@ -116,10 +116,13 @@ action from the LLM and then clean up the tag automatically.
 
 ## Code actions
 
-Hexai provides a code action for working with the current selection in Helix:
+Hexai provides code actions that operate only on the current selection in Helix:
 
-- Rewrite selection: Select code and invoke code actions. Hexai looks for the
-  first instruction inside the selection and rewrites the selection accordingly.
+- Rewrite selection: Hexai looks for the first instruction inside the selection
+  and rewrites the selection accordingly.
+- Resolve diagnostics: With a selection active, Hexai gathers only diagnostics
+  that overlap your selection and fixes them by editing only the selected code.
+  Diagnostics outside the selection are not modified.
 
 Instruction sources (first one found wins):
 - Strict marker: `;text;` (no space after first `;`, none before last `;`).
diff --git a/cmd/hexai/main.go b/cmd/hexai/main.go
index 7cd5296..65cbfdf 100644
--- a/cmd/hexai/main.go
+++ b/cmd/hexai/main.go
@@ -1,131 +1,148 @@
 package main
 
 import (
-    "encoding/json"
-    "flag"
-    "log"
-    "os"
-    "path/filepath"
-    "strings"
+	"encoding/json"
+	"flag"
+	"log"
+	"os"
+	"path/filepath"
+	"strings"
 
-    "hexai/internal"
-    "hexai/internal/logging"
-    "hexai/internal/lsp"
-    "hexai/internal/llm"
+	"hexai/internal"
+	"hexai/internal/llm"
+	"hexai/internal/logging"
+	"hexai/internal/lsp"
 )
 
 func main() {
-    var logPath string
-    var showVersion bool
-    flag.StringVar(&logPath, "log", "/tmp/hexai.log", "path to log file (optional)")
-    flag.BoolVar(&showVersion, "version", false, "print version and exit")
-    flag.Parse()
+	logPath := flag.String("log", "/tmp/hexai.log", "path to log file (optional)")
+	showVersion := flag.Bool("version", false, "print version and exit")
+	flag.Parse()
+	if *showVersion {
+		log.Println(internal.Version)
+		return
+	}
 
-    if showVersion {
-        log.Println(internal.Version)
-        return
-    }
+	// Configure logging (path flag only)
+	logger := log.New(os.Stderr, "hexai-lsp ", log.LstdFlags|log.Lmsgprefix)
+	if *logPath != "" {
+		f, err := os.OpenFile(*logPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
+		if err != nil {
+			logger.Fatalf("failed to open log file: %v", err)
+		}
+		defer f.Close()
+		logger.SetOutput(f)
+	}
+	logging.Bind(logger)
 
-    // Configure logging (path flag only)
-    logger := log.New(os.Stderr, "hexai-lsp ", log.LstdFlags|log.Lmsgprefix)
-    if logPath != "" {
-        f, err := os.OpenFile(logPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
-        if err != nil {
-            logger.Fatalf("failed to open log file: %v", err)
-        }
-        defer f.Close()
-        logger.SetOutput(f)
-    }
-    logging.Bind(logger)
+	// Load config file
+	cfg := loadConfig(logger)
 
-    // Load config file
-    cfg := loadConfig(logger)
+	// Normalize and apply logging config
+	cfg.ContextMode = strings.ToLower(strings.TrimSpace(cfg.ContextMode))
+	if cfg.LogPreviewLimit >= 0 {
+		logging.SetLogPreviewLimit(cfg.LogPreviewLimit)
+	}
 
-    // Normalize and apply logging config
-    cfg.ContextMode = strings.ToLower(strings.TrimSpace(cfg.ContextMode))
-    if cfg.LogPreviewLimit >= 0 {
-        logging.SetLogPreviewLimit(cfg.LogPreviewLimit)
-    }
+	// Build LLM client from config (only OPENAI_API_KEY may come from env)
+	var client llm.Client
+	{
+		llmCfg := llm.Config{
+			Provider:      cfg.Provider,
+			OpenAIBaseURL: cfg.OpenAIBaseURL,
+			OpenAIModel:   cfg.OpenAIModel,
+			OllamaBaseURL: cfg.OllamaBaseURL,
+			OllamaModel:   cfg.OllamaModel,
+		}
+		oaKey := os.Getenv("OPENAI_API_KEY")
+		if c, err := llm.NewFromConfig(llmCfg, oaKey); err != nil {
+			logging.Logf("lsp ", "llm disabled: %v", err)
+		} else {
+			client = c
+			logging.Logf("lsp ", "llm enabled provider=%s model=%s", c.Name(), c.DefaultModel())
+		}
+	}
 
-    // Build LLM client from config (only OPENAI_API_KEY may come from env)
-    var client llm.Client
-    {
-        llmCfg := llm.Config{
-            Provider:      cfg.Provider,
-            OpenAIBaseURL: cfg.OpenAIBaseURL,
-            OpenAIModel:   cfg.OpenAIModel,
-            OllamaBaseURL: cfg.OllamaBaseURL,
-            OllamaModel:   cfg.OllamaModel,
-        }
-        oaKey := os.Getenv("OPENAI_API_KEY")
-        if c, err := llm.NewFromConfig(llmCfg, oaKey); err != nil {
-            logging.Logf("lsp ", "llm disabled: %v", err)
-        } else {
-            client = c
-            logging.Logf("lsp ", "llm enabled provider=%s model=%s", c.Name(), c.DefaultModel())
-        }
-    }
-
-    server := lsp.NewServer(os.Stdin, os.Stdout, logger, logPath != "", cfg.MaxTokens, cfg.ContextMode, cfg.ContextWindowLines, cfg.MaxContextTokens, cfg.NoDiskIO, client)
-    if err := server.Run(); err != nil {
-        logger.Fatalf("server error: %v", err)
-    }
+	server := lsp.NewServer(os.Stdin, os.Stdout, logger, *logPath != "", cfg.MaxTokens, cfg.ContextMode, cfg.ContextWindowLines, cfg.MaxContextTokens, cfg.NoDiskIO, client)
+	if err := server.Run(); err != nil {
+		logger.Fatalf("server error: %v", err)
+	}
 }
 
 // appConfig holds user-configurable settings.
 type appConfig struct {
-    MaxTokens           int    `json:"max_tokens"`
-    ContextMode         string `json:"context_mode"`
-    ContextWindowLines  int    `json:"context_window_lines"`
-    MaxContextTokens    int    `json:"max_context_tokens"`
-    LogPreviewLimit     int    `json:"log_preview_limit"`
-    NoDiskIO            bool   `json:"no_disk_io"`
-    Provider            string `json:"provider"`
-    // Provider-specific options
-    OpenAIBaseURL       string `json:"openai_base_url"`
-    OpenAIModel         string `json:"openai_model"`
-    OllamaBaseURL       string `json:"ollama_base_url"`
-    OllamaModel         string `json:"ollama_model"`
+	MaxTokens          int    `json:"max_tokens"`
+	ContextMode        string `json:"context_mode"`
+	ContextWindowLines int    `json:"context_window_lines"`
+	MaxContextTokens   int    `json:"max_context_tokens"`
+	LogPreviewLimit    int    `json:"log_preview_limit"`
+	NoDiskIO           bool   `json:"no_disk_io"`
+	Provider           string `json:"provider"`
+	// Provider-specific options
+	OpenAIBaseURL string `json:"openai_base_url"`
+	OpenAIModel   string `json:"openai_model"`
+	OllamaBaseURL string `json:"ollama_base_url"`
+	OllamaModel   string `json:"ollama_model"`
 }
 
 func loadConfig(logger *log.Logger) appConfig {
-    // Defaults (mirror prior sensible values)
-    cfg := appConfig{
-        MaxTokens:          4000,
-        ContextMode:        "always-full",
-        ContextWindowLines: 120,
-        MaxContextTokens:   4000,
-        LogPreviewLimit:    100,
-        NoDiskIO:           true,
-    }
-    home, err := os.UserHomeDir()
-    if err != nil {
-        return cfg
-    }
-    path := filepath.Join(home, ".config", "hexai", "config.json")
-    f, err := os.Open(path)
-    if err != nil {
-        return cfg
-    }
-    defer f.Close()
-    dec := json.NewDecoder(f)
-    var fileCfg appConfig
-    if err := dec.Decode(&fileCfg); err != nil {
-        logger.Printf("invalid config file %s: %v", path, err)
-        return cfg
-    }
-    // Merge: file overrides defaults when provided
-    if fileCfg.MaxTokens > 0 { cfg.MaxTokens = fileCfg.MaxTokens }
-    if strings.TrimSpace(fileCfg.ContextMode) != "" { cfg.ContextMode = fileCfg.ContextMode }
-    if fileCfg.ContextWindowLines > 0 { cfg.ContextWindowLines = fileCfg.ContextWindowLines }
-    if fileCfg.MaxContextTokens > 0 { cfg.MaxContextTokens = fileCfg.MaxContextTokens }
-    if fileCfg.LogPreviewLimit >= 0 { cfg.LogPreviewLimit = fileCfg.LogPreviewLimit }
-    cfg.NoDiskIO = fileCfg.NoDiskIO
-    if strings.TrimSpace(fileCfg.Provider) != "" { cfg.Provider = fileCfg.Provider }
-    // Provider-specific options
-    if strings.TrimSpace(fileCfg.OpenAIBaseURL) != "" { cfg.OpenAIBaseURL = fileCfg.OpenAIBaseURL }
-    if strings.TrimSpace(fileCfg.OpenAIModel) != "" { cfg.OpenAIModel = fileCfg.OpenAIModel }
-    if strings.TrimSpace(fileCfg.OllamaBaseURL) != "" { cfg.OllamaBaseURL = fileCfg.OllamaBaseURL }
-    if strings.TrimSpace(fileCfg.OllamaModel) != "" { cfg.OllamaModel = fileCfg.OllamaModel }
-    return cfg
+	// Defaults (mirror prior sensible values)
+	cfg := appConfig{
+		MaxTokens:          4000,
+		ContextMode:        "always-full",
+		ContextWindowLines: 120,
+		MaxContextTokens:   4000,
+		LogPreviewLimit:    100,
+		NoDiskIO:           true,
+	}
+	home, err := os.UserHomeDir()
+	if err != nil {
+		return cfg
+	}
+	path := filepath.Join(home, ".config", "hexai", "config.json")
+	f, err := os.Open(path)
+	if err != nil {
+		return cfg
+	}
+	defer f.Close()
+	dec := json.NewDecoder(f)
+	var fileCfg appConfig
+	if err := dec.Decode(&fileCfg); err != nil {
+		logger.Printf("invalid config file %s: %v", path, err)
+		return cfg
+	}
+	// Merge: file overrides defaults when provided
+	if fileCfg.MaxTokens > 0 {
+		cfg.MaxTokens = fileCfg.MaxTokens
+	}
+	if strings.TrimSpace(fileCfg.ContextMode) != "" {
+		cfg.ContextMode = fileCfg.ContextMode
+	}
+	if fileCfg.ContextWindowLines > 0 {
+		cfg.ContextWindowLines = fileCfg.ContextWindowLines
+	}
+	if fileCfg.MaxContextTokens > 0 {
+		cfg.MaxContextTokens = fileCfg.MaxContextTokens
+	}
+	if fileCfg.LogPreviewLimit >= 0 {
+		cfg.LogPreviewLimit = fileCfg.LogPreviewLimit
+	}
+	cfg.NoDiskIO = fileCfg.NoDiskIO
+	if strings.TrimSpace(fileCfg.Provider) != "" {
+		cfg.Provider = fileCfg.Provider
+	}
+	// Provider-specific options
+	if strings.TrimSpace(fileCfg.OpenAIBaseURL) != "" {
+		cfg.OpenAIBaseURL = fileCfg.OpenAIBaseURL
+	}
+	if strings.TrimSpace(fileCfg.OpenAIModel) != "" {
+		cfg.OpenAIModel = fileCfg.OpenAIModel
+	}
+	if strings.TrimSpace(fileCfg.OllamaBaseURL) != "" {
+		cfg.OllamaBaseURL = fileCfg.OllamaBaseURL
+	}
+	if strings.TrimSpace(fileCfg.OllamaModel) != "" {
+		cfg.OllamaModel = fileCfg.OllamaModel
+	}
+	return cfg
 }
diff --git a/internal/lsp/handlers.go b/internal/lsp/handlers.go
index 9d0e672..dce0b8d 100644
--- a/internal/lsp/handlers.go
+++ b/internal/lsp/handlers.go
@@ -65,10 +65,6 @@ func (s *Server) handleCodeAction(req Request) {
         if len(req.ID) != 0 { s.reply(req.ID, []CodeAction{}, nil) }
         return
     }
-    if s.llmClient == nil {
-        if len(req.ID) != 0 { s.reply(req.ID, []CodeAction{}, nil) }
-        return
-    }
     // Extract selected text
     d := s.getDocument(p.TextDocument.URI)
     if d == nil || len(d.lines) == 0 {
@@ -76,37 +72,62 @@ func (s *Server) handleCodeAction(req Request) {
         return
     }
     sel := extractRangeText(d, p.Range)
-    if strings.TrimSpace(sel) == "" {
-        if len(req.ID) != 0 { s.reply(req.ID, []CodeAction{}, nil) }
-        return
-    }
-    // Derive instruction from selection comments (prefer first), including ;text; marker
-    instr, cleaned := instructionFromSelection(sel)
-    if strings.TrimSpace(instr) == "" {
-        // No instruction; do not offer an action
+    if strings.TrimSpace(sel) == "" || s.llmClient == nil {
         if len(req.ID) != 0 { s.reply(req.ID, []CodeAction{}, nil) }
         return
     }
-    // Build prompt for rewrite of cleaned selection according to instruction
-    sys := "You are a precise code refactoring engine. Rewrite the given code strictly according to the instruction. Return only the updated code with no prose or backticks. Preserve formatting where reasonable."
-    user := fmt.Sprintf("Instruction: %s\n\nSelected code to transform:\n%s", instr, cleaned)
-    ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-    defer cancel()
-    messages := []llm.Message{{Role: "system", Content: sys}, {Role: "user", Content: user}}
-    text, err := s.llmClient.Chat(ctx, messages, llm.WithMaxTokens(s.maxTokens), llm.WithTemperature(0.1))
-    if err != nil {
-        logging.Logf("lsp ", "codeAction llm error: %v", err)
-        if len(req.ID) != 0 { s.reply(req.ID, []CodeAction{}, nil) }
-        return
+
+    actions := make([]CodeAction, 0, 2)
+
+    // Action 1: Rewrite selection based on first instruction in selection
+    if instr, cleaned := instructionFromSelection(sel); strings.TrimSpace(instr) != "" {
+        sys := "You are a precise code refactoring engine. Rewrite the given code strictly according to the instruction. Return only the updated code with no prose or backticks. Preserve formatting where reasonable."
+        user := fmt.Sprintf("Instruction: %s\n\nSelected code to transform:\n%s", instr, cleaned)
+        ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+        defer cancel()
+        messages := []llm.Message{{Role: "system", Content: sys}, {Role: "user", Content: user}}
+        if text, err := s.llmClient.Chat(ctx, messages, llm.WithMaxTokens(s.maxTokens), llm.WithTemperature(0.1)); err == nil {
+            out := strings.TrimSpace(text)
+            if out != "" {
+                edit := WorkspaceEdit{Changes: map[string][]TextEdit{p.TextDocument.URI: {{Range: p.Range, NewText: out}}}}
+                actions = append(actions, CodeAction{Title: "Hexai: rewrite selection", Kind: "refactor.rewrite", Edit: &edit})
+            }
+        } else {
+            logging.Logf("lsp ", "codeAction rewrite llm error: %v", err)
+        }
     }
-    out := strings.TrimSpace(text)
-    if out == "" {
-        if len(req.ID) != 0 { s.reply(req.ID, []CodeAction{}, nil) }
-        return
+
+    // Action 2: Resolve diagnostics within selection
+    if diags := s.diagnosticsInRange(p.Context, p.Range); len(diags) > 0 {
+        // Compose a prompt listing diagnostics relevant to the selected code
+        sys := "You are a precise code fixer. Resolve the given diagnostics by editing only the selected code. Return only the corrected code with no prose or backticks. Keep behavior and style, and avoid unrelated changes."
+        var b strings.Builder
+        b.WriteString("Diagnostics to resolve (selection only):\n")
+        for i, dgn := range diags {
+            // Minimal, user-facing summary; include source if present
+            if dgn.Source != "" {
+                fmt.Fprintf(&b, "%d. [%s] %s\n", i+1, dgn.Source, dgn.Message)
+            } else {
+                fmt.Fprintf(&b, "%d. %s\n", i+1, dgn.Message)
+            }
+        }
+        b.WriteString("\nSelected code:\n")
+        b.WriteString(sel)
+        ctx, cancel := context.WithTimeout(context.Background(), 12*time.Second)
+        defer cancel()
+        messages := []llm.Message{{Role: "system", Content: sys}, {Role: "user", Content: b.String()}}
+        if text, err := s.llmClient.Chat(ctx, messages, llm.WithMaxTokens(s.maxTokens), llm.WithTemperature(0.1)); err == nil {
+            out := strings.TrimSpace(text)
+            if out != "" {
+                edit := WorkspaceEdit{Changes: map[string][]TextEdit{p.TextDocument.URI: {{Range: p.Range, NewText: out}}}}
+                actions = append(actions, CodeAction{Title: "Hexai: resolve diagnostics", Kind: "quickfix", Edit: &edit})
+            }
+        } else {
+            logging.Logf("lsp ", "codeAction diagnostics llm error: %v", err)
+        }
     }
-    edit := WorkspaceEdit{Changes: map[string][]TextEdit{p.TextDocument.URI: {{Range: p.Range, NewText: out}}}}
-    action := CodeAction{Title: "Hexai: rewrite selection", Kind: "refactor.rewrite", Edit: &edit}
-    if len(req.ID) != 0 { s.reply(req.ID, []CodeAction{action}, nil) }
+
+    if len(req.ID) != 0 { s.reply(req.ID, actions, nil) }
 }
 
 // instructionFromSelection extracts the first instruction from selection text.
@@ -194,6 +215,45 @@ func findStrictSemicolonTag(line string) (string, int, int, bool) {
     return "", 0, 0, false
 }
 
+// diagnosticsInRange parses the CodeAction context and returns diagnostics
+// that overlap the given selection range. If the context is missing or does
+// not contain diagnostics, returns an empty slice.
+func (s *Server) diagnosticsInRange(ctxRaw json.RawMessage, sel Range) []Diagnostic {
+    if len(ctxRaw) == 0 { return nil }
+    var ctx CodeActionContext
+    if err := json.Unmarshal(ctxRaw, &ctx); err != nil { return nil }
+    if len(ctx.Diagnostics) == 0 { return nil }
+    out := make([]Diagnostic, 0, len(ctx.Diagnostics))
+    for _, d := range ctx.Diagnostics {
+        if rangesOverlap(d.Range, sel) {
+            out = append(out, d)
+        }
+    }
+    return out
+}
+
+// rangesOverlap reports whether two LSP ranges overlap at all.
+func rangesOverlap(a, b Range) bool {
+    // Normalize ordering
+    if greaterPos(a.Start, a.End) { a.Start, a.End = a.End, a.Start }
+    if greaterPos(b.Start, b.End) { b.Start, b.End = b.End, b.Start }
+    // a ends before b starts
+    if lessPos(a.End, b.Start) { return false }
+    // b ends before a starts
+    if lessPos(b.End, a.Start) { return false }
+    return true
+}
+
+func lessPos(p, q Position) bool {
+    if p.Line != q.Line { return p.Line < q.Line }
+    return p.Character < q.Character
+}
+
+func greaterPos(p, q Position) bool {
+    if p.Line != q.Line { return p.Line > q.Line }
+    return p.Character > q.Character
+}
+
 // extractRangeText returns the exact text within the given document range.
 func extractRangeText(d *document, r Range) string {
     if r.Start.Line == r.End.Line {
diff --git a/internal/lsp/handlers_test.go b/internal/lsp/handlers_test.go
index 1b5080a..613835a 100644
--- a/internal/lsp/handlers_test.go
+++ b/internal/lsp/handlers_test.go
@@ -1,6 +1,7 @@
 package lsp
 
 import (
+    "encoding/json"
     "strings"
     "testing"
 )
@@ -254,3 +255,32 @@ func TestStripDuplicateAssignmentPrefix(t *testing.T) {
         t.Fatalf("dup strip '=' failed: got %q", got2)
     }
 }
+
+func TestRangesOverlap(t *testing.T) {
+    a := Range{Start: Position{Line: 1, Character: 2}, End: Position{Line: 3, Character: 0}}
+    b := Range{Start: Position{Line: 2, Character: 0}, End: Position{Line: 4, Character: 1}}
+    if !rangesOverlap(a, b) { t.Fatalf("expected overlap") }
+    c := Range{Start: Position{Line: 4, Character: 1}, End: Position{Line: 5, Character: 0}}
+    if rangesOverlap(a, c) { t.Fatalf("expected no overlap") }
+}
+
+func TestDiagnosticsInRange_Filtering(t *testing.T) {
+    s := newTestServer()
+    sel := Range{Start: Position{Line: 10, Character: 0}, End: Position{Line: 12, Character: 5}}
+    // Build a fake context payload with three diagnostics: one inside, one outside, one touching boundary
+    ctx := CodeActionContext{Diagnostics: []Diagnostic{
+        {Range: Range{Start: Position{Line: 11, Character: 0}, End: Position{Line: 11, Character: 10}}, Message: "inside"},
+        {Range: Range{Start: Position{Line: 2, Character: 0}, End: Position{Line: 3, Character: 0}}, Message: "outside"},
+        {Range: Range{Start: Position{Line: 12, Character: 5}, End: Position{Line: 12, Character: 8}}, Message: "touch"},
+    }}
+    data, _ := json.Marshal(ctx)
+    got := s.diagnosticsInRange(json.RawMessage(data), sel)
+    if len(got) != 2 {
+        t.Fatalf("expected 2 diagnostics in range, got %d", len(got))
+    }
+    msgs := []string{got[0].Message, got[1].Message}
+    joined := strings.Join(msgs, ",")
+    if !strings.Contains(joined, "inside") || !strings.Contains(joined, "touch") {
+        t.Fatalf("unexpected diagnostics: %v", msgs)
+    }
+}
diff --git a/internal/lsp/server.go b/internal/lsp/server.go
index bfdbca2..c6c3812 100644
--- a/internal/lsp/server.go
+++ b/internal/lsp/server.go
@@ -1,62 +1,62 @@
 package lsp
 
 import (
-    "bufio"
-    "encoding/json"
-    "hexai/internal/llm"
-    "hexai/internal/logging"
-    "io"
-    "log"
-    "sync"
-    "time"
+	"bufio"
+	"encoding/json"
+	"hexai/internal/llm"
+	"hexai/internal/logging"
+	"io"
+	"log"
+	"sync"
+	"time"
 )
 
 // Server implements a minimal LSP over stdio.
 type Server struct {
-	in         *bufio.Reader
-	out        io.Writer
-	logger     *log.Logger
-	exited     bool
-	mu         sync.RWMutex
-	docs       map[string]*document
-	logContext bool
-    llmClient  llm.Client
-    lastInput  time.Time
-    maxTokens  int
-    contextMode      string
-    windowLines      int
-    maxContextTokens int
-    noDiskIO         bool
-    // LLM request stats
-    llmReqTotal        int64
-    llmSentBytesTotal  int64
-    llmRespTotal       int64
-    llmRespBytesTotal  int64
-    startTime          time.Time
+	in               *bufio.Reader
+	out              io.Writer
+	logger           *log.Logger
+	exited           bool
+	mu               sync.RWMutex
+	docs             map[string]*document
+	logContext       bool
+	llmClient        llm.Client
+	lastInput        time.Time
+	maxTokens        int
+	contextMode      string
+	windowLines      int
+	maxContextTokens int
+	noDiskIO         bool
+	// LLM request stats
+	llmReqTotal       int64
+	llmSentBytesTotal int64
+	llmRespTotal      int64
+	llmRespBytesTotal int64
+	startTime         time.Time
 }
 
 func NewServer(r io.Reader, w io.Writer, logger *log.Logger, logContext bool, maxTokens int, contextMode string, windowLines int, maxContextTokens int, noDiskIO bool, client llm.Client) *Server {
-    s := &Server{in: bufio.NewReader(r), out: w, logger: logger, docs: make(map[string]*document), logContext: logContext}
-    if maxTokens <= 0 {
-        maxTokens = 500
-    }
-    s.maxTokens = maxTokens
-    if contextMode == "" {
-        contextMode = "file-on-new-func"
-    }
-    if windowLines <= 0 {
-        windowLines = 120
-    }
-    if maxContextTokens <= 0 {
-        maxContextTokens = 2000
-    }
-    s.contextMode = contextMode
-    s.windowLines = windowLines
-    s.maxContextTokens = maxContextTokens
-    s.noDiskIO = noDiskIO
-    s.startTime = time.Now()
-    s.llmClient = client
-    return s
+	s := &Server{in: bufio.NewReader(r), out: w, logger: logger, docs: make(map[string]*document), logContext: logContext}
+	if maxTokens <= 0 {
+		maxTokens = 500
+	}
+	s.maxTokens = maxTokens
+	if contextMode == "" {
+		contextMode = "file-on-new-func"
+	}
+	if windowLines <= 0 {
+		windowLines = 120
+	}
+	if maxContextTokens <= 0 {
+		maxContextTokens = 2000
+	}
+	s.contextMode = contextMode
+	s.windowLines = windowLines
+	s.maxContextTokens = maxContextTokens
+	s.noDiskIO = noDiskIO
+	s.startTime = time.Now()
+	s.llmClient = client
+	return s
 }
 
 func (s *Server) Run() error {
@@ -70,9 +70,9 @@ func (s *Server) Run() error {
 		}
 		var req Request
 		if err := json.Unmarshal(body, &req); err != nil {
-            logging.Logf("lsp ", "invalid JSON: %v", err)
-            continue
-        }
+			logging.Logf("lsp ", "invalid JSON: %v", err)
+			continue
+		}
 		if req.Method == "" {
 			// A response from client; ignore
 			continue
diff --git a/internal/lsp/types.go b/internal/lsp/types.go
index e41371d..00e483e 100644
--- a/internal/lsp/types.go
+++ b/internal/lsp/types.go
@@ -113,7 +113,7 @@ type CompletionParams struct {
 type CodeActionParams struct {
     TextDocument TextDocumentIdentifier `json:"textDocument"`
     Range        Range                  `json:"range"`
-    Context      any                    `json:"context,omitempty"`
+    Context      json.RawMessage        `json:"context,omitempty"`
 }
 
 type WorkspaceEdit struct {
@@ -126,6 +126,19 @@ type CodeAction struct {
     Edit  *WorkspaceEdit `json:"edit,omitempty"`
 }
 
+// Diagnostics (subset needed for code action context)
+type Diagnostic struct {
+    Range    Range       `json:"range"`
+    Message  string      `json:"message"`
+    Severity int         `json:"severity,omitempty"`
+    Code     interface{} `json:"code,omitempty"`
+    Source   string      `json:"source,omitempty"`
+}
+
+type CodeActionContext struct {
+    Diagnostics []Diagnostic `json:"diagnostics"`
+}
+
 // Range defines a text range in a document.
 type Range struct {
 	Start Position `json:"start"`