From aa7204c968f8cd4f2df36369dae4fe77b487f39e Mon Sep 17 00:00:00 2001
From: Paul Buetow <paul@buetow.org>
Date: Wed, 3 Sep 2025 16:47:54 +0300
Subject: lsp: add 'Hexai: document code' action to add doc comments to
 selected code

---
 internal/lsp/handlers_codeaction.go | 40 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 39 insertions(+), 1 deletion(-)

(limited to 'internal/lsp')

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") {
-- 
cgit v1.2.3