8c7c0db · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

feat(mcp): notifications support - tools/list_changed (#967)

8c7c0db林玮 (Jade Lin)1 year ago

internal/csync/maps.go+19

compare with main

unifiedsplit

@@ -27,6 +27,25 @@ func NewMapFrom[K comparable, V any](m map[K]V) *Map[K, V] {

27 }

28}

30+// NewLazyMap creates a new lazy-loaded map. The provided load function is

31+// executed in a separate goroutine to populate the map.

32+func NewLazyMap[K comparable, V any](load func() map[K]V) *Map[K, V] {

33+ m := &Map[K, V]{}

34+ m.mu.Lock()

35+ go func() {

36+ m.inner = load()

37+ m.mu.Unlock()

38+ }()

39+ return m

40+}

41+

42+// Reset replaces the inner map with the new one.

43+func (m *Map[K, V]) Reset(input map[K]V) {

44+ m.mu.Lock()

45+ defer m.mu.Unlock()

46+ m.inner = input

47+}

48+

30// Set sets the value for the specified key in the map.

49// Set sets the value for the specified key in the map.

31func (m *Map[K, V]) Set(key K, value V) {

50func (m *Map[K, V]) Set(key K, value V) {

32 m.mu.Lock()

51 m.mu.Lock()

internal/csync/maps_test.go+52

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

5 "maps"

6 "sync"

7 "testing"

8+ "testing/synctest"

9+ "time"

9 "github.com/stretchr/testify/require"

11 "github.com/stretchr/testify/require"

10)

12)

@@ -36,6 +38,56 @@ func TestNewMapFrom(t *testing.T) {

36 require.Equal(t, 1, value)

38 require.Equal(t, 1, value)

37}

39}

41+func TestNewLazyMap(t *testing.T) {

42+ t.Parallel()

43+

44+ synctest.Test(t, func(t *testing.T) {

45+ t.Helper()

46+

47+ waiter := sync.Mutex{}

48+ waiter.Lock()

49+ loadCalled := false

50+

51+ loadFunc := func() map[string]int {

52+ waiter.Lock()

53+ defer waiter.Unlock()

54+ loadCalled = true

55+ return map[string]int{

56+ "key1": 1,

57+ "key2": 2,

58+ }

59+ }

60+

61+ m := NewLazyMap(loadFunc)

62+ require.NotNil(t, m)

63+

64+ waiter.Unlock() // Allow the load function to proceed

65+ time.Sleep(100 * time.Millisecond)

66+ require.True(t, loadCalled)

67+ require.Equal(t, 2, m.Len())

68+

69+ value, ok := m.Get("key1")

70+ require.True(t, ok)

71+ require.Equal(t, 1, value)

72+ })

73+}

74+

75+func TestMap_Reset(t *testing.T) {

76+ t.Parallel()

77+

78+ m := NewMapFrom(map[string]int{

79+ "a": 10,

80+ })

81+

82+ m.Reset(map[string]int{

83+ "b": 20,

84+ })

85+ value, ok := m.Get("b")

86+ require.True(t, ok)

87+ require.Equal(t, 20, value)

88+ require.Equal(t, 1, m.Len())

89+}

90+

