git.iko.soy/crush
on main help sign in
refactor(mcp): some more decoupling
54076c0Carlos Alexandro Becker
internal/agent/coordinator.go+1−2
compare with main
unifiedsplit
@@ -18,7 +18,6 @@ import (
18 "github.com/charmbracelet/catwalk/pkg/catwalk"
18 "github.com/charmbracelet/catwalk/pkg/catwalk"
19 "github.com/charmbracelet/crush/internal/agent/prompt"
19 "github.com/charmbracelet/crush/internal/agent/prompt"
20 "github.com/charmbracelet/crush/internal/agent/tools"
20 "github.com/charmbracelet/crush/internal/agent/tools"
21− "github.com/charmbracelet/crush/internal/agent/tools/mcp"
22 "github.com/charmbracelet/crush/internal/config"
21 "github.com/charmbracelet/crush/internal/config"
23 "github.com/charmbracelet/crush/internal/csync"
22 "github.com/charmbracelet/crush/internal/csync"
24 "github.com/charmbracelet/crush/internal/history"
23 "github.com/charmbracelet/crush/internal/history"
@@ -345,7 +344,7 @@ func (c *coordinator) buildTools(ctx context.Context, agent config.Agent) ([]fan
345 }
344 }
346 }
345 }
347
346
348− for tool := range mcp.GetMCPTools() {
347+ for _, tool := range tools.GetMCPTools(c.permissions, c.cfg.WorkingDir()) {
349 if agent.AllowedMCP == nil {
348 if agent.AllowedMCP == nil {
350 // No MCP restrictions
349 // No MCP restrictions
351 filteredTools = append(filteredTools, tool)
350 filteredTools = append(filteredTools, tool)
internal/agent/tools/mcp-tools.gonew file+110
@@ -0,0 +1,110 @@
1+package tools
2+
3+import (
4+ "context"
5+ "fmt"
6+
7+ "charm.land/fantasy"
8+ "github.com/charmbracelet/crush/internal/agent/tools/mcp"
9+ "github.com/charmbracelet/crush/internal/permission"
10+)
11+
12+// GetMCPTools gets all the currently available MCP tools.
13+func GetMCPTools(permissions permission.Service, wd string) []*Tool {
14+ var result []*Tool
15+ for name, tool := range mcp.GetMCPTools() {
16+ result = append(result, &Tool{
17+ mcpName: name,
18+ tool: tool,
19+ permissions: permissions,
20+ workingDir: wd,
21+ })
22+ }
23+ return result
24+}
25+
26+// Tool is a tool from a MCP.
27+type Tool struct {
28+ mcpName string
29+ tool *mcp.Tool
30+ permissions permission.Service
31+ workingDir string
32+ providerOptions fantasy.ProviderOptions
33+}
34+
35+func (m *Tool) SetProviderOptions(opts fantasy.ProviderOptions) {
36+ m.providerOptions = opts
37+}
38+
39+func (m *Tool) ProviderOptions() fantasy.ProviderOptions {
40+ return m.providerOptions
41+}
42+
43+func (m *Tool) Name() string {
44+ return fmt.Sprintf("mcp_%s_%s", m.mcpName, m.tool.Name)
45+}
46+
47+func (m *Tool) MCP() string {
48+ return m.mcpName
49+}
50+
51+func (m *Tool) MCPToolName() string {
52+ return m.tool.Name
53+}
54+
55+func (m *Tool) Info() fantasy.ToolInfo {
56+ parameters := make(map[string]any)
57+ required := make([]string, 0)
58+
59+ if input, ok := m.tool.InputSchema.(map[string]any); ok {
60+ if props, ok := input["properties"].(map[string]any); ok {
61+ parameters = props
62+ }
63+ if req, ok := input["required"].([]any); ok {
64+ // Convert []any -> []string when elements are strings
65+ for _, v := range req {
66+ if s, ok := v.(string); ok {
67+ required = append(required, s)
68+ }
69+ }
70+ } else if reqStr, ok := input["required"].([]string); ok {
71+ // Handle case where it's already []string
72+ required = reqStr
73+ }
74+ }
75+
76+ return fantasy.ToolInfo{
77+ Name: fmt.Sprintf("mcp_%s_%s", m.mcpName, m.tool.Name),
78+ Description: m.tool.Description,
79+ Parameters: parameters,
80+ Required: required,
81+ }
82+}
83+
84+func (m *Tool) Run(ctx context.Context, params fantasy.ToolCall) (fantasy.ToolResponse, error) {
85+ sessionID := GetSessionFromContext(ctx)
86+ if sessionID == "" {
87+ return fantasy.ToolResponse{}, fmt.Errorf("session ID is required for creating a new file")
88+ }
89+ permissionDescription := fmt.Sprintf("execute %s with the following parameters:", m.Info().Name)
90+ p := m.permissions.Request(
91+ permission.CreatePermissionRequest{
92+ SessionID: sessionID,
93+ ToolCallID: params.ID,
94+ Path: m.workingDir,
95+ ToolName: m.Info().Name,
96+ Action: "execute",
97+ Description: permissionDescription,
98+ Params: params.Input,
99+ },
100+ )
101+ if !p {
102+ return fantasy.ToolResponse{}, permission.ErrorPermissionDenied
103+ }
104+
105+ content, err := mcp.RunTool(ctx, m.mcpName, m.tool.Name, params.Input)
106+ if err != nil {
107+ return fantasy.NewTextErrorResponse(err.Error()), nil
108+ }
109+ return fantasy.NewTextResponse(content), nil
110+}
internal/agent/tools/mcp/init.go+3−1
@@ -1,3 +1,5 @@
1+// Package mcp provides functionality for managing Model Context Protocol (MCP)
2+// clients within the Crush application.
1package mcp
3package mcp
2
4
3import (
5import (
@@ -160,7 +162,7 @@ func Initialize(ctx context.Context, permissions permission.Service, cfg *config
160 return
162 return
161 }
163 }
162
164
163− tools, err := getTools(ctx, name, permissions, session, cfg.WorkingDir())
165+ tools, err := getTools(ctx, session)
164 if err != nil {
166 if err != nil {
165 slog.Error("error listing tools", "error", err)
167 slog.Error("error listing tools", "error", err)
166 updateState(name, StateError, err, nil, Counts{})
168 updateState(name, StateError, err, nil, Counts{})
internal/agent/tools/mcp/tools.go+17−107
@@ -7,120 +7,39 @@ import (
7 "iter"
7 "iter"
8 "strings"
8 "strings"
9
9
10− "charm.land/fantasy"
11− "github.com/charmbracelet/crush/internal/agent/tools"
12 "github.com/charmbracelet/crush/internal/csync"
10 "github.com/charmbracelet/crush/internal/csync"
13− "github.com/charmbracelet/crush/internal/permission"
14 "github.com/modelcontextprotocol/go-sdk/mcp"
11 "github.com/modelcontextprotocol/go-sdk/mcp"
15)
12)
16
13
14+type Tool = mcp.Tool
15+
17var (
16var (
18 allTools = csync.NewMap[string, *Tool]()
17 allTools = csync.NewMap[string, *Tool]()
19 client2Tools = csync.NewMap[string, []*Tool]()
18 client2Tools = csync.NewMap[string, []*Tool]()
20)
19)
21
20
22−// GetMCPTools returns all available MCP tools.
21+// GetTools returns all available MCP tools.
23−func GetMCPTools() iter.Seq[*Tool] {
22+func GetTools() iter.Seq2[string, *Tool] {
24− return allTools.Seq()
23+ return allTools.Seq2()
25−}
26−
27−type Tool struct {
28− mcpName string
29− tool *mcp.Tool
30− permissions permission.Service
31− workingDir string
32− providerOptions fantasy.ProviderOptions
33−}
34−
35−func (m *Tool) SetProviderOptions(opts fantasy.ProviderOptions) {
36− m.providerOptions = opts
37−}
38−
39−func (m *Tool) ProviderOptions() fantasy.ProviderOptions {
40− return m.providerOptions
41}
24}
42
25
43−func (m *Tool) Name() string {
26+// RunTool runs an MCP tool with the given input parameters.
44− return fmt.Sprintf("mcp_%s_%s", m.mcpName, m.tool.Name)
45−}
46−
47−func (m *Tool) MCP() string {
48− return m.mcpName
49−}
50−
51−func (m *Tool) MCPToolName() string {
52− return m.tool.Name
53−}
54−
55−func (m *Tool) Info() fantasy.ToolInfo {
56− parameters := make(map[string]any)
57− required := make([]string, 0)
58−
59− if input, ok := m.tool.InputSchema.(map[string]any); ok {
60− if props, ok := input["properties"].(map[string]any); ok {
61− parameters = props
62− }
63− if req, ok := input["required"].([]any); ok {
64− // Convert []any -> []string when elements are strings
65− for _, v := range req {
66− if s, ok := v.(string); ok {
67− required = append(required, s)
68− }
69− }
70− } else if reqStr, ok := input["required"].([]string); ok {
71− // Handle case where it's already []string
72− required = reqStr
73− }
74− }
75−
76− return fantasy.ToolInfo{
77− Name: fmt.Sprintf("mcp_%s_%s", m.mcpName, m.tool.Name),
78− Description: m.tool.Description,
79− Parameters: parameters,
80− Required: required,
81− }
82−}
83−
84−func (m *Tool) Run(ctx context.Context, params fantasy.ToolCall) (fantasy.ToolResponse, error) {
85− sessionID := tools.GetSessionFromContext(ctx)
86− if sessionID == "" {
87− return fantasy.ToolResponse{}, fmt.Errorf("session ID is required for creating a new file")
88− }
89− permissionDescription := fmt.Sprintf("execute %s with the following parameters:", m.Info().Name)
90− p := m.permissions.Request(
91− permission.CreatePermissionRequest{
92− SessionID: sessionID,
93− ToolCallID: params.ID,
94− Path: m.workingDir,
95− ToolName: m.Info().Name,
96− Action: "execute",
97− Description: permissionDescription,
98− Params: params.Input,
99− },
100− )
101− if !p {
102− return fantasy.ToolResponse{}, permission.ErrorPermissionDenied
103− }
104−
105− return runTool(ctx, m.mcpName, m.tool.Name, params.Input)
106−}
107−
108−func runTool(ctx context.Context, name, toolName string, input string) (fantasy.ToolResponse, error) {
27+func RunTool(ctx context.Context, name, toolName string, input string) (string, error) {
109 var args map[string]any
28 var args map[string]any
110 if err := json.Unmarshal([]byte(input), &args); err != nil {
29 if err := json.Unmarshal([]byte(input), &args); err != nil {
111− return fantasy.NewTextErrorResponse(fmt.Sprintf("error parsing parameters: %s", err)), nil
30+ return "", fmt.Errorf("error parsing parameters: %s", err)
112 }
31 }
113
32
114 c, err := getOrRenewClient(ctx, name)
33 c, err := getOrRenewClient(ctx, name)
115 if err != nil {
34 if err != nil {
116− return fantasy.NewTextErrorResponse(err.Error()), nil
35+ return "", err
117 }
36 }
118 result, err := c.CallTool(ctx, &mcp.CallToolParams{
37 result, err := c.CallTool(ctx, &mcp.CallToolParams{
119 Name: toolName,
38 Name: toolName,
120 Arguments: args,
39 Arguments: args,
121 })
40 })
122 if err != nil {
41 if err != nil {
123− return fantasy.NewTextErrorResponse(err.Error()), nil
42+ return "", err
124 }
43 }
125
44
126 output := make([]string, 0, len(result.Content))
45 output := make([]string, 0, len(result.Content))
@@ -131,27 +50,18 @@ func runTool(ctx context.Context, name, toolName string, input string) (fantasy.
131 output = append(output, fmt.Sprintf("%v", v))
50 output = append(output, fmt.Sprintf("%v", v))
132 }
51 }
133 }
52 }
134− return fantasy.NewTextResponse(strings.Join(output, "\n")), nil
53+ return strings.Join(output, "\n"), nil
135}
54}
136
55
137−func getTools(ctx context.Context, name string, permissions permission.Service, c *mcp.ClientSession, workingDir string) ([]*Tool, error) {
56+func getTools(ctx context.Context, session *mcp.ClientSession) ([]*Tool, error) {
138− if c.InitializeResult().Capabilities.Tools == nil {
57+ if session.InitializeResult().Capabilities.Tools == nil {
139 return nil, nil
58 return nil, nil
140 }
59 }
141− result, err := c.ListTools(ctx, &mcp.ListToolsParams{})
60+ result, err := session.ListTools(ctx, &mcp.ListToolsParams{})
142 if err != nil {
61 if err != nil {
143 return nil, err
62 return nil, err
144 }
63 }
145− mcpTools := make([]*Tool, 0, len(result.Tools))
146− for _, tool := range result.Tools {
147− mcpTools = append(mcpTools, &Tool{
148− mcpName: name,
149− tool: tool,
150− permissions: permissions,
151− workingDir: workingDir,
152− })
153− }
154− return mcpTools, nil
64+ return result.Tools, nil
155}
65}
156
66
157// updateTools updates the global mcpTools and mcpClient2Tools maps
67// updateTools updates the global mcpTools and mcpClient2Tools maps
@@ -161,9 +71,9 @@ func updateTools(mcpName string, tools []*Tool) {
161 } else {
71 } else {
162 client2Tools.Set(mcpName, tools)
72 client2Tools.Set(mcpName, tools)
163 }
73 }
164− for _, tools := range client2Tools.Seq2() {
74+ for name, tools := range client2Tools.Seq2() {
165 for _, t := range tools {
75 for _, t := range tools {
166− allTools.Set(t.Name(), t)
76+ allTools.Set(name, t)
167 }
77 }
168 }
78 }
169}
79}