Skip to content

fix(runtime): repopulate StateView per-server tools on reconnect (MCP-2094) #637

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Dumbris merged 1 commit into from
Jun 14, 2026
Merged

fix(runtime): repopulate StateView per-server tools on reconnect (MCP-2094) #637

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
122 changes: 60 additions & 62 deletions internal/runtime/supervisor/supervisor.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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)
})
}

Expand Down Expand Up @@ -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
Expand Down
134 changes: 127 additions & 7 deletions internal/runtime/supervisor/supervisor_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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",
Expand Down
Loading