diff --git a/internal/runtime/supervisor/supervisor.go b/internal/runtime/supervisor/supervisor.go index 560379019..acf10b547 100644 --- a/internal/runtime/supervisor/supervisor.go +++ b/internal/runtime/supervisor/supervisor.go @@ -630,31 +630,7 @@ func (s *Supervisor) updateStateView(name string, state *ServerState) { status.ToolCount = state.ToolCount // Phase 7.1: Convert ToolMetadata to ToolInfo and cache in StateView - if state.Tools != nil { - status.Tools = make([]stateview.ToolInfo, len(state.Tools)) - for i, tool := range state.Tools { - // Parse ParamsJSON into InputSchema - var inputSchema map[string]interface{} - if tool.ParamsJSON != "" { - // ParamsJSON is already a JSON string, we'll store it as-is - // The API endpoint will parse it if needed - inputSchema = map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{}, // TODO: Parse ParamsJSON - } - } - - status.Tools[i] = stateview.ToolInfo{ - Name: tool.Name, - Description: tool.Description, - InputSchema: inputSchema, - Annotations: tool.Annotations, - OutputSchemaJSON: tool.OutputSchemaJSON, - } - } - } else { - status.Tools = nil - } + status.Tools = toolInfosFromMetadata(state.Tools) // Map connection state to string // Use detailed state from ConnectionInfo when available to avoid mislabeling disconnected servers as "connecting" @@ -765,6 +741,37 @@ func (s *Supervisor) SetErrorCodeNotifier(fn func(code string)) { s.errorCodeNotifier = fn } +// toolInfosFromMetadata converts cached upstream tool metadata into the +// StateView ToolInfo representation. Shared by reconcile, background-discovery +// refresh, and reconnect repopulation so all three paths produce an identical +// StateView tool set (MCP-2094). +func toolInfosFromMetadata(tools []*config.ToolMetadata) []stateview.ToolInfo { + if tools == nil { + return nil + } + infos := make([]stateview.ToolInfo, len(tools)) + for i, tool := range tools { + // Parse ParamsJSON into InputSchema + var inputSchema map[string]interface{} + if tool.ParamsJSON != "" { + // ParamsJSON is already a JSON string; the API endpoint parses it if needed. + inputSchema = map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{}, // TODO: Parse ParamsJSON + } + } + + infos[i] = stateview.ToolInfo{ + Name: tool.Name, + Description: tool.Description, + InputSchema: inputSchema, + Annotations: tool.Annotations, + OutputSchemaJSON: tool.OutputSchemaJSON, + } + } + return infos +} + // RefreshToolsFromDiscovery updates both the Supervisor snapshot and StateView with tools from background discovery. // This is called after DiscoverAndIndexTools completes to populate the UI cache. func (s *Supervisor) RefreshToolsFromDiscovery(tools []*config.ToolMetadata) error { @@ -811,39 +818,16 @@ func (s *Supervisor) RefreshToolsFromDiscovery(tools []*config.ToolMetadata) err // Update StateView for each server for serverName, serverTools := range toolsByServer { s.stateView.UpdateServer(serverName, func(status *stateview.ServerStatus) { - // Defensive check: Only update if we have more or equal tools than currently shown - // This prevents overwriting valid tools with stale data from delayed discoveries - // Exception: Always update if current tools is 0 (initial population) - if len(status.Tools) > 0 && len(serverTools) < len(status.Tools) { - s.logger.Debug("StateView already has more tools, skipping update to prevent stale data", - zap.String("server", serverName), - zap.Int("current_tools", len(status.Tools)), - zap.Int("new_tools", len(serverTools))) - return - } - + // StateView mirrors the Supervisor snapshot (updated unconditionally + // above), so apply discovery results as last-writer-wins. A prior + // size-based guard skipped updates whenever the new set was smaller, + // which pinned StateView to a stale higher count when a server + // legitimately dropped tools — diverging from the snapshot and from + // the bleve index. Servers with zero discovered tools never reach + // this loop (they're absent from toolsByServer), so a size guard + // could not protect against empty/stale discoveries anyway (MCP-2094). status.ToolCount = len(serverTools) - status.Tools = make([]stateview.ToolInfo, len(serverTools)) - - for i, tool := range serverTools { - // Parse ParamsJSON into InputSchema - var inputSchema map[string]interface{} - if tool.ParamsJSON != "" { - // ParamsJSON is already a JSON string - inputSchema = map[string]interface{}{ - "type": "object", - "properties": map[string]interface{}{}, // TODO: Parse ParamsJSON - } - } - - status.Tools[i] = stateview.ToolInfo{ - Name: tool.Name, - Description: tool.Description, - InputSchema: inputSchema, - Annotations: tool.Annotations, - OutputSchemaJSON: tool.OutputSchemaJSON, - } - } + status.Tools = toolInfosFromMetadata(serverTools) }) } @@ -942,13 +926,27 @@ func (s *Supervisor) updateSnapshotFromEvent(event Event) { if connected { t := event.Timestamp status.ConnectedAt = &t - // Don't populate Tools here - background indexing will handle it - // Only update the count if tools haven't been discovered yet - // This prevents overwriting tool data from background discovery + // Repopulate the per-server tool set from the retained + // Supervisor snapshot so StateView stays the consistent + // source of truth across a reconnect/unquarantine. The + // disconnect branch below clears StateView.Tools, but the + // snapshot keeps them (Tools are never cleared in the + // snapshot), so we can restore immediately instead of + // waiting for background discovery to re-run. Background + // discovery (RefreshToolsFromDiscovery) later overwrites + // with fresh data. Without this, StateView consumers that + // don't use the #635 read fallback (tray counts, SSE + // servers.changed, health/diagnostics) report 0 tools for a + // connected server that has tools (MCP-2094). if len(status.Tools) == 0 { - status.ToolCount = toolCount + if len(state.Tools) > 0 { + status.Tools = toolInfosFromMetadata(state.Tools) + status.ToolCount = len(state.Tools) + } else { + status.ToolCount = toolCount + } } - // If tools are already populated, keep the existing count + // If tools are already populated, keep the existing set/count // CRITICAL: Clear error when connected, even if connInfo is unavailable // This ensures stale OAuth/connection errors don't persist after successful reconnection diff --git a/internal/runtime/supervisor/supervisor_test.go b/internal/runtime/supervisor/supervisor_test.go index 07c7f20c5..e3b4a45b5 100644 --- a/internal/runtime/supervisor/supervisor_test.go +++ b/internal/runtime/supervisor/supervisor_test.go @@ -15,13 +15,13 @@ import ( // MockUpstreamAdapter is a test double for UpstreamAdapter type MockUpstreamAdapter struct { - mu sync.Mutex - addedServers map[string]*config.ServerConfig - removedServers []string - connected map[string]bool - disconnected []string - eventCh chan Event - states map[string]*ServerState + mu sync.Mutex + addedServers map[string]*config.ServerConfig + removedServers []string + connected map[string]bool + disconnected []string + eventCh chan Event + states map[string]*ServerState } func NewMockUpstreamAdapter() *MockUpstreamAdapter { @@ -556,6 +556,126 @@ func TestSupervisor_RefreshToolsFromDiscovery_EmptyTools(t *testing.T) { } } +// TestSupervisor_ReconnectRepopulatesStateViewTools verifies that after a +// server disconnects (which clears the StateView per-server tool set) and then +// reconnects, the StateView tool set is repopulated from the retained +// Supervisor snapshot instead of being left empty until background discovery +// re-runs. This is the root-cause regression behind MCP-2083 (tracked as +// MCP-2094): StateView consumers that don't route through the #635 read +// fallback (tray counts, SSE servers.changed, health/diagnostics) would +// otherwise report 0 tools for a connected server that has tools. +func TestSupervisor_ReconnectRepopulatesStateViewTools(t *testing.T) { + cfg := &config.Config{ + Listen: "127.0.0.1:8080", + Servers: []*config.ServerConfig{ + {Name: "server1", Enabled: true}, + }, + } + + configSvc := configsvc.NewService(cfg, "/tmp/config.json", zap.NewNop()) + defer configSvc.Close() + + mockUpstream := NewMockUpstreamAdapter() + defer mockUpstream.Close() + _ = mockUpstream.AddServer("server1", cfg.Servers[0]) + + sup := New(configSvc, mockUpstream, zap.NewNop()) + + // Populate initial state + tools (simulating a connected server whose + // background discovery has completed). + _ = sup.reconcile(configSvc.Current()) + tools := []*config.ToolMetadata{ + {Name: "tool1", ServerName: "server1", Description: "Test tool 1"}, + {Name: "tool2", ServerName: "server1", Description: "Test tool 2"}, + } + if err := sup.RefreshToolsFromDiscovery(tools); err != nil { + t.Fatalf("RefreshToolsFromDiscovery failed: %v", err) + } + mockUpstream.SetServerTools("server1", tools) + + // Sanity: StateView shows the 2 tools. + if got := len(sup.StateView().Snapshot().Servers["server1"].Tools); got != 2 { + t.Fatalf("setup: expected 2 tools in StateView, got %d", got) + } + + // Disconnect: StateView deliberately clears the per-server tool set. + sup.updateSnapshotFromEvent(Event{ + Type: EventServerDisconnected, + ServerName: "server1", + Timestamp: time.Now(), + Payload: map[string]interface{}{"connected": false}, + }) + if got := len(sup.StateView().Snapshot().Servers["server1"].Tools); got != 0 { + t.Fatalf("after disconnect: expected StateView tools cleared, got %d", got) + } + + // Reconnect: StateView must be repopulated from the retained Supervisor + // snapshot immediately, without waiting for background discovery to re-run. + sup.updateSnapshotFromEvent(Event{ + Type: EventServerConnected, + ServerName: "server1", + Timestamp: time.Now(), + Payload: map[string]interface{}{"connected": true}, + }) + + status := sup.StateView().Snapshot().Servers["server1"] + if status.ToolCount != 2 { + t.Errorf("after reconnect: expected ToolCount 2, got %d", status.ToolCount) + } + if len(status.Tools) != 2 { + t.Errorf("after reconnect: expected StateView repopulated with 2 tools, got %d", len(status.Tools)) + } + if len(status.Tools) == 2 && status.Tools[0].Name != "tool1" { + t.Errorf("after reconnect: expected first tool 'tool1', got %q", status.Tools[0].Name) + } +} + +// TestSupervisor_RefreshToolsFromDiscovery_ShrinkingToolSet verifies that a +// later discovery reporting fewer (but non-empty) tools updates StateView +// rather than being silently skipped. The old size-based guard pinned +// StateView to a stale higher count, diverging from the Supervisor snapshot +// (which is updated unconditionally). Part of MCP-2094. +func TestSupervisor_RefreshToolsFromDiscovery_ShrinkingToolSet(t *testing.T) { + cfg := &config.Config{ + Listen: "127.0.0.1:8080", + Servers: []*config.ServerConfig{ + {Name: "server1", Enabled: true}, + }, + } + + configSvc := configsvc.NewService(cfg, "/tmp/config.json", zap.NewNop()) + defer configSvc.Close() + + mockUpstream := NewMockUpstreamAdapter() + defer mockUpstream.Close() + + sup := New(configSvc, mockUpstream, zap.NewNop()) + _ = sup.reconcile(configSvc.Current()) + + threeTools := []*config.ToolMetadata{ + {Name: "tool1", ServerName: "server1"}, + {Name: "tool2", ServerName: "server1"}, + {Name: "tool3", ServerName: "server1"}, + } + if err := sup.RefreshToolsFromDiscovery(threeTools); err != nil { + t.Fatalf("RefreshToolsFromDiscovery (3 tools) failed: %v", err) + } + + // Upstream now legitimately exposes only one tool. + oneTool := []*config.ToolMetadata{{Name: "tool1", ServerName: "server1"}} + if err := sup.RefreshToolsFromDiscovery(oneTool); err != nil { + t.Fatalf("RefreshToolsFromDiscovery (1 tool) failed: %v", err) + } + + status := sup.StateView().Snapshot().Servers["server1"] + if status.ToolCount != 1 { + t.Errorf("expected ToolCount 1 after shrink, got %d", status.ToolCount) + } + if len(status.Tools) != 1 { + t.Errorf("expected StateView to reflect 1 tool after shrink, got %d", len(status.Tools)) + } +} + func TestSupervisor_InspectionExemption_GrantAndRevoke(t *testing.T) { cfg := &config.Config{ Listen: "127.0.0.1:8080",