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

4 files changed, 186 insertions, 83 deletions

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"`