From 7abb7c9177d34f3b2a1773624f0da7daa8c8e2de Mon Sep 17 00:00:00 2001 From: Paul Buetow Date: Tue, 19 Aug 2025 23:00:05 +0300 Subject: lsp/codeactions: make actions lazy and resolve on selection\n\n- Advertise CodeAction resolveProvider and implement codeAction/resolve\n- Return lightweight actions with data; no LLM call during listing\n- On resolve, perform LLM and populate WorkspaceEdit\n- Update tests to cover lazy+resolve flow --- internal/lsp/codeaction_test.go | 18 ++-- internal/lsp/handlers.go | 184 ++++++++++++++++++++++++++-------------- internal/lsp/types.go | 19 +++-- 3 files changed, 147 insertions(+), 74 deletions(-) diff --git a/internal/lsp/codeaction_test.go b/internal/lsp/codeaction_test.go index e9abbb8..59b16d8 100644 --- a/internal/lsp/codeaction_test.go +++ b/internal/lsp/codeaction_test.go @@ -15,15 +15,20 @@ func (f fakeLLM) Chat(_ context.Context, _ []llm.Message, _ ...llm.RequestOption func (f fakeLLM) Name() string { return "fake" } func (f fakeLLM) DefaultModel() string { return "fake-model" } -func TestBuildRewriteCodeAction_ReturnsEdit(t *testing.T) { +func TestBuildRewriteCodeAction_LazyAndResolves(t *testing.T) { s := newTestServer() s.llmClient = fakeLLM{resp: "REWRITTEN"} p := CodeActionParams{TextDocument: TextDocumentIdentifier{URI: "file:///t.go"}, Range: Range{Start: Position{Line: 1, Character: 2}, End: Position{Line: 3, Character: 4}}} sel := ";rewrite;\nold code" ca := s.buildRewriteCodeAction(p, sel) if ca == nil { t.Fatalf("expected code action") } - if ca.Edit == nil || len(ca.Edit.Changes) == 0 { t.Fatalf("expected workspace edit with changes") } - edits := ca.Edit.Changes[p.TextDocument.URI] + // Should be lazy (no edit yet) + if ca.Edit != nil { t.Fatalf("expected nil Edit before resolve") } + if len(ca.Data) == 0 { t.Fatalf("expected data payload for lazy resolve") } + // Resolve now + resolved, ok := s.resolveCodeAction(*ca) + if !ok || resolved.Edit == nil { t.Fatalf("expected resolve to produce edit") } + edits := resolved.Edit.Changes[p.TextDocument.URI] if len(edits) != 1 { t.Fatalf("expected 1 edit, got %d", len(edits)) } if edits[0].Range != p.Range { t.Fatalf("edit range mismatch: got %+v want %+v", edits[0].Range, p.Range) } if edits[0].NewText == "" { t.Fatalf("expected non-empty replacement text") } @@ -37,7 +42,7 @@ func TestBuildRewriteCodeAction_NoInstruction(t *testing.T) { if ca := s.buildRewriteCodeAction(p, sel); ca != nil { t.Fatalf("expected nil action when no instruction present") } } -func TestBuildDiagnosticsCodeAction_ReturnsEdit(t *testing.T) { +func TestBuildDiagnosticsCodeAction_LazyAndResolves(t *testing.T) { s := newTestServer() s.llmClient = fakeLLM{resp: "FIXED"} p := CodeActionParams{TextDocument: TextDocumentIdentifier{URI: "file:///t.go"}, Range: Range{Start: Position{Line: 10}, End: Position{Line: 12, Character: 5}}} @@ -50,7 +55,10 @@ func TestBuildDiagnosticsCodeAction_ReturnsEdit(t *testing.T) { sel := "some selected code" ca := s.buildDiagnosticsCodeAction(p, sel) if ca == nil { t.Fatalf("expected diagnostics code action") } - if ca.Edit == nil || len(ca.Edit.Changes) == 0 { t.Fatalf("expected workspace edit") } + if ca.Edit != nil { t.Fatalf("expected lazy action without edit") } + if len(ca.Data) == 0 { t.Fatalf("expected data payload for lazy diagnostics action") } + resolved, ok := s.resolveCodeAction(*ca) + if !ok || resolved.Edit == nil { t.Fatalf("expected resolve to produce edit") } } func TestBuildDiagnosticsCodeAction_NoDiagnostics(t *testing.T) { diff --git a/internal/lsp/handlers.go b/internal/lsp/handlers.go index 45eaec0..ee7c33a 100644 --- a/internal/lsp/handlers.go +++ b/internal/lsp/handlers.go @@ -15,7 +15,7 @@ import ( ) func (s *Server) handle(req Request) { - switch req.Method { + switch req.Method { case "initialize": s.handleInitialize(req) case "initialized": @@ -32,9 +32,11 @@ func (s *Server) handle(req Request) { s.handleDidClose(req) case "textDocument/completion": s.handleCompletion(req) - case "textDocument/codeAction": - s.handleCodeAction(req) - default: + case "textDocument/codeAction": + s.handleCodeAction(req) + case "codeAction/resolve": + s.handleCodeActionResolve(req) + default: if len(req.ID) != 0 { s.reply(req.ID, nil, &RespError{Code: -32601, Message: fmt.Sprintf("method not found: %s", req.Method)}) } @@ -46,17 +48,17 @@ func (s *Server) handleInitialize(req Request) { if s.llmClient != nil { version = version + " [" + s.llmClient.Name() + ":" + s.llmClient.DefaultModel() + "]" } - res := InitializeResult{ - Capabilities: ServerCapabilities{ - TextDocumentSync: 1, // 1 = TextDocumentSyncKindFull - CompletionProvider: &CompletionOptions{ - ResolveProvider: false, - TriggerCharacters: s.triggerChars, - }, - CodeActionProvider: true, - }, - ServerInfo: &ServerInfo{Name: "hexai", Version: version}, - } + res := InitializeResult{ + Capabilities: ServerCapabilities{ + TextDocumentSync: 1, // 1 = TextDocumentSyncKindFull + CompletionProvider: &CompletionOptions{ + ResolveProvider: false, + TriggerCharacters: s.triggerChars, + }, + CodeActionProvider: CodeActionOptions{ResolveProvider: true}, + }, + ServerInfo: &ServerInfo{Name: "hexai", Version: version}, + } s.reply(req.ID, res, nil) } @@ -96,57 +98,113 @@ func (s *Server) handleCodeAction(req Request) { } func (s *Server) buildRewriteCodeAction(p CodeActionParams, sel string) *CodeAction { - 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}} - 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{p.TextDocument.URI: {{Range: p.Range, NewText: out}}}} - ca := CodeAction{Title: "Hexai: rewrite selection", Kind: "refactor.rewrite", Edit: &edit} - return &ca - } - } else { - logging.Logf("lsp ", "codeAction rewrite llm error: %v", err) - } - } - return nil + if instr, cleaned := instructionFromSelection(sel); strings.TrimSpace(instr) != "" { + payload := struct{ + Type string `json:"type"` + URI string `json:"uri"` + Range Range `json:"range"` + Instruction string `json:"instruction"` + Selection string `json:"selection"` + }{Type: "rewrite", URI: p.TextDocument.URI, Range: p.Range, Instruction: instr, Selection: cleaned} + raw, _ := json.Marshal(payload) + ca := CodeAction{Title: "Hexai: rewrite selection", Kind: "refactor.rewrite", Data: raw} + return &ca + } + return nil } func (s *Server) buildDiagnosticsCodeAction(p CodeActionParams, sel string) *CodeAction { - diags := s.diagnosticsInRange(p.Context, p.Range) - if len(diags) == 0 { - return nil - } - 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 { - 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()}} - 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{p.TextDocument.URI: {{Range: p.Range, NewText: out}}}} - ca := CodeAction{Title: "Hexai: resolve diagnostics", Kind: "quickfix", Edit: &edit} - return &ca - } - } else { - logging.Logf("lsp ", "codeAction diagnostics llm error: %v", err) - } - return nil + diags := s.diagnosticsInRange(p.Context, p.Range) + if len(diags) == 0 { + return nil + } + payload := struct{ + Type string `json:"type"` + URI string `json:"uri"` + Range Range `json:"range"` + Selection string `json:"selection"` + Diagnostics []Diagnostic `json:"diagnostics"` + }{Type: "diagnostics", URI: p.TextDocument.URI, Range: p.Range, Selection: sel, Diagnostics: diags} + raw, _ := json.Marshal(payload) + ca := CodeAction{Title: "Hexai: resolve diagnostics", Kind: "quickfix", Data: raw} + return &ca +} + +func (s *Server) resolveCodeAction(ca CodeAction) (CodeAction, bool) { + if s.llmClient == nil || len(ca.Data) == 0 { + return ca, false + } + var payload struct{ + Type string `json:"type"` + URI string `json:"uri"` + Range Range `json:"range"` + Instruction string `json:"instruction,omitempty"` + Selection string `json:"selection"` + Diagnostics []Diagnostic `json:"diagnostics,omitempty"` + } + if err := json.Unmarshal(ca.Data, &payload); err != nil { + return ca, false + } + switch payload.Type { + case "rewrite": + 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", payload.Instruction, 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 rewrite llm error: %v", err) + } + case "diagnostics": + 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 payload.Diagnostics { + 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(payload.Selection) + ctx, cancel := context.WithTimeout(context.Background(), 12*time.Second) + defer cancel() + messages := []llm.Message{{Role: "system", Content: sys}, {Role: "user", Content: b.String()}} + 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 diagnostics llm error: %v", err) + } + } + return ca, false +} + +func (s *Server) handleCodeActionResolve(req Request) { + var ca CodeAction + if err := json.Unmarshal(req.Params, &ca); err != nil { + if len(req.ID) != 0 { + s.reply(req.ID, ca, nil) + } + return + } + if resolved, ok := s.resolveCodeAction(ca); ok { + s.reply(req.ID, resolved, nil) + return + } + s.reply(req.ID, ca, nil) } func (s *Server) llmRequestOpts() []llm.RequestOption { diff --git a/internal/lsp/types.go b/internal/lsp/types.go index ce98ac0..868a1a2 100644 --- a/internal/lsp/types.go +++ b/internal/lsp/types.go @@ -35,9 +35,10 @@ type ServerInfo struct { } type ServerCapabilities struct { - TextDocumentSync any `json:"textDocumentSync,omitempty"` - CompletionProvider *CompletionOptions `json:"completionProvider,omitempty"` - CodeActionProvider bool `json:"codeActionProvider,omitempty"` + TextDocumentSync any `json:"textDocumentSync,omitempty"` + CompletionProvider *CompletionOptions `json:"completionProvider,omitempty"` + // bool | CodeActionOptions + CodeActionProvider any `json:"codeActionProvider,omitempty"` } type CompletionOptions struct { @@ -63,6 +64,11 @@ type CompletionItem struct { Documentation string `json:"documentation,omitempty"` } +// Code action options +type CodeActionOptions struct { + ResolveProvider bool `json:"resolveProvider,omitempty"` +} + // LSP param types (subset) type TextDocumentItem struct { URI string `json:"uri"` @@ -122,9 +128,10 @@ type WorkspaceEdit struct { } type CodeAction struct { - Title string `json:"title"` - Kind string `json:"kind,omitempty"` - Edit *WorkspaceEdit `json:"edit,omitempty"` + Title string `json:"title"` + Kind string `json:"kind,omitempty"` + Edit *WorkspaceEdit `json:"edit,omitempty"` + Data json.RawMessage `json:"data,omitempty"` } // Diagnostics (subset needed for code action context) -- cgit v1.2.3