9423d10 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

Merge pull request #208 from charmbracelet/get-tools-mu

9423d10Kujtim Hoxha1 year ago

cmd/root.go−20

compare with main

unifiedsplit

@@ -6,14 +6,11 @@ import (

6 "io"

7 "log/slog"

8 "os"

9− "time"

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

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

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

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

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

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

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

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

15− "github.com/charmbracelet/crush/internal/llm/agent"

16− "github.com/charmbracelet/crush/internal/log"

17 "github.com/charmbracelet/crush/internal/tui"

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

18 "github.com/charmbracelet/crush/internal/version"

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

19 "github.com/charmbracelet/fang"

16 "github.com/charmbracelet/fang"

@@ -92,9 +89,6 @@ to assist developers in writing, debugging, and understanding code directly from

92 }

89 }

93 defer app.Shutdown()

90 defer app.Shutdown()

95− // Initialize MCP tools early for both modes

96− initMCPTools(ctx, app, cfg)

97−

98 prompt, err = maybePrependStdin(prompt)

92 prompt, err = maybePrependStdin(prompt)

99 if err != nil {

93 if err != nil {

100 slog.Error(fmt.Sprintf("Failed to read from stdin: %v", err))

94 slog.Error(fmt.Sprintf("Failed to read from stdin: %v", err))

@@ -126,20 +120,6 @@ to assist developers in writing, debugging, and understanding code directly from

126 },

120 },

127}

121}

128

122

129−func initMCPTools(ctx context.Context, app *app.App, cfg *config.Config) {

130− go func() {

131− defer log.RecoverPanic("MCP-goroutine", nil)

132−

133− // Create a context with timeout for the initial MCP tools fetch

134− ctxWithTimeout, cancel := context.WithTimeout(ctx, 30*time.Second)

135− defer cancel()

136−

137− // Set this up once with proper error handling

138− agent.GetMcpTools(ctxWithTimeout, app.Permissions, cfg)

139− slog.Info("MCP message handling goroutine exiting")

140− }()

141−}

142−

