diff --git a/internal/runtime/event_bus.go b/internal/runtime/event_bus.go index 6cfea91a..caf550c2 100644 --- a/internal/runtime/event_bus.go +++ b/internal/runtime/event_bus.go @@ -641,6 +641,14 @@ func (r *Runtime) EmitActivityPromptGet(serverName, promptName, sessionID, reque r.publishEvent(newEvent(EventTypeActivityPromptGet, payload)) } +// emitUpstreamPromptsChanged signals that a connected upstream changed its +// advertised prompt list at runtime (F13). Debounced by promptsRefreshDebouncer, +// so it fires at most once per window regardless of how many upstreams changed. +// server.listenForRoutingModeRefresh subscribes and calls RefreshPrompts once. +func (r *Runtime) emitUpstreamPromptsChanged() { + r.publishEvent(newEvent(EventTypeUpstreamPromptsChanged, nil)) +} + // EmitActivityConfigChange emits an event when configuration changes (Spec 024). // action is one of: server_added, server_removed, server_updated, settings_changed // source indicates how the change was triggered: "mcp", "cli", or "api" diff --git a/internal/runtime/events.go b/internal/runtime/events.go index 93e910f2..1e001962 100644 --- a/internal/runtime/events.go +++ b/internal/runtime/events.go @@ -50,6 +50,11 @@ const ( EventTypeActivityConfigChange EventType = "activity.config_change" // EventTypeActivityPromptGet is emitted when an upstream prompts/get completes (Finding F10). EventTypeActivityPromptGet EventType = "activity.prompt_get.completed" + // EventTypeUpstreamPromptsChanged is emitted (debounced) when a connected + // upstream sends notifications/prompts/list_changed, so the aggregated prompt + // set is rebuilt without waiting for an unrelated servers.changed (F13). + // Carries no payload — RefreshPrompts re-aggregates from live state. + EventTypeUpstreamPromptsChanged EventType = "upstream.prompts_changed" // Spec 026: Sensitive data detection event // EventTypeSensitiveDataDetected is emitted when sensitive data is detected in a tool call. diff --git a/internal/runtime/lifecycle.go b/internal/runtime/lifecycle.go index df7e3ea9..af024855 100644 --- a/internal/runtime/lifecycle.go +++ b/internal/runtime/lifecycle.go @@ -141,6 +141,29 @@ func (r *Runtime) StartBackgroundInitialization() { return r.DiscoverAndIndexToolsForServer(ctx, serverName) }) r.logger.Info("Tool discovery callback registered on upstream manager") + + // F13: keep the aggregated prompt list fresh when an upstream adds/removes + // a prompt at runtime (notifications/prompts/list_changed). Prompts are + // aggregated inside the MCP server layer, so we cannot RefreshPrompts from + // here — instead publish a debounced EventTypeUpstreamPromptsChanged that + // listenForRoutingModeRefresh turns into a single RefreshPrompts on its own + // goroutine (no new reentrancy). + r.promptsRefresh = newPromptsRefreshDebouncer(promptsRefreshDebounceWindow, r.emitUpstreamPromptsChanged) + r.upstreamManager.SetPromptsChangedCallback(func(serverName string) { + // Only meaningful while aggregation is on. Read live config so a + // hot-reload flip of aggregate_upstream_prompts takes effect without a + // restart; short-circuiting here avoids waking the listener (and + // re-setting built-ins on every routing-mode server) for a feature + // nobody enabled. RefreshPrompts double-guards on the same live flags. + cfg := r.Config() + if cfg == nil || !cfg.EnablePrompts || !cfg.AggregateUpstreamPrompts { + return + } + r.logger.Debug("upstream prompts/list_changed received; scheduling prompt refresh", + zap.String("server", serverName)) + r.promptsRefresh.trigger() + }) + r.logger.Info("Upstream prompts-changed callback registered on upstream manager") } // Watch the config file for external edits (editors, CLI, `jq > tmp && mv`) diff --git a/internal/runtime/prompts_refresh.go b/internal/runtime/prompts_refresh.go new file mode 100644 index 00000000..60bc32d3 --- /dev/null +++ b/internal/runtime/prompts_refresh.go @@ -0,0 +1,44 @@ +package runtime + +import ( + "sync" + "time" +) + +// promptsRefreshDebounceWindow bounds how long a burst of upstream +// notifications/prompts/list_changed notifications is held before a single +// RefreshPrompts fan-out is triggered (F13). RefreshPrompts re-lists prompts +// across EVERY connected server (one 30s-budget ListPrompts per server), and a +// shared config push can make several upstreams fire within milliseconds, so a +// trailing-edge window collapses the burst to one aggregation. 1s trades a +// barely-perceptible staleness for a large reduction in redundant fan-outs. +const promptsRefreshDebounceWindow = time.Second + +// promptsRefreshDebouncer coalesces upstream prompts/list_changed signals into +// at most one fire() per window. Trailing-edge: fire() runs once, `window` after +// the FIRST trigger of a burst; triggers landing inside the armed window are +// absorbed. +type promptsRefreshDebouncer struct { + mu sync.Mutex + timer *time.Timer + window time.Duration + fire func() +} + +func newPromptsRefreshDebouncer(window time.Duration, fire func()) *promptsRefreshDebouncer { + return &promptsRefreshDebouncer{window: window, fire: fire} +} + +func (d *promptsRefreshDebouncer) trigger() { + d.mu.Lock() + defer d.mu.Unlock() + if d.timer != nil { + return // already armed; this trigger is coalesced + } + d.timer = time.AfterFunc(d.window, func() { + d.mu.Lock() + d.timer = nil + d.mu.Unlock() + d.fire() + }) +} diff --git a/internal/runtime/prompts_refresh_test.go b/internal/runtime/prompts_refresh_test.go new file mode 100644 index 00000000..afe71483 --- /dev/null +++ b/internal/runtime/prompts_refresh_test.go @@ -0,0 +1,32 @@ +package runtime + +import ( + "sync/atomic" + "testing" + "time" +) + +// TestPromptsRefreshDebouncer_CoalescesBurst verifies a burst of triggers +// collapses to one fire per window (F13), and that a trigger after the window +// opens a fresh one. +func TestPromptsRefreshDebouncer_CoalescesBurst(t *testing.T) { + var fires int32 + d := newPromptsRefreshDebouncer(50*time.Millisecond, func() { + atomic.AddInt32(&fires, 1) + }) + + for i := 0; i < 10; i++ { + d.trigger() + } + time.Sleep(120 * time.Millisecond) + if got := atomic.LoadInt32(&fires); got != 1 { + t.Fatalf("expected exactly 1 fire for a burst, got %d", got) + } + + // A trigger after the window opens a fresh window and fires again. + d.trigger() + time.Sleep(120 * time.Millisecond) + if got := atomic.LoadInt32(&fires); got != 2 { + t.Fatalf("expected 2 fires after a second window, got %d", got) + } +} diff --git a/internal/runtime/runtime.go b/internal/runtime/runtime.go index 0f75f522..e48f196d 100644 --- a/internal/runtime/runtime.go +++ b/internal/runtime/runtime.go @@ -106,6 +106,9 @@ type Runtime struct { indexManager *index.Manager upstreamManager *upstream.Manager cacheManager *cache.Manager + // promptsRefresh debounces upstream prompts/list_changed notifications into a + // single RefreshPrompts fan-out (F13). Nil until lifecycle registration. + promptsRefresh *promptsRefreshDebouncer // truncator is swapped on config hot-reload (tool_response_limit) while the // MCP serving path reads it via Truncator(). An atomic.Pointer makes the // swap/read race-free (#861) and independent of r.mu, so the accessor stays diff --git a/internal/server/server.go b/internal/server/server.go index 85a85132..9057fc85 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -496,6 +496,16 @@ func (s *Server) listenForRoutingModeRefresh() { if s.mcpProxy != nil { s.mcpProxy.RefreshPrompts() } + case runtime.EventTypeUpstreamPromptsChanged: + // F13: an upstream added/removed a prompt at runtime (debounced + // notifications/prompts/list_changed). Rebuild only the aggregated + // prompt set — RefreshPrompts is a no-op when prompts or aggregation + // are disabled. Runs on THIS single listener goroutine, the same one + // servers.changed/config.reloaded use, so it never races another + // RefreshPrompts (no new reentrancy). + if s.mcpProxy != nil { + s.mcpProxy.RefreshPrompts() + } case runtime.EventTypeSecurityScanSettled: // Spec 086 stage 3 (FR-011): a scan-mode server that was quarantined on // add stays quarantined until its baseline scan settles GREEN. React to diff --git a/internal/upstream/core/client.go b/internal/upstream/core/client.go index 10c6e568..eca45036 100644 --- a/internal/upstream/core/client.go +++ b/internal/upstream/core/client.go @@ -137,6 +137,9 @@ type Client struct { // Notification callback for tools/list_changed onToolsChanged func(serverName string) + + // Notification callback for prompts/list_changed (F13) + onPromptsChanged func(serverName string) } // NewClient creates a new core MCP client @@ -774,6 +777,15 @@ func (c *Client) SetOnToolsChangedCallback(callback func(serverName string)) { c.onToolsChanged = callback } +// SetOnPromptsChangedCallback sets the callback invoked when a +// notifications/prompts/list_changed notification is received from the upstream +// MCP server (F13). Enables reactive re-aggregation of upstream prompts. +func (c *Client) SetOnPromptsChangedCallback(callback func(serverName string)) { + c.mu.Lock() + defer c.mu.Unlock() + c.onPromptsChanged = callback +} + // Helper methods func (c *Client) getServerName() string { diff --git a/internal/upstream/core/connection_lifecycle.go b/internal/upstream/core/connection_lifecycle.go index 75e26f33..88a8d7db 100644 --- a/internal/upstream/core/connection_lifecycle.go +++ b/internal/upstream/core/connection_lifecycle.go @@ -154,33 +154,13 @@ func (c *Client) registerNotificationHandler() { } c.client.OnNotification(func(notification mcp.JSONRPCNotification) { - // Filter for tools/list_changed notifications only - if notification.Method != string(mcp.MethodNotificationToolsListChanged) { - return - } - - c.logger.Info("Received tools/list_changed notification from upstream server", - zap.String("server", c.config.Name)) - - // Log capability status for debugging - if c.serverInfo != nil && c.serverInfo.Capabilities.Tools != nil && c.serverInfo.Capabilities.Tools.ListChanged { - c.logger.Debug("Server advertised tools.listChanged capability", - zap.String("server", c.config.Name)) - } else { - c.logger.Warn("Received tools notification from server that did not advertise listChanged capability", - zap.String("server", c.config.Name)) - } - - // Invoke the callback if set - c.mu.RLock() - callback := c.onToolsChanged - c.mu.RUnlock() - - if callback != nil { - callback(c.config.Name) - } else { - c.logger.Debug("No onToolsChanged callback set - notification ignored", - zap.String("server", c.config.Name)) + switch notification.Method { + case string(mcp.MethodNotificationToolsListChanged): + c.handleToolsListChangedNotification() + case string(mcp.MethodNotificationPromptsListChanged): + c.handlePromptsListChangedNotification() + default: + // Ignore all other notifications (logging, resources, progress, ...). } }) @@ -194,6 +174,62 @@ func (c *Client) registerNotificationHandler() { } } +// handleToolsListChangedNotification forwards a notifications/tools/list_changed +// signal to the onToolsChanged callback. serverInfo and the callback are read +// under the same RLock. +func (c *Client) handleToolsListChangedNotification() { + c.logger.Info("Received tools/list_changed notification from upstream server", + zap.String("server", c.config.Name)) + + c.mu.RLock() + serverInfo := c.serverInfo + callback := c.onToolsChanged + c.mu.RUnlock() + + if serverInfo != nil && serverInfo.Capabilities.Tools != nil && serverInfo.Capabilities.Tools.ListChanged { + c.logger.Debug("Server advertised tools.listChanged capability", + zap.String("server", c.config.Name)) + } else { + c.logger.Warn("Received tools notification from server that did not advertise listChanged capability", + zap.String("server", c.config.Name)) + } + + if callback != nil { + callback(c.config.Name) + } else { + c.logger.Debug("No onToolsChanged callback set - notification ignored", + zap.String("server", c.config.Name)) + } +} + +// handlePromptsListChangedNotification mirrors handleToolsListChangedNotification +// for notifications/prompts/list_changed (F13). It only forwards the signal; the +// managed/manager/runtime layers debounce and re-aggregate. +func (c *Client) handlePromptsListChangedNotification() { + c.logger.Info("Received prompts/list_changed notification from upstream server", + zap.String("server", c.config.Name)) + + c.mu.RLock() + serverInfo := c.serverInfo + callback := c.onPromptsChanged + c.mu.RUnlock() + + if serverInfo != nil && serverInfo.Capabilities.Prompts != nil && serverInfo.Capabilities.Prompts.ListChanged { + c.logger.Debug("Server advertised prompts.listChanged capability", + zap.String("server", c.config.Name)) + } else { + c.logger.Warn("Received prompts notification from server that did not advertise listChanged capability", + zap.String("server", c.config.Name)) + } + + if callback != nil { + callback(c.config.Name) + } else { + c.logger.Debug("No onPromptsChanged callback set - notification ignored", + zap.String("server", c.config.Name)) + } +} + // Disconnect closes the connection func (c *Client) Disconnect() error { return c.DisconnectWithContext(context.Background()) diff --git a/internal/upstream/core/prompts_test.go b/internal/upstream/core/prompts_test.go index 1f5b5fc6..bbda3640 100644 --- a/internal/upstream/core/prompts_test.go +++ b/internal/upstream/core/prompts_test.go @@ -347,3 +347,40 @@ func TestClient_ListPrompts_EndlessCursorTerminatesAtPageCap(t *testing.T) { require.NoError(t, err) assert.Len(t, prompts, maxListPromptsPages, "endless-cursor upstream must terminate at the page cap") } + +// TestClient_HandlePromptsListChanged_FiresCallback is the F13 wiring test: the +// core prompts/list_changed handler must invoke the onPromptsChanged callback +// with the server name. (Real server->client push delivery is covered by manual +// QA; this asserts the dispatch + callback the proxy relies on.) +func TestClient_HandlePromptsListChanged_FiresCallback(t *testing.T) { + upstream := newTestPromptUpstream(t, true, nil) + testServer := mcpserver.NewTestStreamableHTTPServer(upstream) + defer testServer.Close() + + c := connectedTestClient(t, testServer.URL, nil) + + got := make(chan string, 1) + c.SetOnPromptsChangedCallback(func(serverName string) { + got <- serverName + }) + + c.handlePromptsListChangedNotification() + + select { + case name := <-got: + assert.Equal(t, c.config.Name, name) + case <-time.After(time.Second): + t.Fatal("onPromptsChanged callback did not fire") + } +} + +// TestClient_HandlePromptsListChanged_NilCallbackNoPanic ensures the handler is +// safe when no callback is registered. +func TestClient_HandlePromptsListChanged_NilCallbackNoPanic(t *testing.T) { + upstream := newTestPromptUpstream(t, true, nil) + testServer := mcpserver.NewTestStreamableHTTPServer(upstream) + defer testServer.Close() + + c := connectedTestClient(t, testServer.URL, nil) + c.handlePromptsListChangedNotification() // must not panic +} diff --git a/internal/upstream/managed/client.go b/internal/upstream/managed/client.go index 24aed05b..618fbb54 100644 --- a/internal/upstream/managed/client.go +++ b/internal/upstream/managed/client.go @@ -79,6 +79,9 @@ type Client struct { // Tool discovery callback for notifications/tools/list_changed handling toolDiscoveryCallback func(ctx context.Context, serverName string) error + // Prompts-changed callback for notifications/prompts/list_changed handling (F13) + promptsChangedCallback func(serverName string) + // consecutiveHealthFailures counts back-to-back transient health-check // failures. The state-machine only flips to Error once it reaches // healthCheckFailureThreshold; one success resets it. Hard failures @@ -221,6 +224,23 @@ func NewClient(id string, serverConfig *config.ServerConfig, logger *zap.Logger, }() }) + // Wire core prompts/list_changed notifications to the manager-level callback + // (F13). Unlike the tool-discovery callback (which runs a network ListTools + // and therefore spawns a goroutine), the manager callback only schedules a + // debounced refresh — non-blocking — so no goroutine is needed on mcp-go's + // notification goroutine here. + coreClient.SetOnPromptsChangedCallback(func(serverName string) { + mc.mu.RLock() + callback := mc.promptsChangedCallback + mc.mu.RUnlock() + if callback == nil { + mc.logger.Debug("No prompts-changed callback set - notification ignored", + zap.String("server", serverName)) + return + } + callback(serverName) + }) + return mc, nil } @@ -580,6 +600,15 @@ func (mc *Client) SetToolDiscoveryCallback(callback func(ctx context.Context, se mc.toolDiscoveryCallback = callback } +// SetPromptsChangedCallback sets the callback invoked when a +// notifications/prompts/list_changed notification is received from the upstream +// server (F13). +func (mc *Client) SetPromptsChangedCallback(callback func(serverName string)) { + mc.mu.Lock() + defer mc.mu.Unlock() + mc.promptsChangedCallback = callback +} + // acquireListToolsContext claims the in-progress flag for an upstream ListTools // call. When successful it also allocates listToolsWaitCh and resets the cached // last-result, so any concurrent ListTools waiter can safely block on the diff --git a/internal/upstream/manager.go b/internal/upstream/manager.go index 444ecbec..875c1a10 100644 --- a/internal/upstream/manager.go +++ b/internal/upstream/manager.go @@ -129,6 +129,9 @@ type Manager struct { // Tool discovery callback for notifications/tools/list_changed handling toolDiscoveryCallback func(ctx context.Context, serverName string) error + // Prompts-changed callback for notifications/prompts/list_changed handling (F13) + promptsChangedCallback func(serverName string) + // limiters owns the spec-093 concurrency limiter instances (one per upstream // plus the proxy-wide aggregate). Created once and never replaced — hot // reload republishes limits INTO it so occupancy is shared across @@ -356,6 +359,16 @@ func (m *Manager) SetToolDiscoveryCallback(callback func(ctx context.Context, se m.logger.Debug("Tool discovery callback set on manager") } +// SetPromptsChangedCallback sets the callback for triggering aggregated-prompt +// refresh when an upstream sends notifications/prompts/list_changed (F13). It is +// passed to all new clients created by the manager. +func (m *Manager) SetPromptsChangedCallback(callback func(serverName string)) { + m.mu.Lock() + defer m.mu.Unlock() + m.promptsChangedCallback = callback + m.logger.Debug("Prompts-changed callback set on manager") +} + // AddServerConfig adds a server configuration without connecting func (m *Manager) AddServerConfig(id string, serverConfig *config.ServerConfig) error { m.mu.Lock() @@ -434,6 +447,11 @@ func (m *Manager) AddServerConfig(id string, serverConfig *config.ServerConfig) client.SetToolDiscoveryCallback(m.toolDiscoveryCallback) } + // Set up prompts-changed callback for notifications/prompts/list_changed (F13) + if m.promptsChangedCallback != nil { + client.SetPromptsChangedCallback(m.promptsChangedCallback) + } + // Spec 093: install admission control before the client becomes reachable, // so no dispatch can ever see a client without its limiter wiring. client.SetAdmissionControl(m.limiters, m.currentRejectObserver())