39func TestMap_Set(t *testing.T) {

91func TestMap_Set(t *testing.T) {

40 t.Parallel()

92 t.Parallel()

internal/llm/agent/agent.go+94−34

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

5 "errors"

6 "fmt"

7 "log/slog"

8+ "maps"

8 "slices"

9 "slices"

9 "strings"

10 "strings"

10 "time"

11 "time"

@@ -65,11 +66,13 @@ type agent struct {

65 sessions session.Service

66 sessions session.Service

66 messages message.Service

67 messages message.Service

67 permissions permission.Service

68 permissions permission.Service

68− mcpTools []McpTool

69+ baseTools *csync.Map[string, tools.BaseTool]

70+ mcpTools *csync.Map[string, tools.BaseTool]

71+ lspClients *csync.Map[string, *lsp.Client]

70− tools *csync.LazySlice[tools.BaseTool]

71 // We need this to be able to update it when model changes

73 // We need this to be able to update it when model changes

72− agentToolFn func() (tools.BaseTool, error)

74+ agentToolFn func() (tools.BaseTool, error)

75+ cleanupFuncs []func()

74 provider provider.Provider

77 provider provider.Provider

75 providerID string

78 providerID string

@@ -171,14 +174,16 @@ func NewAgent(

171 return nil, err

174 return nil, err

172 }

175 }

173

176

174− toolFn := func() []tools.BaseTool {

177+ baseToolsFn := func() map[string]tools.BaseTool {

175− slog.Info("Initializing agent tools", "agent", agentCfg.ID)

178+ slog.Info("Initializing agent base tools", "agent", agentCfg.ID)

176 defer func() {

179 defer func() {

177− slog.Info("Initialized agent tools", "agent", agentCfg.ID)

180+ slog.Info("Initialized agent base tools", "agent", agentCfg.ID)

178 }()

181 }()

179

182

183+ // Base tools available to all agents

180 cwd := cfg.WorkingDir()

184 cwd := cfg.WorkingDir()

185+ result := make(map[string]tools.BaseTool)

181− allTools := []tools.BaseTool{

186+ for _, tool := range []tools.BaseTool{

182 tools.NewBashTool(permissions, cwd, cfg.Options.Attribution),

187 tools.NewBashTool(permissions, cwd, cfg.Options.Attribution),

183 tools.NewDownloadTool(permissions, cwd),

188 tools.NewDownloadTool(permissions, cwd),

184 tools.NewEditTool(lspClients, permissions, history, cwd),

189 tools.NewEditTool(lspClients, permissions, history, cwd),

@@ -190,36 +195,25 @@ func NewAgent(

190 tools.NewSourcegraphTool(),

195 tools.NewSourcegraphTool(),

191 tools.NewViewTool(lspClients, permissions, cwd),

196 tools.NewViewTool(lspClients, permissions, cwd),

192 tools.NewWriteTool(lspClients, permissions, history, cwd),

197 tools.NewWriteTool(lspClients, permissions, history, cwd),

198+ } {

199+ result[tool.Name()] = tool

193 }

200 }

201+ return result

202+ }

203+ mcpToolsFn := func() map[string]tools.BaseTool {

204+ slog.Info("Initializing agent mcp tools", "agent", agentCfg.ID)

205+ defer func() {

206+ slog.Info("Initialized agent mcp tools", "agent", agentCfg.ID)

207+ }()

194

208

195 mcpToolsOnce.Do(func() {

209 mcpToolsOnce.Do(func() {

196− mcpTools = doGetMCPTools(ctx, permissions, cfg)

210+ doGetMCPTools(ctx, permissions, cfg)

197 })

211 })

198

212

199− withCoderTools := func(t []tools.BaseTool) []tools.BaseTool {

213+ return maps.Collect(mcpTools.Seq2())

200− if agentCfg.ID == "coder" {

201− t = append(t, mcpTools...)

202− if lspClients.Len() > 0 {

203− t = append(t, tools.NewDiagnosticsTool(lspClients))

204− }

205− }

206− return t

207− }

208−

209− if agentCfg.AllowedTools == nil {

210− return withCoderTools(allTools)

211− }

212−

213− var filteredTools []tools.BaseTool

214− for _, tool := range allTools {

215− if slices.Contains(agentCfg.AllowedTools, tool.Name()) {

216− filteredTools = append(filteredTools, tool)

217− }

218− }

219− return withCoderTools(filteredTools)

220 }

214 }

221

215

222− return &agent{

216+ a := &agent{

223 Broker: pubsub.NewBroker[AgentEvent](),

217 Broker: pubsub.NewBroker[AgentEvent](),

224 agentCfg: agentCfg,

218 agentCfg: agentCfg,

225 provider: agentProvider,

219 provider: agentProvider,

@@ -231,10 +225,14 @@ func NewAgent(

231 summarizeProviderID: string(providerCfg.ID),

225 summarizeProviderID: string(providerCfg.ID),

232 agentToolFn: agentToolFn,

226 agentToolFn: agentToolFn,

233 activeRequests: csync.NewMap[string, context.CancelFunc](),

227 activeRequests: csync.NewMap[string, context.CancelFunc](),

234− tools: csync.NewLazySlice(toolFn),

228+ mcpTools: csync.NewLazyMap(mcpToolsFn),

229+ baseTools: csync.NewLazyMap(baseToolsFn),

235 promptQueue: csync.NewMap[string, []string](),

230 promptQueue: csync.NewMap[string, []string](),

236 permissions: permissions,

231 permissions: permissions,

232+ lspClients: lspClients,

233+ }

234+ a.setupEvents(ctx)

237− }, nil

235+ return a, nil

238}

236}

239

237

240func (a *agent) Model() catwalk.Model {

238func (a *agent) Model() catwalk.Model {

@@ -517,7 +515,18 @@ func (a *agent) createUserMessage(ctx context.Context, sessionID, content string

517}

515}

518

516

519func (a *agent) getAllTools() ([]tools.BaseTool, error) {

517func (a *agent) getAllTools() ([]tools.BaseTool, error) {

520− allTools := slices.Collect(a.tools.Seq())

518+ var allTools []tools.BaseTool

519+ for tool := range a.baseTools.Seq() {

520+ if a.agentCfg.AllowedTools == nil || slices.Contains(a.agentCfg.AllowedTools, tool.Name()) {

521+ allTools = append(allTools, tool)

522+ }

523+ }

524+ if a.agentCfg.ID == "coder" {

525+ allTools = slices.AppendSeq(allTools, a.mcpTools.Seq())

526+ if a.lspClients.Len() > 0 {

527+ allTools = append(allTools, tools.NewDiagnosticsTool(a.lspClients))

528+ }

529+ }

521 if a.agentToolFn != nil {

530 if a.agentToolFn != nil {

522 agentTool, agentToolErr := a.agentToolFn()

531 agentTool, agentToolErr := a.agentToolFn()

523 if agentToolErr != nil {

532 if agentToolErr != nil {

@@ -591,7 +600,7 @@ loop:

591 default:

600 default:

592 // Continue processing

601 // Continue processing

593 var tool tools.BaseTool

602 var tool tools.BaseTool

594− allTools, _ := a.getAllTools()

603+ allTools, _ = a.getAllTools()

595 for _, availableTool := range allTools {

604 for _, availableTool := range allTools {

596 if availableTool.Info().Name == toolCall.Name {

605 if availableTool.Info().Name == toolCall.Name {

597 tool = availableTool

606 tool = availableTool

@@ -960,6 +969,12 @@ func (a *agent) CancelAll() {

960 a.Cancel(key) // key is sessionID

969 a.Cancel(key) // key is sessionID

961 }

970 }

962

971

972+ for _, cleanup := range a.cleanupFuncs {

973+ if cleanup != nil {

974+ cleanup()

975+ }

976+ }

977+

963 timeout := time.After(5 * time.Second)

978 timeout := time.After(5 * time.Second)

964 for a.IsBusy() {

979 for a.IsBusy() {

965 select {

980 select {

@@ -1071,3 +1086,48 @@ func (a *agent) UpdateModel() error {

1071

1086

1072 return nil

1087 return nil

1073}

1088}

1089+

1090+func (a *agent) setupEvents(ctx context.Context) {

1091+ ctx, cancel := context.WithCancel(ctx)

1092+

1093+ go func() {

1094+ subCh := SubscribeMCPEvents(ctx)

1095+

1096+ for {

1097+ select {

1098+ case event, ok := <-subCh:

1099+ if !ok {

1100+ slog.Debug("MCPEvents subscription channel closed")

1101+ return

1102+ }

1103+ switch event.Payload.Type {

1104+ case MCPEventToolsListChanged:

1105+ name := event.Payload.Name

1106+ c, ok := mcpClients.Get(name)

1107+ if !ok {

1108+ slog.Warn("MCP client not found for tools update", "name", name)

1109+ continue

1110+ }

1111+ cfg := config.Get()

1112+ tools, err := getTools(ctx, name, a.permissions, c, cfg.WorkingDir())

1113+ if err != nil {

1114+ slog.Error("error listing tools", "error", err)

1115+ updateMCPState(name, MCPStateError, err, nil, 0)

1116+ _ = c.Close()

1117+ continue

1118+ }

1119+ updateMcpTools(name, tools)

1120+ a.mcpTools.Reset(maps.Collect(mcpTools.Seq2()))

1121+ updateMCPState(name, MCPStateConnected, nil, c, a.mcpTools.Len())

1122+ default:

1123+ continue

1124+ }

1125+ case <-ctx.Done():

1126+ slog.Debug("MCPEvents subscription cancelled")

1127+ return

1128+ }

1129+ }

1130+ }()

1131+

1132+ a.cleanupFuncs = append(a.cleanupFuncs, cancel)

1133+}

internal/llm/agent/mcp-tools.go+52−13

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

8 "fmt"

9 "log/slog"

10 "maps"

11− "slices"

12 "strings"

11 "strings"

13 "sync"

12 "sync"

14 "time"

13 "time"

@@ -54,7 +53,8 @@ func (s MCPState) String() string {

54type MCPEventType string

53type MCPEventType string

56const (

55const (

57− MCPEventStateChanged MCPEventType = "state_changed"

56+ MCPEventStateChanged MCPEventType = "state_changed"

57+ MCPEventToolsListChanged MCPEventType = "tools_list_changed"

58)

60// MCPEvent represents an event in the MCP system

@@ -77,11 +77,12 @@ type MCPClientInfo struct {

77}

79var (

80− mcpToolsOnce sync.Once

80+ mcpToolsOnce sync.Once

81− mcpTools []tools.BaseTool

81+ mcpTools = csync.NewMap[string, tools.BaseTool]()

82+ mcpClient2Tools = csync.NewMap[string, []tools.BaseTool]()

82− mcpClients = csync.NewMap[string, *client.Client]()

83+ mcpClients = csync.NewMap[string, *client.Client]()

83− mcpStates = csync.NewMap[string, MCPClientInfo]()

84+ mcpStates = csync.NewMap[string, MCPClientInfo]()

84− mcpBroker = pubsub.NewBroker[MCPEvent]()

85+ mcpBroker = pubsub.NewBroker[MCPEvent]()

85)

86)

87type McpTool struct {

88type McpTool struct {

@@ -237,8 +238,12 @@ func updateMCPState(name string, state MCPState, err error, client *client.Clien

237 Client: client,

238 Client: client,

238 ToolCount: toolCount,

239 ToolCount: toolCount,

239 }

240 }

240− if state == MCPStateConnected {

241+ switch state {

242+ case MCPStateConnected:

241 info.ConnectedAt = time.Now()

243 info.ConnectedAt = time.Now()

244+ case MCPStateError:

245+ updateMcpTools(name, nil)

246+ mcpClients.Del(name)

242 }

247 }

243 mcpStates.Set(name, info)

248 mcpStates.Set(name, info)

244

249

@@ -252,6 +257,14 @@ func updateMCPState(name string, state MCPState, err error, client *client.Clien

252 })

257 })

253}

258}

254

259

260+// publishMCPEventToolsListChanged publishes a tool list changed event

261+func publishMCPEventToolsListChanged(name string) {

262+ mcpBroker.Publish(pubsub.UpdatedEvent, MCPEvent{

263+ Type: MCPEventToolsListChanged,

264+ Name: name,

265+ })

266+}

267+

255// CloseMCPClients closes all MCP clients. This should be called during application shutdown.

268// CloseMCPClients closes all MCP clients. This should be called during application shutdown.

256func CloseMCPClients() error {

269func CloseMCPClients() error {

257 var errs []error

270 var errs []error

@@ -274,10 +287,8 @@ var mcpInitRequest = mcp.InitializeRequest{

274 },

287 },

275}

288}

276

289

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

290+func doGetMCPTools(ctx context.Context, permissions permission.Service, cfg *config.Config) {

278 var wg sync.WaitGroup

291 var wg sync.WaitGroup

279− result := csync.NewSlice[tools.BaseTool]()

280−

281 // Initialize states for all configured MCPs

292 // Initialize states for all configured MCPs

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

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

283 if m.Disabled {

294 if m.Disabled {

@@ -316,6 +327,8 @@ func doGetMCPTools(ctx context.Context, permissions permission.Service, cfg *con

316 return

327 return

317 }

328 }

318

329

330+ mcpClients.Set(name, c)

331+

319 tools, err := getTools(ctx, name, permissions, c, cfg.WorkingDir())

332 tools, err := getTools(ctx, name, permissions, c, cfg.WorkingDir())

320 if err != nil {

333 if err != nil {

321 slog.Error("error listing tools", "error", err)

334 slog.Error("error listing tools", "error", err)

@@ -324,13 +337,26 @@ func doGetMCPTools(ctx context.Context, permissions permission.Service, cfg *con

324 return

337 return

325 }

338 }

326

339

340+ updateMcpTools(name, tools)

327 mcpClients.Set(name, c)

341 mcpClients.Set(name, c)

328 updateMCPState(name, MCPStateConnected, nil, c, len(tools))

342 updateMCPState(name, MCPStateConnected, nil, c, len(tools))

329− result.Append(tools...)

330 }(name, m)

343 }(name, m)

331 }

344 }

332 wg.Wait()

345 wg.Wait()

333− return slices.Collect(result.Seq())

346+}

347+

348+// updateMcpTools updates the global mcpTools and mcpClient2Tools maps

349+func updateMcpTools(mcpName string, tools []tools.BaseTool) {

350+ if len(tools) == 0 {

351+ mcpClient2Tools.Del(mcpName)

352+ } else {

353+ mcpClient2Tools.Set(mcpName, tools)

354+ }

355+ for _, tools := range mcpClient2Tools.Seq2() {

356+ for _, t := range tools {

357+ mcpTools.Set(t.Name(), t)

358+ }

359+ }

334}

360}

335

361

336func createAndInitializeClient(ctx context.Context, name string, m config.MCPConfig, resolver config.VariableResolver) (*client.Client, error) {

362func createAndInitializeClient(ctx context.Context, name string, m config.MCPConfig, resolver config.VariableResolver) (*client.Client, error) {

@@ -341,11 +367,22 @@ func createAndInitializeClient(ctx context.Context, name string, m config.MCPCon

341 return nil, err

367 return nil, err

342 }

368 }

343

369

370+ c.OnNotification(func(n mcp.JSONRPCNotification) {

371+ slog.Debug("Received MCP notification", "name", name, "notification", n)

372+ switch n.Method {

373+ case "notifications/tools/list_changed":

374+ publishMCPEventToolsListChanged(name)

375+ default:

376+ slog.Debug("Unhandled MCP notification", "name", name, "method", n.Method)

377+ }

378+ })

379+

344 // XXX: ideally we should be able to use context.WithTimeout here, but,

380 // XXX: ideally we should be able to use context.WithTimeout here, but,

345 // the SSE MCP client will start failing once that context is canceled.

381 // the SSE MCP client will start failing once that context is canceled.

346 timeout := mcpTimeout(m)

382 timeout := mcpTimeout(m)

347 mcpCtx, cancel := context.WithCancel(ctx)

383 mcpCtx, cancel := context.WithCancel(ctx)

348 cancelTimer := time.AfterFunc(timeout, cancel)

384 cancelTimer := time.AfterFunc(timeout, cancel)

385+

349 if err := c.Start(mcpCtx); err != nil {

386 if err := c.Start(mcpCtx); err != nil {

350 updateMCPState(name, MCPStateError, maybeTimeoutErr(err, timeout), nil, 0)

387 updateMCPState(name, MCPStateError, maybeTimeoutErr(err, timeout), nil, 0)

351 slog.Error("error starting mcp client", "error", err, "name", name)

388 slog.Error("error starting mcp client", "error", err, "name", name)

@@ -353,6 +390,7 @@ func createAndInitializeClient(ctx context.Context, name string, m config.MCPCon

353 cancel()

390 cancel()

354 return nil, err

391 return nil, err

355 }

392 }

393+

356 if _, err := c.Initialize(mcpCtx, mcpInitRequest); err != nil {

394 if _, err := c.Initialize(mcpCtx, mcpInitRequest); err != nil {

357 updateMCPState(name, MCPStateError, maybeTimeoutErr(err, timeout), nil, 0)

395 updateMCPState(name, MCPStateError, maybeTimeoutErr(err, timeout), nil, 0)

358 slog.Error("error initializing mcp client", "error", err, "name", name)

396 slog.Error("error initializing mcp client", "error", err, "name", name)

@@ -360,6 +398,7 @@ func createAndInitializeClient(ctx context.Context, name string, m config.MCPCon

360 cancel()

398 cancel()

361 return nil, err

399 return nil, err

362 }

400 }

401+

363 cancelTimer.Stop()

402 cancelTimer.Stop()

364 slog.Info("Initialized mcp client", "name", name)

403 slog.Info("Initialized mcp client", "name", name)

365 return c, nil

404 return c, nil