lsp: add 'Hexai: document code' action to add doc comments to selected code

2 files changed, 40 insertions, 1 deletions

diff --git a/docs/usage-examples.md b/docs/usage-examples.md
index 3e96ade..9abb0b8 100644
--- a/docs/usage-examples.md
+++ b/docs/usage-examples.md
@@ -66,6 +66,7 @@ Operate on the current selection in Helix:
 - Rewrite selection: finds the first instruction inside the selection and rewrites accordingly.
 - Resolve diagnostics: gathers only diagnostics overlapping the selection and fixes them by editing the selected code; diagnostics outside the selection are not changed.
 - Implement unit test (Go): when editing a `.go` file, adds a code action to generate a unit test for the function under the cursor. If `<file>_test.go` exists, appends a new `Test*`; otherwise creates the test file with `package` and `import "testing"`.
+- Document code: adds idiomatic documentation comments to the selected code, preserving behavior and returning only the documented code.
 
 Instruction sources (first match wins):
 
diff --git a/internal/lsp/handlers_codeaction.go b/internal/lsp/handlers_codeaction.go
index ad11861..5740264 100644
--- a/internal/lsp/handlers_codeaction.go
+++ b/internal/lsp/handlers_codeaction.go
@@ -30,13 +30,16 @@ func (s *Server) handleCodeAction(req Request) {
 	}
     sel := extractRangeText(d, p.Range)
 
-    actions := make([]CodeAction, 0, 3)
+    actions := make([]CodeAction, 0, 4)
     if a := s.buildRewriteCodeAction(p, sel); a != nil {
         actions = append(actions, *a)
     }
     if a := s.buildDiagnosticsCodeAction(p, sel); a != nil {
         actions = append(actions, *a)
     }
+    if a := s.buildDocumentCodeAction(p, sel); a != nil {
+        actions = append(actions, *a)
+    }
     if a := s.buildGoUnitTestCodeAction(p); a != nil {
         actions = append(actions, *a)
     }
@@ -136,6 +139,22 @@ func (s *Server) resolveCodeAction(ca CodeAction) (CodeAction, bool) {
 		} else {
 			logging.Logf("lsp ", "codeAction diagnostics llm error: %v", err)
 		}
+    case "document":
+        sys := "You are a precise code documentation engine. Add idiomatic documentation comments to the given code. Preserve exact behavior and formatting as much as possible. Return only the updated code with comments, no prose or backticks."
+        user := "Add documentation comments to this code:\n" + payload.Selection
+        ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+        defer cancel()
+        messages := []llm.Message{{Role: "system", Content: sys}, {Role: "user", Content: user}}
+        opts := s.llmRequestOpts()
+        if text, err := s.llmClient.Chat(ctx, messages, opts...); err == nil {
+            if out := stripCodeFences(strings.TrimSpace(text)); out != "" {
+                edit := WorkspaceEdit{Changes: map[string][]TextEdit{payload.URI: {{Range: payload.Range, NewText: out}}}}
+                ca.Edit = &edit
+                return ca, true
+            }
+        } else {
+            logging.Logf("lsp ", "codeAction document llm error: %v", err)
+        }
     case "go_test":
         if edit, jumpURI, jumpRange, ok := s.resolveGoTest(payload.URI, payload.Range.Start); ok {
             ca.Edit = &edit
@@ -248,6 +267,25 @@ func (s *Server) buildGoUnitTestCodeAction(p CodeActionParams) *CodeAction {
     return &ca
 }
 
+// buildDocumentCodeAction offers to document the selected code by injecting comments.
+func (s *Server) buildDocumentCodeAction(p CodeActionParams, sel string) *CodeAction {
+    if s.llmClient == nil {
+        return nil
+    }
+    if strings.TrimSpace(sel) == "" {
+        return nil
+    }
+    payload := struct {
+        Type      string `json:"type"`
+        URI       string `json:"uri"`
+        Range     Range  `json:"range"`
+        Selection string `json:"selection"`
+    }{Type: "document", URI: p.TextDocument.URI, Range: p.Range, Selection: sel}
+    raw, _ := json.Marshal(payload)
+    ca := CodeAction{Title: "Hexai: document code", Kind: "refactor.rewrite", Data: raw}
+    return &ca
+}
+
 func (s *Server) resolveGoTest(uri string, pos Position) (WorkspaceEdit, string, Range, bool) {
     path := strings.TrimPrefix(uri, "file://")
     if !strings.HasSuffix(path, ".go") || strings.HasSuffix(path, "_test.go") {