From 875f38cc334a14b66b12d2a1372fd2943089700c Mon Sep 17 00:00:00 2001
From: Carlos Alexandro Becker <caarlos0@users.noreply.github.com>
Date: Thu, 9 Oct 2025 16:58:26 -0300
Subject: [PATCH] refactor: more cleanup

Signed-off-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com>
---
 .../tui/components/dialogs/commands/loader.go | 32 ++++++-------
 internal/tui/tui.go                           | 47 +++----------------
 2 files changed, 21 insertions(+), 58 deletions(-)

diff --git a/internal/tui/components/dialogs/commands/loader.go b/internal/tui/components/dialogs/commands/loader.go
index d5449a8f41d3bc990d3db49b96ddebd2a695d9b2..774888bcb91f3a0e9dab82c3e9dd7f7a134306f9 100644
--- a/internal/tui/components/dialogs/commands/loader.go
+++ b/internal/tui/components/dialogs/commands/loader.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"io/fs"
+	"log/slog"
 	"os"
 	"path/filepath"
 	"regexp"
@@ -214,45 +215,40 @@ func LoadMCPPrompts() []Command {
 	commands := make([]Command, 0, len(prompts))
 
 	for key, prompt := range prompts {
-		p := prompt
 		clientName, promptName, ok := strings.Cut(key, ":")
 		if !ok {
+			slog.Warn("prompt not found", "key", key)
 			continue
 		}
 		commands = append(commands, Command{
 			ID:          key,
 			Title:       clientName + ":" + promptName,
-			Description: p.Description,
-			Handler:     createMCPPromptHandler(key, promptName, p),
+			Description: prompt.Description,
+			Handler:     createMCPPromptHandler(clientName, promptName, prompt),
 		})
 	}
 
 	return commands
 }
 
-func createMCPPromptHandler(key, promptName string, prompt *mcp.Prompt) func(Command) tea.Cmd {
+func createMCPPromptHandler(clientName, promptName string, prompt *mcp.Prompt) func(Command) tea.Cmd {
 	return func(cmd Command) tea.Cmd {
 		if len(prompt.Arguments) == 0 {
-			return executeMCPPromptWithoutArgs(key, promptName)
+			return execMCPPrompt(clientName, promptName, nil)
 		}
 		return util.CmdHandler(ShowMCPPromptArgumentsDialogMsg{
-			PromptID:   cmd.ID,
-			PromptName: promptName,
+			Prompt: prompt,
+			OnSubmit: func(args map[string]string) tea.Cmd {
+				return execMCPPrompt(clientName, promptName, args)
+			},
 		})
 	}
 }
 
-func executeMCPPromptWithoutArgs(key, promptName string) tea.Cmd {
+func execMCPPrompt(clientName, promptName string, args map[string]string) tea.Cmd {
 	return func() tea.Msg {
-		// key format is "clientName:promptName"
-		parts := strings.SplitN(key, ":", 2)
-		if len(parts) != 2 {
-			return util.ReportError(fmt.Errorf("invalid prompt key: %s", key))
-		}
-		clientName := parts[0]
-
 		ctx := context.Background()
-		result, err := agent.GetMCPPromptContent(ctx, clientName, promptName, nil)
+		result, err := agent.GetMCPPromptContent(ctx, clientName, promptName, args)
 		if err != nil {
 			return util.ReportError(err)
 		}
@@ -274,6 +270,6 @@ func executeMCPPromptWithoutArgs(key, promptName string) tea.Cmd {
 }
 
 type ShowMCPPromptArgumentsDialogMsg struct {
-	PromptID   string
-	PromptName string
+	Prompt   *mcp.Prompt
+	OnSubmit func(arg map[string]string) tea.Cmd
 }
diff --git a/internal/tui/tui.go b/internal/tui/tui.go
index 353b55dd6be75c360ff50f31db61c47cd32e96b2..e0e5ad64b3054223c9b110385950a7ca4b87e701 100644
--- a/internal/tui/tui.go
+++ b/internal/tui/tui.go
@@ -3,7 +3,6 @@ package tui
 import (
 	"context"
 	"fmt"
-	"log/slog"
 	"math/rand"
 	"strings"
 	"time"
@@ -35,7 +34,6 @@ import (
 	"github.com/charmbracelet/crush/internal/tui/styles"
 	"github.com/charmbracelet/crush/internal/tui/util"
 	"github.com/charmbracelet/lipgloss/v2"
-	"github.com/modelcontextprotocol/go-sdk/mcp"
 	"golang.org/x/text/cases"
 	"golang.org/x/text/language"
 )
@@ -175,48 +173,17 @@ func (a *appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 			},
 		)
 	case commands.ShowMCPPromptArgumentsDialogMsg:
-		prompt, ok := agent.GetMCPPrompt(msg.PromptID)
-		clientName, _, ok := strings.Cut(msg.PromptID, ":")
-		if !ok {
-			slog.Warn("prompt not found", "prompt_id", msg.PromptID, "prompt_name", msg.PromptName)
-			util.ReportWarn(fmt.Sprintf("Prompt %s not found", msg.PromptName))
-			return a, nil
-		}
-		args := make([]commands.Argument, 0, len(prompt.Arguments))
-		for _, arg := range prompt.Arguments {
+		args := make([]commands.Argument, 0, len(msg.Prompt.Arguments))
+		for _, arg := range msg.Prompt.Arguments {
 			args = append(args, commands.Argument(*arg))
 		}
 		dialog := commands.NewCommandArgumentsDialog(
-			msg.PromptID,
-			prompt.Title,
-			prompt.Name,
-			prompt.Description,
+			msg.Prompt.Name,
+			msg.Prompt.Title,
+			msg.Prompt.Name,
+			msg.Prompt.Description,
 			args,
-			func(args map[string]string) tea.Cmd {
-				return func() tea.Msg {
-					ctx := context.Background()
-					result, err := agent.GetMCPPromptContent(ctx, clientName, prompt.Name, args)
-					if err != nil {
-						return util.ReportError(err)
-					}
-
-					var content strings.Builder
-					for _, msg := range result.Messages {
-						if msg.Role != "user" {
-							continue
-						}
-						textContent, ok := msg.Content.(*mcp.TextContent)
-						if !ok {
-							continue
-						}
-						_, _ = content.WriteString(textContent.Text)
-						_, _ = content.WriteString("\n")
-					}
-					return commands.CommandRunCustomMsg{
-						Content: content.String(),
-					}
-				}
-			},
+			msg.OnSubmit,
 		)
 		return a, util.CmdHandler(
 			dialogs.OpenDialogMsg{