From bff9994dc98b7a0a3fde66568e533591ebfefcff Mon Sep 17 00:00:00 2001 From: Cheikh Seck Date: Sat, 14 Mar 2026 16:10:23 +0000 Subject: [PATCH] Fix prompt handling and background listings --- cmd/godex/main.go | 63 ++------ cmd/godex/prompt.go | 284 +++++++++++++++++++++++++++++++++++ go.mod | 20 ++- go.sum | 39 +++++ internal/mcp/bash.go | 42 ++++-- internal/providers/gemini.go | 24 +++ internal/providers/ollama.go | 29 +++- 7 files changed, 426 insertions(+), 75 deletions(-) create mode 100644 cmd/godex/prompt.go diff --git a/cmd/godex/main.go b/cmd/godex/main.go index 32764e9..b2316a3 100644 --- a/cmd/godex/main.go +++ b/cmd/godex/main.go @@ -21,8 +21,6 @@ import ( godexcontext "github.com/cheikh-seck/godex/internal/context" "github.com/cheikh-seck/godex/internal/mcp" "github.com/cheikh-seck/godex/internal/wizard" - - "github.com/peterh/liner" ) const ( @@ -39,6 +37,8 @@ type MCPServer interface { Close() error } +var slashCommands = []string{"/add-path ", "/paths", "/tools", "/exit", "/quit", "/save", "/save-exit", "/killbg", "/bg", "/clear", "/help"} + type sessionEntry struct { Prompt string Response string @@ -117,7 +117,7 @@ func main() { fmt.Printf("MCP Servers: %d\n", len(servers)) fmt.Println("Commands: /paths, /add-path , /exit, Up/Down for history, Ctrl+C to cancel") fmt.Println("Type your prompt or /help for more options.") - fmt.Println("Multiline: Enter to add new line, Enter again on empty line to submit") + fmt.Println("Multiline: paste with newlines to enter multiline, Enter on empty line to submit") // Get working directory for session files wd, _ := os.Getwd() @@ -134,11 +134,9 @@ func main() { fmt.Println("[Agents] Loaded AGENTS.md") } - rl := NewLiner() - defer rl.Close() - // Track session for summary var sessionEntries []sessionEntry + var history []string for { // Show prompt with background process count @@ -148,46 +146,22 @@ func main() { prompt = fmt.Sprintf("[%d bg] >", bgCount) } - input, err := rl.Prompt(prompt) + input, err := readPrompt(prompt, history) if err != nil { - if err == liner.ErrPromptAborted { - fmt.Println("\n[Cancelled]") + if err == ErrPromptAborted { + fmt.Println("\n[Cancelled] Use /quit to exit or /save to save and exit.") agent.CancelPrompt(provider) continue } break } - // Handle multiline paste: check if input already contains newlines (paste captured full text) - // If not a command (doesn't start with /), wait for more input until user submits empty line - accumulated := input - - // If input already has newlines, it's a multiline paste - use as-is - if !strings.Contains(input, "\n") && !strings.HasPrefix(input, "/") { - // Keep prompting for more input until user submits empty line - for { - moreInput, err := rl.Prompt("... ") - if err != nil { - if err == liner.ErrPromptAborted { - accumulated = "" // Discard on cancellation - } - break - } - if moreInput == "" { - break - } - accumulated = accumulated + "\n" + moreInput - } - } - - input = accumulated - input = strings.TrimSpace(input) if input == "" { continue } - rl.AppendHistory(input) + history = appendHistory(history, input) if input == "/exit" || input == "/quit" { break @@ -311,6 +285,7 @@ User request: %s`, toolsDesc, sessionContext, wd, wd, input) go func() { select { case <-stopSignal: + agent.CancelPrompt(provider) roundCancel() case <-roundCtx.Done(): } @@ -850,26 +825,6 @@ func runSinglePrompt(ctx context.Context, provider *config.Provider, prompt stri return nil } -var slashCommands = []string{"/add-path ", "/paths", "/tools", "/exit", "/quit", "/save", "/save-exit", "/killbg", "/bg", "/clear", "/help"} - -func NewLiner() *liner.State { - l := liner.NewLiner() - l.SetCtrlCAborts(true) - l.SetCompleter(func(line string) []string { - if !strings.HasPrefix(line, "/") { - return nil - } - var matches []string - for _, cmd := range slashCommands { - if strings.HasPrefix(cmd, line) { - matches = append(matches, cmd) - } - } - return matches - }) - return l -} - func getToolsDescription(servers []MCPServer) string { if len(servers) == 0 { return "No tools available." diff --git a/cmd/godex/prompt.go b/cmd/godex/prompt.go new file mode 100644 index 0000000..a267840 --- /dev/null +++ b/cmd/godex/prompt.go @@ -0,0 +1,284 @@ +package main + +import ( + "errors" + "strings" + "unicode/utf8" + + "github.com/charmbracelet/bubbles/textinput" + tea "github.com/charmbracelet/bubbletea" +) + +var ErrPromptAborted = errors.New("prompt aborted") + +type promptModel struct { + basePrompt string + contPrompt string + input textinput.Model + lines []string + multiline bool + history []string + historyIndex int + historyDraft string + submitted bool + aborted bool + completions []string + completeSeed string + completeIdx int +} + +func newPromptModel(prompt string, history []string) promptModel { + ti := textinput.New() + ti.Prompt = prompt + ti.Focus() + ti.CharLimit = 0 + + return promptModel{ + basePrompt: prompt, + contPrompt: "... ", + input: ti, + history: history, + historyIndex: len(history), + } +} + +func readPrompt(prompt string, history []string) (string, error) { + m := newPromptModel(prompt, history) + p := tea.NewProgram(m) + finalModel, err := p.Run() + if err != nil { + return "", err + } + result := finalModel.(promptModel) + if result.aborted { + return "", ErrPromptAborted + } + if !result.submitted { + return "", nil + } + return result.value(), nil +} + +func (m promptModel) Init() tea.Cmd { + return tea.Batch(textinput.Blink, tea.EnableBracketedPaste) +} + +func (m promptModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { + switch msg := msg.(type) { + case tea.KeyMsg: + if msg.Type == tea.KeyCtrlC { + m.aborted = true + return m, tea.Quit + } + if msg.Type == tea.KeyTab { + if m.applyCompletion() { + return m, nil + } + } + if msg.Type == tea.KeyEnter { + if !m.multiline && len(m.lines) == 0 { + m.submitted = true + return m, tea.Quit + } + if m.input.Value() == "" { + m.submitted = true + return m, tea.Quit + } + m.lines = append(m.lines, m.input.Value()) + m.input.SetValue("") + m.input.SetCursor(0) + m.multiline = true + m.updatePrompt() + return m, nil + } + + if msg.Paste || (msg.Type == tea.KeyRunes && strings.ContainsRune(string(msg.Runes), '\n')) { + m.handlePaste(string(msg.Runes)) + return m, nil + } + + if msg.Type == tea.KeyRunes || msg.Type == tea.KeyBackspace || msg.Type == tea.KeyDelete { + m.resetCompletion() + } + + if !m.multiline && len(m.lines) == 0 { + switch msg.Type { + case tea.KeyUp: + m.historyUp() + m.updatePrompt() + return m, nil + case tea.KeyDown: + m.historyDown() + m.updatePrompt() + return m, nil + } + } + } + + var cmd tea.Cmd + m.input, cmd = m.input.Update(msg) + m.updatePrompt() + return m, cmd +} + +func (m promptModel) View() string { + var b strings.Builder + if len(m.lines) > 0 { + b.WriteString(m.basePrompt) + b.WriteString(m.lines[0]) + b.WriteByte('\n') + for _, line := range m.lines[1:] { + b.WriteString(m.contPrompt) + b.WriteString(line) + b.WriteByte('\n') + } + } + b.WriteString(m.input.View()) + return b.String() +} + +func (m promptModel) value() string { + if len(m.lines) == 0 { + return m.input.Value() + } + lines := append([]string{}, m.lines...) + if m.input.Value() != "" { + lines = append(lines, m.input.Value()) + } + return strings.Join(lines, "\n") +} + +func (m *promptModel) updatePrompt() { + if len(m.lines) > 0 || m.multiline { + m.input.Prompt = m.contPrompt + } else { + m.input.Prompt = m.basePrompt + } +} + +func (m *promptModel) resetCompletion() { + m.completions = nil + m.completeSeed = "" + m.completeIdx = 0 +} + +func (m *promptModel) handlePaste(paste string) { + if paste == "" { + return + } + paste = strings.ReplaceAll(paste, "\r\n", "\n") + paste = strings.ReplaceAll(paste, "\r", "\n") + + if !strings.Contains(paste, "\n") { + m.insertAtCursor(paste) + return + } + + parts := strings.Split(paste, "\n") + left, right := splitAtRune(m.input.Value(), m.input.Position()) + + first := left + parts[0] + last := parts[len(parts)-1] + right + + m.lines = append(m.lines, first) + if len(parts) > 2 { + m.lines = append(m.lines, parts[1:len(parts)-1]...) + } + + m.input.SetValue(last) + m.input.SetCursor(utf8.RuneCountInString(parts[len(parts)-1])) + m.multiline = true + m.updatePrompt() +} + +func (m *promptModel) insertAtCursor(s string) { + if s == "" { + return + } + left, right := splitAtRune(m.input.Value(), m.input.Position()) + m.input.SetValue(left + s + right) + m.input.SetCursor(m.input.Position() + utf8.RuneCountInString(s)) +} + +func splitAtRune(s string, pos int) (string, string) { + runes := []rune(s) + if pos < 0 { + pos = 0 + } + if pos > len(runes) { + pos = len(runes) + } + return string(runes[:pos]), string(runes[pos:]) +} + +func (m *promptModel) historyUp() { + if len(m.history) == 0 { + return + } + if m.historyIndex == len(m.history) { + m.historyDraft = m.input.Value() + } + if m.historyIndex > 0 { + m.historyIndex-- + m.input.SetValue(m.history[m.historyIndex]) + m.input.SetCursor(len([]rune(m.history[m.historyIndex]))) + } +} + +func (m *promptModel) historyDown() { + if len(m.history) == 0 { + return + } + if m.historyIndex < len(m.history)-1 { + m.historyIndex++ + m.input.SetValue(m.history[m.historyIndex]) + m.input.SetCursor(len([]rune(m.history[m.historyIndex]))) + return + } + if m.historyIndex == len(m.history)-1 { + m.historyIndex = len(m.history) + m.input.SetValue(m.historyDraft) + m.input.SetCursor(len([]rune(m.historyDraft))) + } +} + +func (m *promptModel) applyCompletion() bool { + value := m.input.Value() + if !strings.HasPrefix(value, "/") { + return false + } + if strings.Contains(value, " ") { + return false + } + if value != m.completeSeed || len(m.completions) == 0 { + m.completeSeed = value + m.completions = nil + for _, cmd := range slashCommands { + if strings.HasPrefix(cmd, value) { + m.completions = append(m.completions, cmd) + } + } + m.completeIdx = 0 + } + if len(m.completions) == 0 { + return false + } + if m.completeIdx >= len(m.completions) { + m.completeIdx = 0 + } + next := m.completions[m.completeIdx] + m.completeIdx++ + m.input.SetValue(next) + m.input.SetCursor(len([]rune(next))) + return true +} + +func appendHistory(history []string, item string) []string { + if len(history) == 0 { + return append(history, item) + } + if history[len(history)-1] == item { + return history + } + return append(history, item) +} diff --git a/go.mod b/go.mod index 3c26553..6ff7c33 100644 --- a/go.mod +++ b/go.mod @@ -3,10 +3,11 @@ module github.com/cheikh-seck/godex go 1.25.7 require ( + github.com/charmbracelet/bubbles v1.0.0 + github.com/charmbracelet/bubbletea v1.3.10 github.com/chromedp/chromedp v0.14.2 github.com/jesseduffield/gocui v0.3.0 github.com/mark3labs/mcp-go v0.45.0 - github.com/peterh/liner v1.2.2 github.com/sabhiram/go-gitignore v0.0.0-20210923224102-525f6e181f06 golang.org/x/net v0.38.0 google.golang.org/genai v1.49.0 @@ -17,12 +18,21 @@ require ( cloud.google.com/go v0.116.0 // indirect cloud.google.com/go/auth v0.9.3 // indirect cloud.google.com/go/compute/metadata v0.5.0 // indirect + github.com/atotto/clipboard v0.1.4 // indirect + github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect github.com/bahlo/generic-list-go v0.2.0 // indirect github.com/buger/jsonparser v1.1.1 // indirect + github.com/charmbracelet/colorprofile v0.4.1 // indirect + github.com/charmbracelet/lipgloss v1.1.0 // indirect + github.com/charmbracelet/x/ansi v0.11.6 // indirect + github.com/charmbracelet/x/cellbuf v0.0.15 // indirect + github.com/charmbracelet/x/term v0.2.2 // indirect github.com/chromedp/cdproto v0.0.0-20250724212937-08a3db8b4327 // indirect github.com/chromedp/sysutil v1.1.0 // indirect + github.com/clipperhouse/displaywidth v0.9.0 // indirect github.com/clipperhouse/stringish v0.1.1 // indirect github.com/clipperhouse/uax29/v2 v2.5.0 // indirect + github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect github.com/go-json-experiment/json v0.0.0-20250725192818-e39067aee2d2 // indirect github.com/gobwas/httphead v0.1.0 // indirect github.com/gobwas/pool v0.2.1 // indirect @@ -34,11 +44,19 @@ require ( github.com/googleapis/enterprise-certificate-proxy v0.3.4 // indirect github.com/gorilla/websocket v1.5.3 // indirect github.com/invopop/jsonschema v0.13.0 // indirect + github.com/lucasb-eyer/go-colorful v1.3.0 // indirect github.com/mailru/easyjson v0.7.7 // indirect + github.com/mattn/go-isatty v0.0.20 // indirect + github.com/mattn/go-localereader v0.0.1 // indirect github.com/mattn/go-runewidth v0.0.19 // indirect + github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect + github.com/muesli/cancelreader v0.2.2 // indirect + github.com/muesli/termenv v0.16.0 // indirect github.com/nsf/termbox-go v1.1.1 // indirect + github.com/rivo/uniseg v0.4.7 // indirect github.com/spf13/cast v1.7.1 // indirect github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect + github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect github.com/yosida95/uritemplate/v3 v3.0.2 // indirect go.opencensus.io v0.24.0 // indirect golang.org/x/crypto v0.36.0 // indirect diff --git a/go.sum b/go.sum index aef7652..ed21e09 100644 --- a/go.sum +++ b/go.sum @@ -6,11 +6,29 @@ cloud.google.com/go/auth v0.9.3/go.mod h1:7z6VY+7h3KUdRov5F1i8NDP5ZzWKYmEPO842Bg cloud.google.com/go/compute/metadata v0.5.0 h1:Zr0eK8JbFv6+Wi4ilXAR8FJ3wyNdpxHKJNPos6LTZOY= cloud.google.com/go/compute/metadata v0.5.0/go.mod h1:aHnloV2TPI38yx4s9+wAZhHykWvVCfu7hQbF+9CWoiY= github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= +github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4= +github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI= +github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k= +github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8= github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs= github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0= github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU= +github.com/charmbracelet/bubbles v1.0.0 h1:12J8/ak/uCZEMQ6KU7pcfwceyjLlWsDLAxB5fXonfvc= +github.com/charmbracelet/bubbles v1.0.0/go.mod h1:9d/Zd5GdnauMI5ivUIVisuEm3ave1XwXtD1ckyV6r3E= +github.com/charmbracelet/bubbletea v1.3.10 h1:otUDHWMMzQSB0Pkc87rm691KZ3SWa4KUlvF9nRvCICw= +github.com/charmbracelet/bubbletea v1.3.10/go.mod h1:ORQfo0fk8U+po9VaNvnV95UPWA1BitP1E0N6xJPlHr4= +github.com/charmbracelet/colorprofile v0.4.1 h1:a1lO03qTrSIRaK8c3JRxJDZOvhvIeSco3ej+ngLk1kk= +github.com/charmbracelet/colorprofile v0.4.1/go.mod h1:U1d9Dljmdf9DLegaJ0nGZNJvoXAhayhmidOdcBwAvKk= +github.com/charmbracelet/lipgloss v1.1.0 h1:vYXsiLHVkK7fp74RkV7b2kq9+zDLoEU4MZoFqR/noCY= +github.com/charmbracelet/lipgloss v1.1.0/go.mod h1:/6Q8FR2o+kj8rz4Dq0zQc3vYf7X+B0binUUBwA0aL30= +github.com/charmbracelet/x/ansi v0.11.6 h1:GhV21SiDz/45W9AnV2R61xZMRri5NlLnl6CVF7ihZW8= +github.com/charmbracelet/x/ansi v0.11.6/go.mod h1:2JNYLgQUsyqaiLovhU2Rv/pb8r6ydXKS3NIttu3VGZQ= +github.com/charmbracelet/x/cellbuf v0.0.15 h1:ur3pZy0o6z/R7EylET877CBxaiE1Sp1GMxoFPAIztPI= +github.com/charmbracelet/x/cellbuf v0.0.15/go.mod h1:J1YVbR7MUuEGIFPCaaZ96KDl5NoS0DAWkskup+mOY+Q= +github.com/charmbracelet/x/term v0.2.2 h1:xVRT/S2ZcKdhhOuSP4t5cLi5o+JxklsoEObBSgfgZRk= +github.com/charmbracelet/x/term v0.2.2/go.mod h1:kF8CY5RddLWrsgVwpw4kAa6TESp6EB5y3uxGLeCqzAI= github.com/chromedp/cdproto v0.0.0-20250724212937-08a3db8b4327 h1:UQ4AU+BGti3Sy/aLU8KVseYKNALcX9UXY6DfpwQ6J8E= github.com/chromedp/cdproto v0.0.0-20250724212937-08a3db8b4327/go.mod h1:NItd7aLkcfOA/dcMXvl8p1u+lQqioRMq/SqDp71Pb/k= github.com/chromedp/chromedp v0.14.2 h1:r3b/WtwM50RsBZHMUm9fsNhhzRStTHrKdr2zmwbZSzM= @@ -18,6 +36,8 @@ github.com/chromedp/chromedp v0.14.2/go.mod h1:rHzAv60xDE7VNy/MYtTUrYreSc0ujt2O1 github.com/chromedp/sysutil v1.1.0 h1:PUFNv5EcprjqXZD9nJb9b/c9ibAbxiYo4exNWZyipwM= github.com/chromedp/sysutil v1.1.0/go.mod h1:WiThHUdltqCNKGc4gaU50XgYjwjYIhKWoHGPTUfWTJ8= github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw= +github.com/clipperhouse/displaywidth v0.9.0 h1:Qb4KOhYwRiN3viMv1v/3cTBlz3AcAZX3+y9OLhMtAtA= +github.com/clipperhouse/displaywidth v0.9.0/go.mod h1:aCAAqTlh4GIVkhQnJpbL0T/WfcrJXHcj8C0yjYcjOZA= github.com/clipperhouse/stringish v0.1.1 h1:+NSqMOr3GR6k1FdRhhnXrLfztGzuG+VuFDfatpWHKCs= github.com/clipperhouse/stringish v0.1.1/go.mod h1:v/WhFtE1q0ovMta2+m+UbpZ+2/HEXNWYXQgCt4hdOzA= github.com/clipperhouse/uax29/v2 v2.5.0 h1:x7T0T4eTHDONxFJsL94uKNKPHrclyFI0lm7+w94cO8U= @@ -30,6 +50,8 @@ github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymF github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= github.com/envoyproxy/go-control-plane v0.9.4/go.mod h1:6rpuAdCZL397s3pYoYcLgu1mIlRU8Am5FuJP05cCM98= github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c= +github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f h1:Y/CXytFA4m6baUTXGLOoWe4PQhGxaX0KpnayAqC48p4= +github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f/go.mod h1:vw97MGsxSvLiUE2X8qFplwetxpGLQrlU1Q9AUEIzCaM= github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8= github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0= github.com/go-json-experiment/json v0.0.0-20250725192818-e39067aee2d2 h1:iizUGZ9pEquQS5jTGkh4AqeeHCMbfbjeb0zMt0aEFzs= @@ -82,14 +104,26 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= github.com/ledongthuc/pdf v0.0.0-20220302134840-0c2507a12d80 h1:6Yzfa6GP0rIo/kULo2bwGEkFvCePZ3qHDDTC3/J9Swo= github.com/ledongthuc/pdf v0.0.0-20220302134840-0c2507a12d80/go.mod h1:imJHygn/1yfhB7XSJJKlFZKl/J+dCPAknuiaGOshXAs= +github.com/lucasb-eyer/go-colorful v1.3.0 h1:2/yBRLdWBZKrf7gB40FoiKfAWYQ0lqNcbuQwVHXptag= +github.com/lucasb-eyer/go-colorful v1.3.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0= github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0= github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc= github.com/mark3labs/mcp-go v0.45.0 h1:s0S8qR/9fWaQ3pHxz7pm1uQ0DrswoSnRIxKIjbiQtkc= github.com/mark3labs/mcp-go v0.45.0/go.mod h1:YnJfOL382MIWDx1kMY+2zsRHU/q78dBg9aFb8W6Thdw= +github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= +github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/mattn/go-localereader v0.0.1 h1:ygSAOl7ZXTx4RdPYinUpg6W99U8jWvWi9Ye2JC/oIi4= +github.com/mattn/go-localereader v0.0.1/go.mod h1:8fBrzywKY7BI3czFoHkuzRoWE9C+EiG4R1k4Cjx5p88= github.com/mattn/go-runewidth v0.0.3/go.mod h1:LwmH8dsx7+W8Uxz3IHJYH5QSwggIsqBzpuz5H//U1FU= github.com/mattn/go-runewidth v0.0.9/go.mod h1:H031xJmbD/WCDINGzjvQ9THkh0rPKHF+m2gUSrubnMI= github.com/mattn/go-runewidth v0.0.19 h1:v++JhqYnZuu5jSKrk9RbgF5v4CGUjqRfBm05byFGLdw= github.com/mattn/go-runewidth v0.0.19/go.mod h1:XBkDxAl56ILZc9knddidhrOlY5R/pDhgLpndooCuJAs= +github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 h1:ZK8zHtRHOkbHy6Mmr5D264iyp3TiX5OmNcI5cIARiQI= +github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6/go.mod h1:CJlz5H+gyd6CUWT45Oy4q24RdLyn7Md9Vj2/ldJBSIo= +github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA= +github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo= +github.com/muesli/termenv v0.16.0 h1:S5AlUN9dENB57rsbnkPyfdGuWIlkmzJjbFf0Tf5FWUc= +github.com/muesli/termenv v0.16.0/go.mod h1:ZRfOIKPFDYQoDFF4Olj7/QJbW60Ol/kL1pU3VfY/Cnk= github.com/nsf/termbox-go v1.1.1 h1:nksUPLCb73Q++DwbYUBEglYBRPZyoXJdrj5L+TkjyZY= github.com/nsf/termbox-go v1.1.1/go.mod h1:T0cTdVuOwf7pHQNtfhnEbzHbcNyCEcVU4YPpouCbVxo= github.com/orisano/pixelmatch v0.0.0-20220722002657-fb0b55479cde h1:x0TT0RDC7UhAVbbWWBzr41ElhJx5tXPWkIHA2HWPRuw= @@ -99,6 +133,8 @@ github.com/peterh/liner v1.2.2/go.mod h1:xFwJyiKIXJZUKItq5dGHZSTBRAuG/CpeNpWLyiN github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA= +github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ= +github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZVejAe8= github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= github.com/sabhiram/go-gitignore v0.0.0-20210923224102-525f6e181f06 h1:OkMGxebDjyw0ULyrTYWeN0UNCCkmCWfjPnIA2W6oviI= @@ -116,6 +152,8 @@ github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsT github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw= +github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no= +github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM= github.com/yosida95/uritemplate/v3 v3.0.2 h1:Ed3Oyj9yrmi9087+NczuL5BwkIc4wvTb5zIM+UJPGz4= github.com/yosida95/uritemplate/v3 v3.0.2/go.mod h1:ILOh0sOhIJR3+L/8afwt/kE++YT040gmv5BQTMR2HP4= go.opencensus.io v0.24.0 h1:y73uSU6J157QMP2kn2r30vwW1A2W2WFwSCGnAVxeaD0= @@ -146,6 +184,7 @@ golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5h golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20210809222454-d867a43fc93e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20211117180635-dee7805ff2e1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.38.0 h1:3yZWxaJjBmCWXqhN1qh02AkOnCQ1poK6oF+a7xWL6Gc= diff --git a/internal/mcp/bash.go b/internal/mcp/bash.go index a1b5665..d143ae3 100644 --- a/internal/mcp/bash.go +++ b/internal/mcp/bash.go @@ -17,7 +17,12 @@ func GetWorkingDir() (string, error) { type BashServer struct { allowedPaths []string tools []Tool - backgroundPIDs []int + backgroundPIDs []bgProcess +} + +type bgProcess struct { + PID int + Command string } func NewBashServer(allowedPaths []string) *BashServer { @@ -124,7 +129,10 @@ func (s *BashServer) runCommand(args map[string]interface{}) (string, error) { if err != nil { return "", fmt.Errorf("failed to start background command: %v", err) } - s.backgroundPIDs = append(s.backgroundPIDs, cmd.Process.Pid) + s.backgroundPIDs = append(s.backgroundPIDs, bgProcess{ + PID: cmd.Process.Pid, + Command: strings.TrimSpace(command), + }) return fmt.Sprintf("Started background process (PID: %d)\nOutput will be in nohup.out", cmd.Process.Pid), nil } @@ -167,7 +175,7 @@ func (s *BashServer) killCommand(args map[string]interface{}) (string, error) { // Remove from tracking for i, p := range s.backgroundPIDs { - if p == int(pid) { + if p.PID == int(pid) { s.backgroundPIDs = append(s.backgroundPIDs[:i], s.backgroundPIDs[i+1:]...) break } @@ -183,16 +191,16 @@ func (s *BashServer) killAllBackground(args map[string]interface{}) (string, err var killed []int var failed []int - for _, pid := range s.backgroundPIDs { - proc, err := os.FindProcess(pid) + for _, procInfo := range s.backgroundPIDs { + proc, err := os.FindProcess(procInfo.PID) if err == nil { if err := proc.Kill(); err == nil { - killed = append(killed, pid) + killed = append(killed, procInfo.PID) } else { - failed = append(failed, pid) + failed = append(failed, procInfo.PID) } } else { - failed = append(failed, pid) + failed = append(failed, procInfo.PID) } } @@ -285,12 +293,16 @@ func (s *BashServer) ListBackground() (string, error) { return "No background processes", nil } var pids []string - for _, pid := range s.backgroundPIDs { - proc, err := os.FindProcess(pid) + for _, procInfo := range s.backgroundPIDs { + proc, err := os.FindProcess(procInfo.PID) if err == nil { err = proc.Signal(nil) // Check if process is alive if err == nil { - pids = append(pids, fmt.Sprintf("PID %d (running)", pid)) + if procInfo.Command != "" { + pids = append(pids, fmt.Sprintf("PID %d (running) - %s", procInfo.PID, procInfo.Command)) + } else { + pids = append(pids, fmt.Sprintf("PID %d (running)", procInfo.PID)) + } } } } @@ -306,13 +318,13 @@ func (s *BashServer) BackgroundCount() int { } func (s *BashServer) cleanupDeadProcesses() { - var alive []int - for _, pid := range s.backgroundPIDs { - proc, err := os.FindProcess(pid) + var alive []bgProcess + for _, procInfo := range s.backgroundPIDs { + proc, err := os.FindProcess(procInfo.PID) if err == nil { err = proc.Signal(nil) // Signal 0 checks if process exists if err == nil { - alive = append(alive, pid) + alive = append(alive, procInfo) } } } diff --git a/internal/providers/gemini.go b/internal/providers/gemini.go index 4c3f455..5c93af4 100644 --- a/internal/providers/gemini.go +++ b/internal/providers/gemini.go @@ -5,6 +5,7 @@ import ( "fmt" "os" "strings" + "sync" "google.golang.org/genai" @@ -17,6 +18,9 @@ type geminiProvider struct { client *genai.Client model string temperature *float64 + cancelMu sync.Mutex + cancelFunc context.CancelFunc + cancelGen uint64 } func init() { @@ -72,6 +76,21 @@ func (g *geminiProvider) Send(ctx context.Context, prompt string) (string, error } } + g.cancelMu.Lock() + ctx, cancel := context.WithCancel(ctx) + g.cancelGen++ + gen := g.cancelGen + g.cancelFunc = cancel + g.cancelMu.Unlock() + + defer func() { + g.cancelMu.Lock() + if g.cancelGen == gen { + g.cancelFunc = nil + } + g.cancelMu.Unlock() + }() + resp, err := g.client.Models.GenerateContent(ctx, g.model, genai.Text(prompt), config) if err != nil { return "", err @@ -91,6 +110,11 @@ func (g *geminiProvider) SetThinkCallback(fn func(string)) { } func (g *geminiProvider) Cancel() { + g.cancelMu.Lock() + defer g.cancelMu.Unlock() + if g.cancelFunc != nil { + g.cancelFunc() + } } func (g *geminiProvider) CallTool(ctx context.Context, name string, args map[string]interface{}) (string, error) { diff --git a/internal/providers/ollama.go b/internal/providers/ollama.go index 2a91377..aa1a13f 100644 --- a/internal/providers/ollama.go +++ b/internal/providers/ollama.go @@ -27,8 +27,10 @@ type ollamaProvider struct { client *http.Client messages []map[string]string mu sync.Mutex + sendMu sync.Mutex OnThink func(string) cancelFunc context.CancelFunc + cancelGen uint64 } func init() { @@ -60,15 +62,28 @@ func newOllamaProvider(cfg *config.Provider) (Provider, error) { } func (o *ollamaProvider) Send(ctx context.Context, prompt string) (string, error) { - o.mu.Lock() - defer o.mu.Unlock() + o.sendMu.Lock() + defer o.sendMu.Unlock() + o.mu.Lock() o.messages = append(o.messages, map[string]string{ "role": "user", "content": prompt, }) - ctx, o.cancelFunc = context.WithCancel(ctx) + ctx, cancel := context.WithCancel(ctx) + o.cancelGen++ + gen := o.cancelGen + o.cancelFunc = cancel + o.mu.Unlock() + + defer func() { + o.mu.Lock() + if o.cancelGen == gen { + o.cancelFunc = nil + } + o.mu.Unlock() + }() endpoint := o.baseURL + "/api/chat" reqBody := map[string]any{ @@ -137,17 +152,19 @@ func (o *ollamaProvider) Send(ctx context.Context, prompt string) (string, error return "", fmt.Errorf("ollama returned empty response") } + o.mu.Lock() o.messages = append(o.messages, map[string]string{ "role": "assistant", "content": response, }) + o.mu.Unlock() return response, nil } func (o *ollamaProvider) Close() error { - o.mu.Lock() - defer o.mu.Unlock() + o.sendMu.Lock() + defer o.sendMu.Unlock() ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() @@ -174,7 +191,9 @@ func (o *ollamaProvider) Close() error { } defer resp.Body.Close() + o.mu.Lock() o.messages = nil + o.mu.Unlock() return nil }