Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
8 changes: 8 additions & 0 deletions internal/runtime/event_bus.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
5 changes: 5 additions & 0 deletions internal/runtime/events.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
23 changes: 23 additions & 0 deletions internal/runtime/lifecycle.go
Original file line number Diff line number Diff line change
Expand Up @@ -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`)
Expand Down
44 changes: 44 additions & 0 deletions internal/runtime/prompts_refresh.go
Original file line number Diff line number Diff line change
@@ -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()
})
}
32 changes: 32 additions & 0 deletions internal/runtime/prompts_refresh_test.go
Original file line number Diff line number Diff line change
@@ -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)
}
}
3 changes: 3 additions & 0 deletions internal/runtime/runtime.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
10 changes: 10 additions & 0 deletions internal/server/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
12 changes: 12 additions & 0 deletions internal/upstream/core/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 {
Expand Down
90 changes: 63 additions & 27 deletions internal/upstream/core/connection_lifecycle.go
Original file line number Diff line number Diff line change
Expand Up @@ -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, ...).
}
})

Expand All @@ -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())
Expand Down
37 changes: 37 additions & 0 deletions internal/upstream/core/prompts_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
29 changes: 29 additions & 0 deletions internal/upstream/managed/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
}

Expand Down Expand Up @@ -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
Expand Down
Loading
Loading