143func Execute() {

123func Execute() {

144 if err := fang.Execute(

124 if err := fang.Execute(

145 context.Background(),

125 context.Background(),

internal/llm/agent/agent.go+1−1

@@ -94,7 +94,7 @@ func NewAgent(

94) (Service, error) {

95 ctx := context.Background()

96 cfg := config.Get()

97− otherTools := GetMcpTools(ctx, permissions, cfg)

97+ otherTools := GetMCPTools(ctx, permissions, cfg)

98 if len(lspClients) > 0 {

99 otherTools = append(otherTools, tools.NewDiagnosticsTool(lspClients))

100 }

internal/llm/agent/mcp-tools.go+62−41

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

5 "encoding/json"

6 "fmt"

7 "log/slog"

8+ "sync"

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

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

10 "github.com/charmbracelet/crush/internal/llm/tools"

11 "github.com/charmbracelet/crush/internal/llm/tools"

@@ -154,8 +155,6 @@ func NewMcpTool(name string, tool mcp.Tool, permissions permission.Service, mcpC

154 }

155 }

155}

156}

156

157

157−var mcpTools []tools.BaseTool

158−

159func getTools(ctx context.Context, name string, m config.MCPConfig, permissions permission.Service, c MCPClient, workingDir string) []tools.BaseTool {

158func getTools(ctx context.Context, name string, m config.MCPConfig, permissions permission.Service, c MCPClient, workingDir string) []tools.BaseTool {

160 var stdioTools []tools.BaseTool

159 var stdioTools []tools.BaseTool

161 initRequest := mcp.InitializeRequest{}

160 initRequest := mcp.InitializeRequest{}

@@ -183,50 +182,72 @@ func getTools(ctx context.Context, name string, m config.MCPConfig, permissions

183 return stdioTools

182 return stdioTools

184}

183}

185

184

185+var (

186+ mcpToolsOnce sync.Once

187+ mcpTools []tools.BaseTool

188+)

189+

186−func GetMcpTools(ctx context.Context, permissions permission.Service, cfg *config.Config) []tools.BaseTool {

190+func GetMCPTools(ctx context.Context, permissions permission.Service, cfg *config.Config) []tools.BaseTool {

187− if len(mcpTools) > 0 {

191+ mcpToolsOnce.Do(func() {

192+ mcpTools = doGetMCPTools(ctx, permissions, cfg)

193+ })

188− return mcpTools

194+ return mcpTools

189− }

195+}

196+

197+func doGetMCPTools(ctx context.Context, permissions permission.Service, cfg *config.Config) []tools.BaseTool {

198+ var mu sync.Mutex

199+ var wg sync.WaitGroup

200+ var result []tools.BaseTool

190 for name, m := range cfg.MCP {

201 for name, m := range cfg.MCP {

191 if m.Disabled {

202 if m.Disabled {

192 slog.Debug("skipping disabled mcp", "name", name)

203 slog.Debug("skipping disabled mcp", "name", name)

193 continue

204 continue

194 }

205 }

206+ wg.Add(1)

207+ go func(name string, m config.MCPConfig) {

208+ defer wg.Done()

195− switch m.Type {

209+ switch m.Type {

196− case config.MCPStdio:

210+ case config.MCPStdio:

197− c, err := client.NewStdioMCPClient(

211+ c, err := client.NewStdioMCPClient(

198− m.Command,

212+ m.Command,

199− m.Env,

213+ m.Env,

200− m.Args...,

214+ m.Args...,

201− )

215+ )

216+ if err != nil {

217+ slog.Error("error creating mcp client", "error", err)

218+ return

219+ }

220+

221+ mu.Lock()

222+ result = append(result, getTools(ctx, name, m, permissions, c, cfg.WorkingDir())...)

223+ mu.Unlock()

224+ case config.MCPHttp:

225+ c, err := client.NewStreamableHttpClient(

226+ m.URL,

227+ transport.WithHTTPHeaders(m.Headers),

228+ )

229+ if err != nil {

230+ slog.Error("error creating mcp client", "error", err)

231+ return

232+ }

233+ mu.Lock()

234+ result = append(result, getTools(ctx, name, m, permissions, c, cfg.WorkingDir())...)

235+ mu.Unlock()

236+ case config.MCPSse:

237+ c, err := client.NewSSEMCPClient(

238+ m.URL,

239+ client.WithHeaders(m.Headers),

240+ )

202− if err != nil {

241+ if err != nil {

203− slog.Error("error creating mcp client", "error", err)

242+ slog.Error("error creating mcp client", "error", err)

204− continue

243+ return

244+ }

245+ mu.Lock()

246+ result = append(result, getTools(ctx, name, m, permissions, c, cfg.WorkingDir())...)

247+ mu.Unlock()

205 }

248 }

206−

249+ }(name, m)

207− mcpTools = append(mcpTools, getTools(ctx, name, m, permissions, c, cfg.WorkingDir())...)

208− case config.MCPHttp:

209− c, err := client.NewStreamableHttpClient(

210− m.URL,

211− transport.WithHTTPHeaders(m.Headers),

212− )

213− if err != nil {

214− slog.Error("error creating mcp client", "error", err)

215− continue

216− }

217− mcpTools = append(mcpTools, getTools(ctx, name, m, permissions, c, cfg.WorkingDir())...)

218− case config.MCPSse:

219− c, err := client.NewSSEMCPClient(

220− m.URL,

221− client.WithHeaders(m.Headers),

222− )

223− if err != nil {

224− slog.Error("error creating mcp client", "error", err)

225− continue

226− }

227− mcpTools = append(mcpTools, getTools(ctx, name, m, permissions, c, cfg.WorkingDir())...)

228− }

229 }

250 }

230−

251+ wg.Wait()

231− return mcpTools

252+ return result

232}

253}