From a2f7cbbfb3b1ba1715095161f22b21482396b06c Mon Sep 17 00:00:00 2001 From: Carlos Alexandro Becker Date: Mon, 3 Nov 2025 15:26:15 -0300 Subject: [PATCH] fix: improve code Signed-off-by: Carlos Alexandro Becker --- internal/agent/tools/fetch.go | 6 ---- internal/agent/tools/mcp-tools.go | 18 ++++++------ internal/agent/tools/mcp/init.go | 8 +++--- internal/agent/tools/mcp/prompts.go | 25 +++++------------ internal/agent/tools/mcp/tools.go | 25 +++++------------ .../tui/components/dialogs/commands/loader.go | 28 +++++++++---------- 6 files changed, 41 insertions(+), 69 deletions(-) diff --git a/internal/agent/tools/fetch.go b/internal/agent/tools/fetch.go index 0701bd2be151213f2fd0e313726d03a5c35833a3..e25404192bae81895a78f205948f83f844082ff5 100644 --- a/internal/agent/tools/fetch.go +++ b/internal/agent/tools/fetch.go @@ -28,12 +28,6 @@ type FetchPermissionsParams struct { Timeout int `json:"timeout,omitempty"` } -type fetchTool struct { - client *http.Client - permissions permission.Service - workingDir string -} - const FetchToolName = "fetch" //go:embed fetch.md diff --git a/internal/agent/tools/mcp-tools.go b/internal/agent/tools/mcp-tools.go index ff3be182e87cb2bff82be3041b62aa4f126c0be2..e75a6ccb7c92c6a4925960e0590eb5f3b2bac47e 100644 --- a/internal/agent/tools/mcp-tools.go +++ b/internal/agent/tools/mcp-tools.go @@ -12,13 +12,15 @@ import ( // GetMCPTools gets all the currently available MCP tools. func GetMCPTools(permissions permission.Service, wd string) []*Tool { var result []*Tool - for name, tool := range mcp.Tools() { - result = append(result, &Tool{ - mcpName: name, - tool: tool, - permissions: permissions, - workingDir: wd, - }) + for mcpName, tools := range mcp.Tools() { + for _, tool := range tools { + result = append(result, &Tool{ + mcpName: mcpName, + tool: tool, + permissions: permissions, + workingDir: wd, + }) + } } return result } @@ -74,7 +76,7 @@ func (m *Tool) Info() fantasy.ToolInfo { } return fantasy.ToolInfo{ - Name: fmt.Sprintf("mcp_%s_%s", m.mcpName, m.tool.Name), + Name: m.Name(), Description: m.tool.Description, Parameters: parameters, Required: required, diff --git a/internal/agent/tools/mcp/init.go b/internal/agent/tools/mcp/init.go index 961ce58071cb90808db52654629fccd7b1030b8c..cb97c0e440d87108d08a0c572c8ec2160ee5cf17 100644 --- a/internal/agent/tools/mcp/init.go +++ b/internal/agent/tools/mcp/init.go @@ -58,12 +58,12 @@ func (s State) String() string { } // EventType represents the type of MCP event -type EventType string +type EventType uint const ( - EventStateChanged EventType = "state_changed" - EventToolsListChanged EventType = "tools_list_changed" - EventPromptsListChanged EventType = "prompts_list_changed" + EventStateChanged EventType = iota + EventToolsListChanged + EventPromptsListChanged ) // Event represents an event in the MCP system diff --git a/internal/agent/tools/mcp/prompts.go b/internal/agent/tools/mcp/prompts.go index e168b71f2137963c4c7e428e21edd11d10ac6a9d..0bd6e665dd80dad90c844d7d31c61c506ea83803 100644 --- a/internal/agent/tools/mcp/prompts.go +++ b/internal/agent/tools/mcp/prompts.go @@ -11,13 +11,10 @@ import ( type Prompt = mcp.Prompt -var ( - allPrompts = csync.NewMap[string, *Prompt]() - clientPrompts = csync.NewMap[string, []*Prompt]() -) +var allPrompts = csync.NewMap[string, []*Prompt]() // Prompts returns all available MCP prompts. -func Prompts() iter.Seq2[string, *Prompt] { +func Prompts() iter.Seq2[string, []*Prompt] { return allPrompts.Seq2() } @@ -65,10 +62,8 @@ func RefreshPrompts(ctx context.Context, name string) { updatePrompts(name, prompts) prev, _ := states.Get(name) - updateState(name, StateConnected, nil, session, Counts{ - Prompts: len(prompts), - Tools: prev.Counts.Tools, - }) + prev.Counts.Prompts = len(prompts) + updateState(name, StateConnected, nil, session, prev.Counts) } func getPrompts(ctx context.Context, c *mcp.ClientSession) ([]*Prompt, error) { @@ -85,14 +80,8 @@ func getPrompts(ctx context.Context, c *mcp.ClientSession) ([]*Prompt, error) { // updatePrompts updates the global mcpPrompts and mcpClient2Prompts maps func updatePrompts(mcpName string, prompts []*Prompt) { if len(prompts) == 0 { - clientPrompts.Del(mcpName) - } else { - clientPrompts.Set(mcpName, prompts) - } - for mcpName, prompts := range clientPrompts.Seq2() { - for _, p := range prompts { - key := mcpName + ":" + p.Name - allPrompts.Set(key, p) - } + allPrompts.Del(mcpName) + return } + allPrompts.Set(mcpName, prompts) } diff --git a/internal/agent/tools/mcp/tools.go b/internal/agent/tools/mcp/tools.go index e39e999352373f7c2a140600a84c7ed70571e1b6..561783e6f7986a772b2d1deb9bcdd6dac2bf8f34 100644 --- a/internal/agent/tools/mcp/tools.go +++ b/internal/agent/tools/mcp/tools.go @@ -14,13 +14,10 @@ import ( type Tool = mcp.Tool -var ( - allTools = csync.NewMap[string, *Tool]() - clientTools = csync.NewMap[string, []*Tool]() -) +var allTools = csync.NewMap[string, []*Tool]() // Tools returns all available MCP tools. -func Tools() iter.Seq2[string, *Tool] { +func Tools() iter.Seq2[string, []*Tool] { return allTools.Seq2() } @@ -72,10 +69,8 @@ func RefreshTools(ctx context.Context, name string) { updateTools(name, tools) prev, _ := states.Get(name) - updateState(name, StateConnected, nil, session, Counts{ - Tools: len(tools), - Prompts: prev.Counts.Prompts, - }) + prev.Counts.Tools = len(tools) + updateState(name, StateConnected, nil, session, prev.Counts) } func getTools(ctx context.Context, session *mcp.ClientSession) ([]*Tool, error) { @@ -89,16 +84,10 @@ func getTools(ctx context.Context, session *mcp.ClientSession) ([]*Tool, error) return result.Tools, nil } -// updateTools updates the global mcpTools and mcpClient2Tools maps func updateTools(name string, tools []*Tool) { if len(tools) == 0 { - clientTools.Del(name) - } else { - clientTools.Set(name, tools) - } - for name, tools := range clientTools.Seq2() { - for _, t := range tools { - allTools.Set(name, t) - } + allTools.Del(name) + return } + allTools.Set(name, tools) } diff --git a/internal/tui/components/dialogs/commands/loader.go b/internal/tui/components/dialogs/commands/loader.go index 9e20aea7480ac58456fa22eb73057d9d6d1115e6..3b51604818a77ca12540fa7adcced0726689f884 100644 --- a/internal/tui/components/dialogs/commands/loader.go +++ b/internal/tui/components/dialogs/commands/loader.go @@ -1,10 +1,10 @@ package commands import ( + "cmp" "context" "fmt" "io/fs" - "log/slog" "os" "path/filepath" "regexp" @@ -223,32 +223,30 @@ type CommandRunCustomMsg struct { func loadMCPPrompts() []Command { var commands []Command - for key, prompt := range mcp.Prompts() { - clientName, promptName, ok := strings.Cut(key, ":") - if !ok { - slog.Warn("prompt not found", "key", key) - continue + for mcpName, prompts := range mcp.Prompts() { + for _, prompt := range prompts { + key := mcpName + ":" + prompt.Name + commands = append(commands, Command{ + ID: key, + Title: cmp.Or(prompt.Title, prompt.Name), + Description: prompt.Description, + Handler: createMCPPromptHandler(mcpName, prompt.Name, prompt), + }) } - commands = append(commands, Command{ - ID: key, - Title: clientName + ":" + promptName, - Description: prompt.Description, - Handler: createMCPPromptHandler(clientName, promptName, prompt), - }) } return commands } -func createMCPPromptHandler(clientName, promptName string, prompt *mcp.Prompt) func(Command) tea.Cmd { +func createMCPPromptHandler(mcpName, promptName string, prompt *mcp.Prompt) func(Command) tea.Cmd { return func(cmd Command) tea.Cmd { if len(prompt.Arguments) == 0 { - return execMCPPrompt(clientName, promptName, nil) + return execMCPPrompt(mcpName, promptName, nil) } return util.CmdHandler(ShowMCPPromptArgumentsDialogMsg{ Prompt: prompt, OnSubmit: func(args map[string]string) tea.Cmd { - return execMCPPrompt(clientName, promptName, args) + return execMCPPrompt(mcpName, promptName, args) }, }) }