e519d6f · crush · git.iko.soy

internal/agent/tools/mcp-tools.go+1−1

compare with main

unifiedsplit

@@ -12,7 +12,7 @@ import (

12// GetMCPTools gets all the currently available MCP tools.

13func GetMCPTools(permissions permission.Service, wd string) []*Tool {

14 var result []*Tool

15− for name, tool := range mcp.GetMCPTools() {

15+ for name, tool := range mcp.Tools() {

16 result = append(result, &Tool{

17 mcpName: name,

18 tool: tool,

internal/agent/tools/mcp/init.go+3

@@ -281,6 +281,9 @@ func createSession(ctx context.Context, name string, m config.MCPConfig, resolve

281 Name: name,

282 })

283 },

284+ LoggingMessageHandler: func(_ context.Context, req *mcp.LoggingMessageRequest) {

285+ slog.Info("mcp log", "name", name, "data", req.Params.Data)

286+ },

284 KeepAlive: time.Minute * 10,

287 KeepAlive: time.Minute * 10,

285 },

288 },

286 )

289 )

internal/agent/tools/mcp/prompts.go+32−17

@@ -3,6 +3,7 @@ package mcp

3import (

4 "context"

5 "iter"

6+ "log/slog"

6

7

7 "github.com/charmbracelet/crush/internal/csync"

8 "github.com/charmbracelet/crush/internal/csync"

8 "github.com/modelcontextprotocol/go-sdk/mcp"

9 "github.com/modelcontextprotocol/go-sdk/mcp"

@@ -11,25 +12,15 @@ import (

11type Prompt = mcp.Prompt

12type Prompt = mcp.Prompt

12

13

13var (

14var (

14− allPrompts = csync.NewMap[string, *Prompt]()

15+ allPrompts = csync.NewMap[string, *Prompt]()

15− client2Prompts = csync.NewMap[string, []*Prompt]()

16+ clientPrompts = csync.NewMap[string, []*Prompt]()

16)

17)

17

18

18−// GetPrompts returns all available MCP prompts.

19+// Prompts returns all available MCP prompts.

19−func GetPrompts() iter.Seq2[string, *Prompt] {

20+func Prompts() iter.Seq2[string, *Prompt] {

20 return allPrompts.Seq2()

21 return allPrompts.Seq2()

21}

22}

22

23

23−// GetPrompt returns a specific MCP prompt by name.

24−func GetPrompt(name string) (*Prompt, bool) {

25− return allPrompts.Get(name)

26−}

27−

28−// GetPromptsByClient returns all prompts for a specific MCP client.

29−func GetPromptsByClient(clientName string) ([]*Prompt, bool) {

30− return client2Prompts.Get(clientName)

31−}

32−

33// GetPromptMessages retrieves the content of an MCP prompt with the given arguments.

24// GetPromptMessages retrieves the content of an MCP prompt with the given arguments.

34func GetPromptMessages(ctx context.Context, clientName, promptName string, args map[string]string) ([]string, error) {

25func GetPromptMessages(ctx context.Context, clientName, promptName string, args map[string]string) ([]string, error) {

35 c, err := getOrRenewClient(ctx, clientName)

26 c, err := getOrRenewClient(ctx, clientName)

@@ -56,6 +47,30 @@ func GetPromptMessages(ctx context.Context, clientName, promptName string, args

56 return messages, nil

47 return messages, nil

57}

48}

58

49

50+// RefreshPrompts gets the updated list of prompts from the MCP and updates the

51+// global state.

52+func RefreshPrompts(ctx context.Context, name string) {

53+ session, ok := sessions.Get(name)

54+ if !ok {

55+ slog.Warn("refresh prompts: no session", "name", name)

56+ return

57+ }

58+

59+ prompts, err := getPrompts(ctx, session)

60+ if err != nil {

61+ updateState(name, StateError, err, nil, Counts{})

62+ return

63+ }

64+

65+ updatePrompts(name, prompts)

66+

67+ prev, _ := states.Get(name)

68+ updateState(name, StateConnected, nil, session, Counts{

69+ Prompts: len(prompts),

70+ Tools: prev.Counts.Tools,

71+ })

72+}

73+

59func getPrompts(ctx context.Context, c *mcp.ClientSession) ([]*Prompt, error) {

74func getPrompts(ctx context.Context, c *mcp.ClientSession) ([]*Prompt, error) {

60 if c.InitializeResult().Capabilities.Prompts == nil {

75 if c.InitializeResult().Capabilities.Prompts == nil {

61 return nil, nil

76 return nil, nil

@@ -70,11 +85,11 @@ func getPrompts(ctx context.Context, c *mcp.ClientSession) ([]*Prompt, error) {

70// updatePrompts updates the global mcpPrompts and mcpClient2Prompts maps

85// updatePrompts updates the global mcpPrompts and mcpClient2Prompts maps

71func updatePrompts(mcpName string, prompts []*Prompt) {

86func updatePrompts(mcpName string, prompts []*Prompt) {

72 if len(prompts) == 0 {

87 if len(prompts) == 0 {

73− client2Prompts.Del(mcpName)

88+ clientPrompts.Del(mcpName)

74 } else {

89 } else {

75− client2Prompts.Set(mcpName, prompts)

90+ clientPrompts.Set(mcpName, prompts)

76 }

91 }

77− for mcpName, prompts := range client2Prompts.Seq2() {

92+ for mcpName, prompts := range clientPrompts.Seq2() {

78 for _, p := range prompts {

93 for _, p := range prompts {

79 key := mcpName + ":" + p.Name

94 key := mcpName + ":" + p.Name

80 allPrompts.Set(key, p)

95 allPrompts.Set(key, p)

internal/agent/tools/mcp/tools.go+33−8

@@ -5,6 +5,7 @@ import (

5 "encoding/json"

6 "fmt"

7 "iter"

8+ "log/slog"

8 "strings"

9 "strings"

9

10

10 "github.com/charmbracelet/crush/internal/csync"

11 "github.com/charmbracelet/crush/internal/csync"

@@ -14,12 +15,12 @@ import (

14type Tool = mcp.Tool

15type Tool = mcp.Tool

15

16

16var (

17var (

17− allTools = csync.NewMap[string, *Tool]()

18+ allTools = csync.NewMap[string, *Tool]()

18− client2Tools = csync.NewMap[string, []*Tool]()

19+ clientTools = csync.NewMap[string, []*Tool]()

19)

20)

20

21

21−// GetTools returns all available MCP tools.

22+// Tools returns all available MCP tools.

22−func GetTools() iter.Seq2[string, *Tool] {

23+func Tools() iter.Seq2[string, *Tool] {

23 return allTools.Seq2()

24 return allTools.Seq2()

24}

25}

25

26

@@ -53,6 +54,30 @@ func RunTool(ctx context.Context, name, toolName string, input string) (string,

53 return strings.Join(output, "\n"), nil

54 return strings.Join(output, "\n"), nil

54}

55}

55

56

57+// RefreshTools gets the updated list of tools from the MCP and updates the

58+// global state.

59+func RefreshTools(ctx context.Context, name string) {

60+ session, ok := sessions.Get(name)

61+ if !ok {

62+ slog.Warn("refresh tools: no session", "name", name)

63+ return

64+ }

65+

66+ tools, err := getTools(ctx, session)

67+ if err != nil {

68+ updateState(name, StateError, err, nil, Counts{})

69+ return

70+ }

71+

72+ updateTools(name, tools)

73+

74+ prev, _ := states.Get(name)

75+ updateState(name, StateConnected, nil, session, Counts{

76+ Tools: len(tools),

77+ Prompts: prev.Counts.Prompts,

78+ })

79+}

80+

56func getTools(ctx context.Context, session *mcp.ClientSession) ([]*Tool, error) {

81func getTools(ctx context.Context, session *mcp.ClientSession) ([]*Tool, error) {

57 if session.InitializeResult().Capabilities.Tools == nil {

82 if session.InitializeResult().Capabilities.Tools == nil {

58 return nil, nil

83 return nil, nil

@@ -65,13 +90,13 @@ func getTools(ctx context.Context, session *mcp.ClientSession) ([]*Tool, error)

65}

90}

66

91

67// updateTools updates the global mcpTools and mcpClient2Tools maps

92// updateTools updates the global mcpTools and mcpClient2Tools maps

68−func updateTools(mcpName string, tools []*Tool) {

93+func updateTools(name string, tools []*Tool) {

69 if len(tools) == 0 {

94 if len(tools) == 0 {

70− client2Tools.Del(mcpName)

95+ clientTools.Del(name)

71 } else {

96 } else {

72− client2Tools.Set(mcpName, tools)

97+ clientTools.Set(name, tools)

73 }

98 }

74− for name, tools := range client2Tools.Seq2() {

99+ for name, tools := range clientTools.Seq2() {

75 for _, t := range tools {

100 for _, t := range tools {

76 allTools.Set(name, t)

101 allTools.Set(name, t)

77 }

102 }

internal/app/app.go+1−4

@@ -97,7 +97,7 @@ func New(ctx context.Context, conn *sql.DB, cfg *config.Config) (*App, error) {

97 }()

98

99 // cleanup database upon app shutdown

100− app.cleanupFuncs = append(app.cleanupFuncs, conn.Close)

100+ app.cleanupFuncs = append(app.cleanupFuncs, conn.Close, mcp.Close)

101

102 // TODO: remove the concept of agent config, most likely.

103 if !cfg.IsConfigured() {

@@ -327,9 +327,6 @@ func (app *App) InitCoderAgent(ctx context.Context) error {

327 slog.Error("Failed to create coder agent", "err", err)

328 return err

329 }

330−

331− // Add MCP client cleanup to shutdown process

332− app.cleanupFuncs = append(app.cleanupFuncs, mcp.Close)

333 return nil

330 return nil

334}

331}

335

332

internal/tui/components/dialogs/commands/loader.go+1−1

@@ -223,7 +223,7 @@ type CommandRunCustomMsg struct {

223

224func loadMCPPrompts() []Command {

225 var commands []Command

226− for key, prompt := range mcp.GetPrompts() {

226+ for key, prompt := range mcp.Prompts() {

227 clientName, promptName, ok := strings.Cut(key, ":")

228 if !ok {

229 slog.Warn("prompt not found", "key", key)

internal/tui/tui.go+23

@@ -10,6 +10,7 @@ import (

10

11 "github.com/charmbracelet/bubbles/v2/key"

12 tea "github.com/charmbracelet/bubbletea/v2"

13+ "github.com/charmbracelet/crush/internal/agent/tools/mcp"

13 "github.com/charmbracelet/crush/internal/app"

14 "github.com/charmbracelet/crush/internal/app"

14 "github.com/charmbracelet/crush/internal/config"

15 "github.com/charmbracelet/crush/internal/config"

15 "github.com/charmbracelet/crush/internal/event"

16 "github.com/charmbracelet/crush/internal/event"

@@ -140,6 +141,14 @@ func (a *appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {

140 a.completions.Update(msg)

141 a.completions.Update(msg)

141 return a, a.handleWindowResize(msg.Width, msg.Height)

142 return a, a.handleWindowResize(msg.Width, msg.Height)

142

143

144+ case pubsub.Event[mcp.Event]:

145+ switch msg.Payload.Type {

146+ case mcp.EventPromptsListChanged:

147+

148+ case mcp.EventToolsListChanged:

149+ return a, a.handleMCPToolsEvent(context.Background(), msg.Payload.Name)

150+ }

151+

143 // Completions messages

152 // Completions messages

144 case completions.OpenCompletionsMsg, completions.FilterCompletionsMsg,

153 case completions.OpenCompletionsMsg, completions.FilterCompletionsMsg,

145 completions.CloseCompletionsMsg, completions.RepositionCompletionsMsg:

154 completions.CloseCompletionsMsg, completions.RepositionCompletionsMsg:

@@ -618,6 +627,20 @@ func (a *appModel) View() tea.View {

618 return view

627 return view

619}

628}

620

629

630+func (a *appModel) handleMCPPromptsEvent(ctx context.Context, name string) tea.Cmd {

631+ return func() tea.Msg {

632+ mcp.RefreshPrompts(ctx, name)

633+ return nil

634+ }

635+}

636+

637+func (a *appModel) handleMCPToolsEvent(ctx context.Context, name string) tea.Cmd {

638+ return func() tea.Msg {

639+ mcp.RefreshTools(ctx, name)

640+ return nil

641+ }

642+}

643+

621// New creates and initializes a new TUI application model.

644// New creates and initializes a new TUI application model.

622func New(app *app.App) *appModel {

645func New(app *app.App) *appModel {

623 chatPage := chat.New(app)

646 chatPage := chat.New(app)