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

1 files changed, 39 insertions, 1 deletions

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") {