feat: add 'Simplify and improve' action; configurable prompts in config; wire into LSP and TUI

8 files changed, 147 insertions, 74 deletions

diff --git a/internal/appconfig/config.go b/internal/appconfig/config.go
index bd511e2..87b5a29 100644
--- a/internal/appconfig/config.go
+++ b/internal/appconfig/config.go
@@ -74,12 +74,14 @@ type App struct {
 	// Code actions
 	PromptCodeActionRewriteSystem     string `json:"-" toml:"-"`
 	PromptCodeActionDiagnosticsSystem string `json:"-" toml:"-"`
-	PromptCodeActionDocumentSystem    string `json:"-" toml:"-"`
-	PromptCodeActionRewriteUser       string `json:"-" toml:"-"`
-	PromptCodeActionDiagnosticsUser   string `json:"-" toml:"-"`
-	PromptCodeActionDocumentUser      string `json:"-" toml:"-"`
-	PromptCodeActionGoTestSystem      string `json:"-" toml:"-"`
-	PromptCodeActionGoTestUser        string `json:"-" toml:"-"`
+    PromptCodeActionDocumentSystem    string `json:"-" toml:"-"`
+    PromptCodeActionRewriteUser       string `json:"-" toml:"-"`
+    PromptCodeActionDiagnosticsUser   string `json:"-" toml:"-"`
+    PromptCodeActionDocumentUser      string `json:"-" toml:"-"`
+    PromptCodeActionGoTestSystem      string `json:"-" toml:"-"`
+    PromptCodeActionGoTestUser        string `json:"-" toml:"-"`
+    PromptCodeActionSimplifySystem    string `json:"-" toml:"-"`
+    PromptCodeActionSimplifyUser      string `json:"-" toml:"-"`
 	// CLI
 	PromptCLIDefaultSystem string `json:"-" toml:"-"`
 	PromptCLIExplainSystem string `json:"-" toml:"-"`
@@ -127,8 +129,10 @@ func newDefaultConfig() App {
 		PromptCodeActionRewriteUser:       "Instruction: {{instruction}}\n\nSelected code to transform:\n{{selection}}",
 		PromptCodeActionDiagnosticsUser:   "Diagnostics to resolve (selection only):\n{{diagnostics}}\n\nSelected code:\n{{selection}}",
 		PromptCodeActionDocumentUser:      "Add documentation comments to this code:\n{{selection}}",
-		PromptCodeActionGoTestSystem:      "You are a precise Go unit test generator. Given a Go function, write one or more Test* functions using the testing package. Do NOT include package or imports, only the test function(s). Prefer table-driven tests. Keep it minimal and idiomatic.",
-		PromptCodeActionGoTestUser:        "Function under test:\n{{function}}",
+        PromptCodeActionGoTestSystem:      "You are a precise Go unit test generator. Given a Go function, write one or more Test* functions using the testing package. Do NOT include package or imports, only the test function(s). Prefer table-driven tests. Keep it minimal and idiomatic.",
+        PromptCodeActionGoTestUser:        "Function under test:\n{{function}}",
+        PromptCodeActionSimplifySystem:    "You are a precise code improvement engine. Simplify and improve the given code while preserving behavior. Return only the improved code with no prose or backticks.",
+        PromptCodeActionSimplifyUser:      "Improve this code:\n{{selection}}",
 
 		PromptCLIDefaultSystem: "You are Hexai CLI. Default to very short, concise answers. If the user asks for commands, output only the commands (one per line) with no commentary or explanation. Only when the word 'explain' appears in the prompt, produce a verbose explanation.",
 		PromptCLIExplainSystem: "You are Hexai CLI. The user requested an explanation. Provide a clear, verbose explanation with reasoning and details. If commands are needed, include them with brief context.",
@@ -256,14 +260,16 @@ type sectionPromptsChat struct {
 }
 
 type sectionPromptsCodeAction struct {
-	RewriteSystem     string `toml:"rewrite_system"`
-	DiagnosticsSystem string `toml:"diagnostics_system"`
-	DocumentSystem    string `toml:"document_system"`
-	RewriteUser       string `toml:"rewrite_user"`
-	DiagnosticsUser   string `toml:"diagnostics_user"`
-	DocumentUser      string `toml:"document_user"`
-	GoTestSystem      string `toml:"go_test_system"`
-	GoTestUser        string `toml:"go_test_user"`
+    RewriteSystem     string `toml:"rewrite_system"`
+    DiagnosticsSystem string `toml:"diagnostics_system"`
+    DocumentSystem    string `toml:"document_system"`
+    RewriteUser       string `toml:"rewrite_user"`
+    DiagnosticsUser   string `toml:"diagnostics_user"`
+    DocumentUser      string `toml:"document_user"`
+    GoTestSystem      string `toml:"go_test_system"`
+    GoTestUser        string `toml:"go_test_user"`
+    SimplifySystem    string `toml:"simplify_system"`
+    SimplifyUser      string `toml:"simplify_user"`
 }
 
 type sectionPromptsCLI struct {
@@ -387,7 +393,7 @@ func (fc *fileConfig) toApp() App {
 		out.PromptChatSystem = fc.Prompts.Chat.System
 	}
 	// code action
-	if (fc.Prompts.CodeAction != sectionPromptsCodeAction{}) {
+    if (fc.Prompts.CodeAction != sectionPromptsCodeAction{}) {
 		if strings.TrimSpace(fc.Prompts.CodeAction.RewriteSystem) != "" {
 			out.PromptCodeActionRewriteSystem = fc.Prompts.CodeAction.RewriteSystem
 		}
@@ -409,10 +415,16 @@ func (fc *fileConfig) toApp() App {
 		if strings.TrimSpace(fc.Prompts.CodeAction.GoTestSystem) != "" {
 			out.PromptCodeActionGoTestSystem = fc.Prompts.CodeAction.GoTestSystem
 		}
-		if strings.TrimSpace(fc.Prompts.CodeAction.GoTestUser) != "" {
-			out.PromptCodeActionGoTestUser = fc.Prompts.CodeAction.GoTestUser
-		}
-	}
+        if strings.TrimSpace(fc.Prompts.CodeAction.GoTestUser) != "" {
+            out.PromptCodeActionGoTestUser = fc.Prompts.CodeAction.GoTestUser
+        }
+        if strings.TrimSpace(fc.Prompts.CodeAction.SimplifySystem) != "" {
+            out.PromptCodeActionSimplifySystem = fc.Prompts.CodeAction.SimplifySystem
+        }
+        if strings.TrimSpace(fc.Prompts.CodeAction.SimplifyUser) != "" {
+            out.PromptCodeActionSimplifyUser = fc.Prompts.CodeAction.SimplifyUser
+        }
+    }
 	// cli
 	if (fc.Prompts.CLI != sectionPromptsCLI{}) {
 		if strings.TrimSpace(fc.Prompts.CLI.DefaultSystem) != "" {
@@ -611,9 +623,15 @@ func (a *App) mergePrompts(other *App) {
 	if strings.TrimSpace(other.PromptCodeActionGoTestSystem) != "" {
 		a.PromptCodeActionGoTestSystem = other.PromptCodeActionGoTestSystem
 	}
-	if strings.TrimSpace(other.PromptCodeActionGoTestUser) != "" {
-		a.PromptCodeActionGoTestUser = other.PromptCodeActionGoTestUser
-	}
+    if strings.TrimSpace(other.PromptCodeActionGoTestUser) != "" {
+        a.PromptCodeActionGoTestUser = other.PromptCodeActionGoTestUser
+    }
+    if strings.TrimSpace(other.PromptCodeActionSimplifySystem) != "" {
+        a.PromptCodeActionSimplifySystem = other.PromptCodeActionSimplifySystem
+    }
+    if strings.TrimSpace(other.PromptCodeActionSimplifyUser) != "" {
+        a.PromptCodeActionSimplifyUser = other.PromptCodeActionSimplifyUser
+    }
 	// CLI
 	if strings.TrimSpace(other.PromptCLIDefaultSystem) != "" {
 		a.PromptCLIDefaultSystem = other.PromptCLIDefaultSystem
diff --git a/internal/hexaiaction/prompts.go b/internal/hexaiaction/prompts.go
index 2e0e4e2..3c33f8a 100644
--- a/internal/hexaiaction/prompts.go
+++ b/internal/hexaiaction/prompts.go
@@ -43,8 +43,14 @@ func runDiagnostics(ctx context.Context, cfg appconfig.App, client chatDoer, dia
 }
 
 func runDocument(ctx context.Context, cfg appconfig.App, client chatDoer, selection string) (string, error) {
-	sys := cfg.PromptCodeActionDocumentSystem
-	user := Render(cfg.PromptCodeActionDocumentUser, map[string]string{"selection": selection})
+    sys := cfg.PromptCodeActionDocumentSystem
+    user := Render(cfg.PromptCodeActionDocumentUser, map[string]string{"selection": selection})
+    return runOnceWithOpts(ctx, client, sys, user, reqOptsFrom(cfg))
+}
+
+func runSimplify(ctx context.Context, cfg appconfig.App, client chatDoer, selection string) (string, error) {
+    sys := cfg.PromptCodeActionSimplifySystem
+    user := Render(cfg.PromptCodeActionSimplifyUser, map[string]string{"selection": selection})
     return runOnceWithOpts(ctx, client, sys, user, reqOptsFrom(cfg))
 }
 
diff --git a/internal/hexaiaction/run.go b/internal/hexaiaction/run.go
index a8d4243..5417d3f 100644
--- a/internal/hexaiaction/run.go
+++ b/internal/hexaiaction/run.go
@@ -46,10 +46,10 @@ func Run(ctx context.Context, stdin io.Reader, stdout, stderr io.Writer) error {
 }
 
 func executeAction(ctx context.Context, kind ActionKind, parts InputParts, cfg appconfig.App, client chatDoer, stderr io.Writer) (string, error) {
-	switch kind {
-	case ActionSkip:
-		return parts.Selection, nil
-	case ActionRewrite:
+    switch kind {
+    case ActionSkip:
+        return parts.Selection, nil
+    case ActionRewrite:
 		instr, cleaned := ExtractInstruction(parts.Selection)
 		if strings.TrimSpace(instr) == "" {
             fmt.Fprintln(stderr, logging.AnsiBase+"hexai-tmux-action: no inline instruction found; echoing input"+logging.AnsiReset)
@@ -62,17 +62,21 @@ func executeAction(ctx context.Context, kind ActionKind, parts InputParts, cfg a
 		cctx, cancel := timeout10s(ctx)
 		defer cancel()
 		return runDiagnostics(cctx, cfg, client, parts.Diagnostics, parts.Selection)
-	case ActionDocument:
-		cctx, cancel := timeout10s(ctx)
-		defer cancel()
-		return runDocument(cctx, cfg, client, parts.Selection)
-	case ActionGoTest:
-		cctx, cancel := timeout8s(ctx)
-		defer cancel()
-		return runGoTest(cctx, cfg, client, parts.Selection)
-	default:
-		return parts.Selection, nil
-	}
+    case ActionDocument:
+        cctx, cancel := timeout10s(ctx)
+        defer cancel()
+        return runDocument(cctx, cfg, client, parts.Selection)
+    case ActionGoTest:
+        cctx, cancel := timeout8s(ctx)
+        defer cancel()
+        return runGoTest(cctx, cfg, client, parts.Selection)
+    case ActionSimplify:
+        cctx, cancel := timeout10s(ctx)
+        defer cancel()
+        return runSimplify(cctx, cfg, client, parts.Selection)
+    default:
+        return parts.Selection, nil
+    }
 }
 
 // client construction is shared via internal/llmutils
diff --git a/internal/hexaiaction/tui.go b/internal/hexaiaction/tui.go
index 7275480..80b1fee 100644
--- a/internal/hexaiaction/tui.go
+++ b/internal/hexaiaction/tui.go
@@ -28,6 +28,7 @@ type model struct {
 func newModel() model {
     items := []list.Item{
         item{title: "Rewrite selection", desc: "", kind: ActionRewrite, hotkey: 'r'},
+        item{title: "Simplify and improve", desc: "", kind: ActionSimplify, hotkey: 'i'},
         item{title: "Document code", desc: "", kind: ActionDocument, hotkey: 'c'},
         item{title: "Generate Go unit test(s)", desc: "", kind: ActionGoTest, hotkey: 't'},
         item{title: "Skip", desc: "", kind: ActionSkip, hotkey: 's'},
@@ -77,7 +78,7 @@ func handleKey(m model, msg tea.KeyMsg) (tea.Model, tea.Cmd) {
         m.list.Select(0)
     case "end":
         if n := len(m.list.Items()); n > 0 { m.list.Select(n - 1) }
-    case "s", "r", "c", "t":
+    case "s", "r", "c", "t", "i":
         items := m.list.Items()
         for i := 0; i < len(items); i++ {
             if it, ok := items[i].(item); ok && strings.ToLower(string(it.hotkey)) == low {
diff --git a/internal/hexaiaction/types.go b/internal/hexaiaction/types.go
index 2dc918f..708c433 100644
--- a/internal/hexaiaction/types.go
+++ b/internal/hexaiaction/types.go
@@ -5,11 +5,12 @@ package hexaiaction
 type ActionKind string
 
 const (
-	ActionSkip        ActionKind = "skip"
-	ActionRewrite     ActionKind = "rewrite"
-	ActionDiagnostics ActionKind = "diagnostics"
-	ActionDocument    ActionKind = "document"
-	ActionGoTest      ActionKind = "gotest"
+    ActionSkip        ActionKind = "skip"
+    ActionRewrite     ActionKind = "rewrite"
+    ActionDiagnostics ActionKind = "diagnostics"
+    ActionDocument    ActionKind = "document"
+    ActionGoTest      ActionKind = "gotest"
+    ActionSimplify    ActionKind = "simplify"
 )
 
 // InputParts represents parsed stdin input for actions.
diff --git a/internal/hexailsp/run.go b/internal/hexailsp/run.go
index 72cf902..e3dfb28 100644
--- a/internal/hexailsp/run.go
+++ b/internal/hexailsp/run.go
@@ -140,5 +140,7 @@ func makeServerOptions(cfg appconfig.App, logContext bool, client llm.Client) ls
 		PromptDocumentUser:      cfg.PromptCodeActionDocumentUser,
 		PromptGoTestSystem:      cfg.PromptCodeActionGoTestSystem,
 		PromptGoTestUser:        cfg.PromptCodeActionGoTestUser,
+		PromptSimplifySystem:    cfg.PromptCodeActionSimplifySystem,
+		PromptSimplifyUser:      cfg.PromptCodeActionSimplifyUser,
 	}
 }
diff --git a/internal/lsp/handlers_codeaction.go b/internal/lsp/handlers_codeaction.go
index 17e92bc..d8dba38 100644
--- a/internal/lsp/handlers_codeaction.go
+++ b/internal/lsp/handlers_codeaction.go
@@ -31,24 +31,42 @@ func (s *Server) handleCodeAction(req Request) {
 	}
 	sel := extractRangeText(d, p.Range)
 
-	actions := make([]CodeAction, 0, 4)
+    actions := make([]CodeAction, 0, 5)
 	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)
-	}
+    if a := s.buildDocumentCodeAction(p, sel); a != nil {
+        actions = append(actions, *a)
+    }
+    if a := s.buildGoUnitTestCodeAction(p); a != nil {
+        actions = append(actions, *a)
+    }
+    if a := s.buildSimplifyCodeAction(p, sel); a != nil {
+        actions = append(actions, *a)
+    }
 	if len(req.ID) != 0 {
 		s.reply(req.ID, actions, nil)
 	}
 }
 
+func (s *Server) buildSimplifyCodeAction(p CodeActionParams, sel string) *CodeAction {
+    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: "simplify", URI: p.TextDocument.URI, Range: p.Range, Selection: sel}
+    raw, _ := json.Marshal(payload)
+    ca := CodeAction{Title: "Hexai: simplify and improve", Kind: "refactor", Data: raw}
+    return &ca
+}
+
 func (s *Server) buildRewriteCodeAction(p CodeActionParams, sel string) *CodeAction {
 	if instr, cleaned := instructionFromSelection(sel); strings.TrimSpace(instr) != "" {
 		payload := struct {
@@ -97,7 +115,7 @@ func (s *Server) resolveCodeAction(ca CodeAction) (CodeAction, bool) {
 	if err := json.Unmarshal(ca.Data, &payload); err != nil {
 		return ca, false
 	}
-	switch payload.Type {
+    switch payload.Type {
 	case "rewrite":
 		sys := s.promptRewriteSystem
 		user := renderTemplate(s.promptRewriteUser, map[string]string{"instruction": payload.Instruction, "selection": payload.Selection})
@@ -155,17 +173,34 @@ func (s *Server) resolveCodeAction(ca CodeAction) (CodeAction, bool) {
 		} 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
-			// After edit is applied, ask client to jump to new test function
-			ca.Command = &Command{Title: "Jump to generated test", Command: "hexai.showDocument", Arguments: []any{jumpURI, jumpRange}}
-			// Also send a server-initiated showDocument shortly after resolve to cover
-			// clients that do not execute commands from code actions.
-			s.deferShowDocument(jumpURI, jumpRange)
-			return ca, true
-		}
-	}
+    case "go_test":
+        if edit, jumpURI, jumpRange, ok := s.resolveGoTest(payload.URI, payload.Range.Start); ok {
+            ca.Edit = &edit
+            // After edit is applied, ask client to jump to new test function
+            ca.Command = &Command{Title: "Jump to generated test", Command: "hexai.showDocument", Arguments: []any{jumpURI, jumpRange}}
+            // Also send a server-initiated showDocument shortly after resolve to cover
+            // clients that do not execute commands from code actions.
+            s.deferShowDocument(jumpURI, jumpRange)
+            return ca, true
+        }
+    case "simplify":
+        sys := s.promptRewriteSystem
+        // Reuse rewrite user template with a fixed instruction
+        user := renderTemplate(s.promptRewriteUser, map[string]string{"instruction": "Simplify and improve the code while preserving behavior. Return only the improved code.", "selection": 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 simplify llm error: %v", err)
+        }
+    }
 	return ca, false
 }
 
diff --git a/internal/lsp/server.go b/internal/lsp/server.go
index 1c3e676..796d6f4 100644
--- a/internal/lsp/server.go
+++ b/internal/lsp/server.go
@@ -78,9 +78,11 @@ type Server struct {
 	promptDocumentSystem    string
 	promptRewriteUser       string
 	promptDiagnosticsUser   string
-	promptDocumentUser      string
-	promptGoTestSystem      string
-	promptGoTestUser        string
+    promptDocumentUser      string
+    promptGoTestSystem      string
+    promptGoTestUser        string
+    promptSimplifySystem    string
+    promptSimplifyUser      string
 }
 
 // ServerOptions collects configuration for NewServer to avoid long parameter lists.
@@ -119,8 +121,10 @@ type ServerOptions struct {
 	PromptRewriteUser       string
 	PromptDiagnosticsUser   string
 	PromptDocumentUser      string
-	PromptGoTestSystem      string
-	PromptGoTestUser        string
+    PromptGoTestSystem      string
+    PromptGoTestUser        string
+    PromptSimplifySystem    string
+    PromptSimplifyUser      string
 }
 
 func NewServer(r io.Reader, w io.Writer, logger *log.Logger, opts ServerOptions) *Server {
@@ -199,9 +203,11 @@ func NewServer(r io.Reader, w io.Writer, logger *log.Logger, opts ServerOptions)
 	s.promptDocumentSystem = opts.PromptDocumentSystem
 	s.promptRewriteUser = opts.PromptRewriteUser
 	s.promptDiagnosticsUser = opts.PromptDiagnosticsUser
-	s.promptDocumentUser = opts.PromptDocumentUser
-	s.promptGoTestSystem = opts.PromptGoTestSystem
-	s.promptGoTestUser = opts.PromptGoTestUser
+    s.promptDocumentUser = opts.PromptDocumentUser
+    s.promptGoTestSystem = opts.PromptGoTestSystem
+    s.promptGoTestUser = opts.PromptGoTestUser
+    s.promptSimplifySystem = opts.PromptSimplifySystem
+    s.promptSimplifyUser = opts.PromptSimplifyUser
 
 	// Assign package-level inline trigger chars for free helper functions
 	if s.inlineOpen != "" {