Merge pull request #1872 from dgageot/mcp

Improve MCP server lifecycle: caching and auto-restart

Author: dgageot

examples/mcp-toolkit.yaml

@@ -41,7 +41,6 @@ agents:
       - type: mcp
         command: docker
         args: ["mcp", "gateway", "run"]
-        tools: ["mcp-activate-profile", "mcp-add", "mcp-config-set", "mcp-create-profile", "mcp-remove"]
 
 permissions:
   allow:
@@ -50,5 +49,6 @@ permissions:
      - mcp-add
      - mcp-config-set
      - mcp-create-profile
+     - mcp-find
      - mcp-remove
 

pkg/app/app.go

@@ -125,6 +125,17 @@ func New(ctx context.Context, rt runtime.Runtime, sess *session.Session, opts ..
 		})
 	}
 
+	// Subscribe to tool list changes so the sidebar updates immediately
+	// when an MCP server adds or removes tools (outside of a RunStream).
+	if tcs, ok := rt.(runtime.ToolsChangeSubscriber); ok {
+		tcs.OnToolsChanged(func(event runtime.Event) {
+			select {
+			case app.events <- event:
+			case <-ctx.Done():
+			}
+		})
+	}
+
 	return app
 }
 

pkg/runtime/runtime.go

@@ -185,6 +185,15 @@ type RAGInitializer interface {
 	StartBackgroundRAGInit(ctx context.Context, sendEvent func(Event))
 }
 
