8c7c0db · crush · git.iko.soy

git.iko.soy/crush

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

Signed-off-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com> Co-authored-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com>

commit8c7c0db22606db51910c2a24dc15319369ce60f6 author林玮 (Jade Lin) <linw1995@icloud.com> dateWed 1 Oct 2025, 20:03 (1 year ago) parentc2caa3c

4 files changed+217−47

compare with main

unifiedsplit

internal/csync/maps.go+19 internal/csync/maps_test.go+52 internal/llm/agent/agent.go+94−34 internal/llm/agent/mcp-tools.go+52−13

internal/csync/maps.go+19

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

2727 }

2828}

2929

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+

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

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

3251 m.mu.Lock()

internal/csync/maps_test.go+52

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

55 "maps"

66 "sync"

77 "testing"

8+ "testing/synctest"

9+ "time"

810

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

1012)

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

3638 require.Equal(t, 1, value)

3739}

3840

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+

3991func TestMap_Set(t *testing.T) {

4092 t.Parallel()

4193

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

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

55 "errors"

66 "fmt"

77 "log/slog"

8+ "maps"

89 "slices"

910 "strings"

1011 "time"

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

6566 sessions session.Service

6667 messages message.Service

6768 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]

6972

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

7173 // 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()

7376

7477 provider provider.Provider

7578 providerID string

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

171174 return nil, err

172175 }

173176

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

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

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

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

176179 defer func() {

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

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

178181 }()

179182

183+ // Base tools available to all agents

180184 cwd := cfg.WorkingDir()

181− allTools := []tools.BaseTool{

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

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

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

183188 tools.NewDownloadTool(permissions, cwd),

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

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

190195 tools.NewSourcegraphTool(),

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

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

198+ } {

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

193200 }

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+ }()

194208

195209 mcpToolsOnce.Do(func() {

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

210+ doGetMCPTools(ctx, permissions, cfg)

197211 })

198212

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

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)

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

220214 }

221215

222− return &agent{

216+ a := &agent{

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

224218 agentCfg: agentCfg,

225219 provider: agentProvider,

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

231225 summarizeProviderID: string(providerCfg.ID),

232226 agentToolFn: agentToolFn,

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

234− tools: csync.NewLazySlice(toolFn),

228+ mcpTools: csync.NewLazyMap(mcpToolsFn),

229+ baseTools: csync.NewLazyMap(baseToolsFn),

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

236231 permissions: permissions,

237− }, nil

232+ lspClients: lspClients,

233+ }

234+ a.setupEvents(ctx)

235+ return a, nil

238236}

239237

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

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

517515}

518516

519517func (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+ }

521530 if a.agentToolFn != nil {

522531 agentTool, agentToolErr := a.agentToolFn()

523532 if agentToolErr != nil {

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

591600 default:

592601 // Continue processing

593602 var tool tools.BaseTool

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

603+ allTools, _ = a.getAllTools()

595604 for _, availableTool := range allTools {

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

597606 tool = availableTool

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

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

961970 }

962971

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

973+ if cleanup != nil {

974+ cleanup()

975+ }

976+ }

977+

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

964979 for a.IsBusy() {

965980 select {

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

10711086

10721087 return nil

10731088}

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 (

88 "fmt"

99 "log/slog"

1010 "maps"

11− "slices"

1211 "strings"

1312 "sync"

1413 "time"

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

5453type MCPEventType string

5554

5655const (

57− MCPEventStateChanged MCPEventType = "state_changed"

56+ MCPEventStateChanged MCPEventType = "state_changed"

57+ MCPEventToolsListChanged MCPEventType = "tools_list_changed"

5858)

5959

6060// MCPEvent represents an event in the MCP system

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

7777}

7878

7979var (

80− mcpToolsOnce sync.Once

81− mcpTools []tools.BaseTool

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

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

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

80+ mcpToolsOnce sync.Once

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

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

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

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

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

8586)

8687

8788type McpTool struct {

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

237238 Client: client,

238239 ToolCount: toolCount,

239240 }

240− if state == MCPStateConnected {

241+ switch state {

242+ case MCPStateConnected:

241243 info.ConnectedAt = time.Now()

244+ case MCPStateError:

245+ updateMcpTools(name, nil)

246+ mcpClients.Del(name)

242247 }

243248 mcpStates.Set(name, info)

244249

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

252257 })

253258}

254259

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+

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

256269func CloseMCPClients() error {

257270 var errs []error

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

274287 },

275288}

276289

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) {

278291 var wg sync.WaitGroup

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

280−

281292 // Initialize states for all configured MCPs

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

283294 if m.Disabled {

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

316327 return

317328 }

318329

330+ mcpClients.Set(name, c)

331+

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

320333 if err != nil {

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

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

324337 return

325338 }

326339

340+ updateMcpTools(name, tools)

327341 mcpClients.Set(name, c)

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

329− result.Append(tools...)

330343 }(name, m)

331344 }

332345 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+ }

334360}

335361

336362func 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

341367 return nil, err

342368 }

343369

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+

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

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

346382 timeout := mcpTimeout(m)

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

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

385+

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

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

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

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

353390 cancel()

354391 return nil, err

355392 }

393+

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

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

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

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

360398 cancel()

361399 return nil, err

362400 }

401+

363402 cancelTimer.Stop()

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

365404 return c, nil