+// ToolsChangeSubscriber is implemented by runtimes that can notify when
+// toolsets report a change in their tool list (e.g. after an MCP
+// ToolListChanged notification). The provided callback is invoked
+// outside of any RunStream, so the UI can update the tool count
+// immediately.
+type ToolsChangeSubscriber interface {
+	OnToolsChanged(handler func(Event))
+}
+
 // LocalRuntime manages the execution of agents
 type LocalRuntime struct {
 	toolMap                     map[string]ToolHandler
@@ -209,6 +218,9 @@ type LocalRuntime struct {
 	// fallbackCooldowns tracks per-agent cooldown state for sticky fallback behavior
 	fallbackCooldowns    map[string]*fallbackCooldownState
 	fallbackCooldownsMux sync.RWMutex
+
+	// onToolsChanged is called when an MCP toolset reports a tool list change.
+	onToolsChanged func(Event)
 }
 
 type streamResult struct {
@@ -748,6 +760,40 @@ func (r *LocalRuntime) ResetStartupInfo() {
 	r.startupInfoEmitted = false
 }
 
+// OnToolsChanged registers a handler that is called when an MCP toolset
+// reports a tool list change outside of a RunStream. This allows the UI
+// to update the tool count immediately.
+func (r *LocalRuntime) OnToolsChanged(handler func(Event)) {
+	r.onToolsChanged = handler
+
+	for _, name := range r.team.AgentNames() {
+		a, err := r.team.Agent(name)
+		if err != nil {
+			continue
+		}
+		for _, ts := range a.ToolSets() {
+			if n, ok := tools.As[tools.ChangeNotifier](ts); ok {
+				n.SetToolsChangedHandler(r.emitToolsChanged)
+			}
+		}
+	}
+}
+
+// emitToolsChanged is the callback registered on MCP toolsets. It re-reads
+// the current agent's full tool list and pushes a ToolsetInfo event.
+func (r *LocalRuntime) emitToolsChanged() {
+	if r.onToolsChanged == nil {
+		return
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	agentTools, err := r.CurrentAgentTools(ctx)
+	if err != nil {
+		return
+	}
+	r.onToolsChanged(ToolsetInfo(len(agentTools), false, r.currentAgent))
+}
+
 // EmitStartupInfo emits initial agent, team, and toolset information for immediate sidebar display.
 // When sess is non-nil and contains token data, a TokenUsageEvent is also emitted so that the
 // sidebar can display context usage percentage on session restore.
@@ -970,6 +1016,11 @@ func (r *LocalRuntime) RunStream(ctx context.Context, sess *session.Session) <-c
 				return
 			}
 
+			// Emit updated tool count. After a ToolListChanged MCP notification
+			// the cache is invalidated, so getTools above re-fetches from the
+			// server and may return a different count.
+			events <- ToolsetInfo(len(agentTools), false, r.currentAgent)
+
 			// Check iteration limit
 			if runtimeMaxIterations > 0 && iteration >= runtimeMaxIterations {
 				slog.Debug(

pkg/runtime/runtime_test.go

@@ -270,8 +270,8 @@ func TestSimple(t *testing.T) {
 
 	// Extract the actual message from MessageAddedEvent to use in comparison
 	// (it contains dynamic fields like CreatedAt that we can't predict)
-	require.Len(t, events, 9)
-	msgAdded := events[6].(*MessageAddedEvent)
+	require.Len(t, events, 10)
+	msgAdded := events[7].(*MessageAddedEvent)
 	require.NotNil(t, msgAdded.Message)
 	require.Equal(t, "Hello", msgAdded.Message.Message.Content)
 	require.Equal(t, chat.MessageRoleAssistant, msgAdded.Message.Message.Role)
@@ -282,6 +282,7 @@ func TestSimple(t *testing.T) {
 		ToolsetInfo(0, false, "root"),
 		UserMessage("Hi", sess.ID, nil, 0),
 		StreamStarted(sess.ID, "root"),
+		ToolsetInfo(0, false, "root"),
 		AgentChoice("root", "Hello"),
 		MessageAdded(sess.ID, msgAdded.Message, "root"),
 		NewTokenUsageEvent(sess.ID, "root", &Usage{InputTokens: 3, OutputTokens: 2, ContextLength: 5, LastMessage: &MessageUsage{
@@ -310,8 +311,8 @@ func TestMultipleContentChunks(t *testing.T) {
 
 	// Extract the actual message from MessageAddedEvent to use in comparison
 	// (it contains dynamic fields like CreatedAt that we can't predict)
-	require.Len(t, events, 13)
-	msgAdded := events[10].(*MessageAddedEvent)
+	require.Len(t, events, 14)
+	msgAdded := events[11].(*MessageAddedEvent)
 	require.NotNil(t, msgAdded.Message)
 
 	expectedEvents := []Event{
@@ -320,6 +321,7 @@ func TestMultipleContentChunks(t *testing.T) {
 		ToolsetInfo(0, false, "root"),
 		UserMessage("Please greet me", sess.ID, nil, 0),
 		StreamStarted(sess.ID, "root"),
+		ToolsetInfo(0, false, "root"),
 		AgentChoice("root", "Hello "),
 		AgentChoice("root", "there, "),
 		AgentChoice("root", "how "),
@@ -350,8 +352,8 @@ func TestWithReasoning(t *testing.T) {
 
 	// Extract the actual message from MessageAddedEvent to use in comparison
 	// (it contains dynamic fields like CreatedAt that we can't predict)
-	require.Len(t, events, 11)
-	msgAdded := events[8].(*MessageAddedEvent)
+	require.Len(t, events, 12)
+	msgAdded := events[9].(*MessageAddedEvent)
 	require.NotNil(t, msgAdded.Message)
 
 	expectedEvents := []Event{
@@ -360,6 +362,7 @@ func TestWithReasoning(t *testing.T) {
 		ToolsetInfo(0, false, "root"),
 		UserMessage("Hi", sess.ID, nil, 0),
 		StreamStarted(sess.ID, "root"),
+		ToolsetInfo(0, false, "root"),
 		AgentChoiceReasoning("root", "Let me think about this..."),
 		AgentChoiceReasoning("root", " I should respond politely."),
 		AgentChoice("root", "Hello, how can I help you?"),
@@ -389,8 +392,8 @@ func TestMixedContentAndReasoning(t *testing.T) {
 
 	// Extract the actual message from MessageAddedEvent to use in comparison
 	// (it contains dynamic fields like CreatedAt that we can't predict)
-	require.Len(t, events, 12)
-	msgAdded := events[9].(*MessageAddedEvent)
+	require.Len(t, events, 13)
+	msgAdded := events[10].(*MessageAddedEvent)
 	require.NotNil(t, msgAdded.Message)
 
 	expectedEvents := []Event{
@@ -399,6 +402,7 @@ func TestMixedContentAndReasoning(t *testing.T) {
 		ToolsetInfo(0, false, "root"),
 		UserMessage("Hi there", sess.ID, nil, 0),
 		StreamStarted(sess.ID, "root"),
+		ToolsetInfo(0, false, "root"),
 		AgentChoiceReasoning("root", "The user wants a greeting"),
 		AgentChoice("root", "Hello!"),
 		AgentChoiceReasoning("root", " I should be friendly"),
@@ -450,16 +454,17 @@ func TestErrorEvent(t *testing.T) {
 		events = append(events, ev)
 	}
 
-	require.Len(t, events, 7)
+	require.Len(t, events, 8)
 	require.IsType(t, &AgentInfoEvent{}, events[0])
 	require.IsType(t, &TeamInfoEvent{}, events[1])
 	require.IsType(t, &ToolsetInfoEvent{}, events[2])
 	require.IsType(t, &UserMessageEvent{}, events[3])
 	require.IsType(t, &StreamStartedEvent{}, events[4])
-	require.IsType(t, &ErrorEvent{}, events[5])
-	require.IsType(t, &StreamStoppedEvent{}, events[6])
+	require.IsType(t, &ToolsetInfoEvent{}, events[5])
+	require.IsType(t, &ErrorEvent{}, events[6])
+	require.IsType(t, &StreamStoppedEvent{}, events[7])
 
-	errorEvent := events[5].(*ErrorEvent)
+	errorEvent := events[6].(*ErrorEvent)
 	require.Contains(t, errorEvent.Error, "simulated error")
 }
 

pkg/tools/capabilities.go

@@ -34,6 +34,12 @@ func GetInstructions(ts ToolSet) string {
 	return ""
 }
 
+// ChangeNotifier is implemented by toolsets that can notify when their
+// tool list changes (e.g. after an MCP ToolListChanged notification).
+type ChangeNotifier interface {
+	SetToolsChangedHandler(handler func())
+}
+
 // ConfigureHandlers sets all applicable handlers on a toolset.
 // It checks for Elicitable and OAuthCapable interfaces and configures them.
 // This is a convenience function that handles the capability checking internally.