diff --git a/core/bifrost.go b/core/bifrost.go index 1fe56dbb58b..25f6e51a9dc 100644 --- a/core/bifrost.go +++ b/core/bifrost.go @@ -4422,6 +4422,26 @@ func (bifrost *Bifrost) ReconnectMCPClient(id string) error { return bifrost.MCPManager.ReconnectClient(id) } +// RefreshMCPClientTools re-discovers an MCP client's tools from its upstream +// server immediately, instead of waiting for the periodic connection +// checker's next tick. Applies to every client type, including the per-call +// ones ReconnectMCPClient rejects. +// +// Parameters: +// - ctx: Context bounding the discovery attempt +// - id: ID of the client to refresh +// +// Returns: +// - int: Number of tools the client is serving after the refresh +// - error: Any discovery error +func (bifrost *Bifrost) RefreshMCPClientTools(ctx context.Context, id string) (int, error) { + if bifrost.MCPManager == nil { + return 0, fmt.Errorf("mcp is not configured in this bifrost instance") + } + + return bifrost.MCPManager.RefreshClientTools(ctx, id) +} + // CloseAndMarkNeedsReauth closes a shared MCP client's live upstream // connection and flips it to needs_reauth, without attempting a new dial. // Used after OAuth credential rotation. diff --git a/core/mcp/clientmanager.go b/core/mcp/clientmanager.go index 9521706e216..ccab61870c1 100644 --- a/core/mcp/clientmanager.go +++ b/core/mcp/clientmanager.go @@ -458,6 +458,187 @@ func (m *MCPManager) ReconnectClient(id string) (retErr error) { return nil } +// writeBackDiscoveredTools installs a freshly discovered tool set on clientID +// and fires the tools-change callback when the set genuinely changed. Every +// discovery path that is not itself a connect (the periodic checker, and +// RefreshClientTools) writes through here, so they share one staleness rule: +// connGeneration is the value captured before discovery started, and a +// mismatch means a reconnect swapped in a different connection meanwhile — +// these results describe a connection that is no longer installed, so they are +// dropped rather than clobbering the fresh one. +// +// Deliberately does not touch State: the connection checker is the sole +// authority over Healthy/Unstable/NeedsReauth (see ClientConnectionChecker's +// doc comment), and a write-back is not a state transition. +// +// Returns whether the write actually landed. +func (m *MCPManager) writeBackDiscoveredTools(clientID string, connGeneration uint64, newTools map[string]schemas.ChatTool, newMapping map[string]string) bool { + // Precompute serialized JSON before the lock (see precomputeToolSerialization), + // so per-request logging/marshal reuse the bytes and the manager mutex isn't + // held across N marshals. + precomputeToolSerialization(newTools) + + m.mu.Lock() + clientState, exists := m.clientMap[clientID] + if !exists { + m.mu.Unlock() + return false + } + if clientState.ConnGeneration != connGeneration { + m.mu.Unlock() + m.logger.Debug("%s Skipping tool write-back for %s: connection was replaced during discovery", MCPLogPrefix, clientID) + return false + } + clientState.ToolMap = newTools + clientState.ToolNameMapping = newMapping + fire := m.toolsChangedCallback(clientState, clientID, newTools, newMapping) + m.mu.Unlock() + + // Fired outside the lock — see toolsChangeCallback's field doc. This is + // what persists the new set to the DB and re-syncs the hosted /mcp + // surface. Gated on genuine content change: the periodic checker reaches + // here on every tick, and most ticks rediscover identical tools. + if fire != nil { + fire() + } + return true +} + +// installedToolCount reports how many tools clientID is currently serving. +// Used when a discovery is dropped as stale: the caller's own result describes +// a tool set that was never installed, so the live map is the honest answer. +func (m *MCPManager) installedToolCount(clientID string) int { + m.mu.RLock() + defer m.mu.RUnlock() + if clientState, ok := m.clientMap[clientID]; ok { + return len(clientState.ToolMap) + } + return 0 +} + +// RefreshClientTools re-discovers clientID's tools from its upstream server +// now, instead of waiting for the connection checker's next steady-state tick +// (ResolveToolSyncInterval — 10 minutes by default). It is the operator's +// answer to "I just changed this MCP server, pick it up", and the only such +// mechanism that applies to per-call clients: ReconnectClient rejects those +// outright, since they hold no persistent connection to re-establish. +// +// Three shapes, mirroring the connection checker's own branches: +// - Live connection (sticky client): tools/list over it. +// - Per-call client: an ephemeral connect-discover-close cycle via +// performAdminToolDiscovery. +// - Sticky client whose connection is currently down: a reconnect, which +// discovers as part of the dial. +// +// Returns the number of tools the client is serving after the refresh. +// +// Does not change the client's State — that stays the connection checker's +// sole authority, so a failed refresh reports an error to the caller without +// marking the client Unstable behind their back. The reconnect branch is the +// exception, and only because connectToMCPClient owns that transition anyway. +func (m *MCPManager) RefreshClientTools(ctx context.Context, clientID string) (int, error) { + if ctx == nil { + ctx = m.ctx + } + + m.mu.RLock() + clientState, exists := m.clientMap[clientID] + var ( + conn *client.Client + config *schemas.MCPClientConfig + connGeneration uint64 + state schemas.MCPConnectionState + ) + if exists && clientState != nil { + conn = clientState.Conn + config = clientState.ExecutionConfig + connGeneration = clientState.ConnGeneration + state = clientState.State + } + m.mu.RUnlock() + + if !exists { + return 0, fmt.Errorf("mcp client %s: %w", clientID, schemas.ErrMCPClientNotFound) + } + if config == nil { + return 0, fmt.Errorf("mcp client %s has no execution config to discover with", clientID) + } + switch state { + case schemas.MCPConnectionStateDisabled: + // A disabled client has no workers and no connection; discovering for + // it would install tools nothing can execute. + return 0, fmt.Errorf("cannot refresh tools for a disabled MCP client, enable the client first: %w", schemas.ErrMCPRefreshNotApplicable) + case schemas.MCPConnectionStateNeedsReauth: + // The credential is confirmed dead, so every discovery shape below + // would fail on auth. Say so plainly instead of surfacing a raw 401. + return 0, fmt.Errorf("mcp client %s needs reauthorization before its tools can be refreshed: %w", config.Name, schemas.ErrMCPRefreshNotApplicable) + } + // A client still awaiting its one-time admin flow must not be refreshed, + // and the reason is stronger than "it would fail": for token_exchange it + // would SUCCEED. That auth type resolves its own client-credentials token, + // so discovery needs no admin at all — and because awaitsAdminVerification + // reads DiscoveredTools == nil as the pending signal for it, a successful + // refresh would persist tools through the tools-change callback and quietly + // promote the client out of pending_verification for good, taking the + // Verify CTA with it. Checked against the config rather than State alone + // for the same reason EnableClient does: State is the thing automatic + // paths keep losing, the config predicate is the durable truth. + if awaitsAdminVerification(config) || state == schemas.MCPConnectionStatePendingVerification { + return 0, fmt.Errorf("mcp client %s is awaiting admin verification, complete that instead of refreshing its tools: %w", config.Name, schemas.ErrMCPRefreshNotApplicable) + } + + switch { + case conn != nil: + // Marked as a check for the same reason the periodic checker marks its + // own list_tools (see ClientConnectionChecker.markAsCheck): this is + // Bifrost's own maintenance traffic, not a caller's inference request, + // and plugins gate on that distinction. + attemptCtx, cancel := context.WithTimeout(ctx, ConnectionCheckTimeout) + defer cancel() + bfCtx := schemas.NewBifrostContext(attemptCtx, schemas.NoDeadline) + bfCtx.SetValue(schemas.BifrostContextKeyMCPHealthCheckRequest, true) + + tools, mapping, err := m.runListToolsWithHooks(bfCtx, conn, config.Name) + if err != nil { + return 0, fmt.Errorf("failed to list tools for MCP client %s: %w", config.Name, err) + } + if !m.writeBackDiscoveredTools(clientID, connGeneration, tools, mapping) { + // Dropped as stale: a reconnect swapped the connection while this + // list was in flight, so these tools were never installed and the + // client is still serving whatever that reconnect discovered. + return m.installedToolCount(clientID), nil + } + return len(tools), nil + + case m.credStore.RequiresPerCallConnection(config): + // Same ephemeral cycle the checker's per-call branch runs; it dials, + // lists, and closes, so there is no connection to keep or reuse. + attemptCtx, cancel := context.WithTimeout(ctx, ConnectionCheckTimeout) + defer cancel() + + tools, mapping, err := m.performAdminToolDiscovery(attemptCtx, config) + if err != nil { + return 0, fmt.Errorf("failed to discover tools for MCP client %s: %w", config.Name, err) + } + if !m.writeBackDiscoveredTools(clientID, connGeneration, tools, mapping) { + // Same staleness guard as the live branch above. + return m.installedToolCount(clientID), nil + } + return len(tools), nil + + default: + // Sticky client with no live connection (still Unstable from an + // earlier failed connect, say). Discovery is part of the dial, so a + // successful reconnect has already installed the fresh tools — + // including firing the tools-change callback — by the time this + // returns. + if err := m.ReconnectClient(clientID); err != nil { + return 0, fmt.Errorf("failed to refresh tools for MCP client %s: %w", config.Name, err) + } + return m.installedToolCount(clientID), nil + } +} + // AddClient adds a new MCP client to the manager. // It validates the client configuration and establishes a connection. // If connection fails, the client entry is retained in Disconnected state and diff --git a/core/mcp/connectionchecker.go b/core/mcp/connectionchecker.go index 42026dbd3fe..9480c8bb9dc 100644 --- a/core/mcp/connectionchecker.go +++ b/core/mcp/connectionchecker.go @@ -376,7 +376,7 @@ func (c *ClientConnectionChecker) checkLiveConnection(conn *client.Client, clien return false } - c.writeBackTools(connGeneration, newTools, newMapping) + c.manager.writeBackDiscoveredTools(c.clientID, connGeneration, newTools, newMapping) c.recordSuccess(clientName, connGeneration) return true } @@ -399,7 +399,7 @@ func (c *ClientConnectionChecker) checkPerCall(config *schemas.MCPClientConfig, return false } - c.writeBackTools(connGeneration, newTools, newMapping) + c.manager.writeBackDiscoveredTools(c.clientID, connGeneration, newTools, newMapping) c.recordSuccess(config.Name, connGeneration) return true } @@ -414,46 +414,6 @@ func (c *ClientConnectionChecker) markAsCheck(ctx context.Context) *schemas.Bifr return bfCtx } -// writeBackTools mirrors the old tool-syncer's generation-guarded -// write-back: if a reconnect swapped in a fresh connection while this check -// was in flight, the fresh generation no longer matches what was captured -// before the check ran, and these (now stale) results are dropped silently -// — the next tick syncs against whatever is current. -func (c *ClientConnectionChecker) writeBackTools(connGeneration uint64, newTools map[string]schemas.ChatTool, newMapping map[string]string) { - // Precompute serialized JSON before the lock (see precomputeToolSerialization), - // so per-request logging/marshal reuse the bytes and the manager mutex isn't - // held across N marshals. - precomputeToolSerialization(newTools) - - c.manager.mu.Lock() - - clientState, exists := c.manager.clientMap[c.clientID] - if !exists { - c.manager.mu.Unlock() - return - } - if clientState.ConnGeneration != connGeneration { - c.manager.mu.Unlock() - c.logger.Debug("%s Skipping tool write-back for %s: connection was replaced during check", MCPLogPrefix, c.clientID) - return - } - clientState.ToolMap = newTools - clientState.ToolNameMapping = newMapping - fire := c.manager.toolsChangedCallback(clientState, c.clientID, newTools, newMapping) - c.manager.mu.Unlock() - - // Fired outside the lock — see toolsChangeCallback's field doc. Covers - // the periodic checker's own refresh for both sticky (checkLiveConnection) - // and per-call (checkPerCall) clients — previously the one path where a - // client's tools could drift out of sync with the DB indefinitely, since - // nothing else revisits a per-call client after its first discovery. - // Gated on genuine content change: this is the highest-frequency firing - // point (every checker tick), and most ticks rediscover identical tools. - if fire != nil { - fire() - } -} - // recordFailure marks the client Unstable — a single transient-classified // failure is enough, no consecutive-failure counter: the retry-with-backoff // each check already went through (ProbeRetryConfig) is what absorbs an diff --git a/core/mcp/interface.go b/core/mcp/interface.go index a2c9e3bb82f..fde84b0a1f2 100644 --- a/core/mcp/interface.go +++ b/core/mcp/interface.go @@ -99,6 +99,12 @@ type MCPManagerInterface interface { // ReconnectClient reconnects an MCP client by ID ReconnectClient(id string) error + // RefreshClientTools re-discovers a client's tools from its upstream + // server now rather than at the connection checker's next tick, and + // returns how many tools it is serving afterwards. Unlike + // ReconnectClient, it applies to per-call clients too. + RefreshClientTools(ctx context.Context, clientID string) (int, error) + // CloseAndMarkNeedsReauth closes a shared client's live upstream // connection and flips it to needs_reauth, without attempting a new // dial. Used after OAuth credential rotation. diff --git a/core/mcp/refreshtools_test.go b/core/mcp/refreshtools_test.go new file mode 100644 index 00000000000..2f18c690e11 --- /dev/null +++ b/core/mcp/refreshtools_test.go @@ -0,0 +1,459 @@ +package mcp + +import ( + "bytes" + "context" + "io" + "net/http" + "net/http/httptest" + "sort" + "testing" + + mcpgo "github.com/mark3labs/mcp-go/mcp" + "github.com/mark3labs/mcp-go/server" + "github.com/maximhq/bifrost/core/schemas" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// refreshToolMapKeys snapshots the prefixed tool names currently cached on a +// client, so a test can watch the manager's view of an upstream tool set +// without reaching into the connection. +func refreshToolMapKeys(m *MCPManager, id string) []string { + m.mu.RLock() + defer m.mu.RUnlock() + cs, ok := m.clientMap[id] + if !ok || cs.ToolMap == nil { + return nil + } + keys := make([]string, 0, len(cs.ToolMap)) + for k := range cs.ToolMap { + keys = append(keys, k) + } + sort.Strings(keys) + return keys +} + +// buildRefreshMCPServer starts a streamable-HTTP MCP server serving a single +// "echo" tool. The returned addTool adds another one after the fact, standing +// in for an upstream server whose tool set changed while Bifrost was connected. +func buildRefreshMCPServer(t *testing.T) (ts *httptest.Server, addTool func(name, description string)) { + t.Helper() + + s := server.NewMCPServer("refresh-tools", "1.0.0", server.WithToolCapabilities(true)) + echoTool := mcpgo.NewTool("echo", + mcpgo.WithDescription("Echo tool"), + mcpgo.WithString("message", mcpgo.Required(), mcpgo.Description("message")), + ) + s.AddTool(echoTool, func(_ context.Context, req mcpgo.CallToolRequest) (*mcpgo.CallToolResult, error) { + msg, _ := req.GetArguments()["message"].(string) + return mcpgo.NewToolResultText(msg), nil + }) + + streamable := server.NewStreamableHTTPServer(s) + ts = httptest.NewServer(http.HandlerFunc(streamable.ServeHTTP)) + t.Cleanup(func() { + // A listening streamable-HTTP transport holds a GET open, and + // httptest.Server.Close waits for outstanding requests. + ts.CloseClientConnections() + ts.Close() + }) + + return ts, func(name, description string) { + s.AddTool( + mcpgo.NewTool(name, mcpgo.WithDescription(description)), + func(context.Context, mcpgo.CallToolRequest) (*mcpgo.CallToolResult, error) { + return mcpgo.NewToolResultText("ok"), nil + }, + ) + } +} + +// newRefreshConfig builds an http client config. Leaving stickiness unset makes +// it per-call (the default for newly created clients); setting it true gives it +// a persistent connection. +func newRefreshConfig(id, name, url string, sticky bool) *schemas.MCPClientConfig { + config := &schemas.MCPClientConfig{ + ID: id, + Name: name, + ConnectionType: schemas.MCPConnectionTypeHTTP, + AuthType: schemas.MCPAuthTypeNone, + ConnectionString: schemas.NewSecretVar(url), + ToolsToExecute: []string{"*"}, + } + if sticky { + config.NeedsSessionStickiness = schemas.Ptr(true) + } + return config +} + +// TestRefreshClientTools_PerCallClient_RediscoversOnDemand is the regression +// test for #6885's third gap: a per-call client — which covers every per-user +// auth type as well as any shared http client running with +// needs_session_stickiness nil/false — had no way at all to pick up an upstream +// tool-set change short of the periodic checker's tick (up to 10 minutes) or a +// gateway restart. ReconnectClient does not apply to these clients by design: +// there is no persistent connection to re-establish. +func TestRefreshClientTools_PerCallClient_RediscoversOnDemand(t *testing.T) { + ts, addTool := buildRefreshMCPServer(t) + + m := NewMCPManager(context.Background(), schemas.MCPConfig{}, nil, &MockLogger{}, nil) + t.Cleanup(func() { _ = m.Cleanup() }) + + config := newRefreshConfig("refresh-percall", "percall", ts.URL, false) + require.NoError(t, m.AddClient(context.Background(), config)) + require.True(t, m.RequiresPerCallConnection(config), + "precondition: this client must be per-call for the test to exercise the gap") + require.Equal(t, []string{"percall-echo"}, refreshToolMapKeys(m, config.ID), + "precondition: the initial tool must be cached after AddClient") + + addTool("ping", "Ping tool added after the client was registered") + + count, err := m.RefreshClientTools(context.Background(), config.ID) + require.NoError(t, err) + require.Equal(t, 2, count, "refresh should report the freshly discovered tool count") + require.Equal(t, []string{"percall-echo", "percall-ping"}, refreshToolMapKeys(m, config.ID), + "refresh should have picked up the tool added upstream") +} + +// TestRefreshClientTools_StickyClient_RelistsOverLiveConnection covers the +// sticky branch: the client already holds a live connection, so the refresh is +// a plain tools/list over it rather than an ephemeral connect-discover-close. +func TestRefreshClientTools_StickyClient_RelistsOverLiveConnection(t *testing.T) { + ts, addTool := buildRefreshMCPServer(t) + + m := NewMCPManager(context.Background(), schemas.MCPConfig{}, nil, &MockLogger{}, nil) + t.Cleanup(func() { _ = m.Cleanup() }) + + config := newRefreshConfig("refresh-sticky", "sticky", ts.URL, true) + require.NoError(t, m.AddClient(context.Background(), config)) + require.False(t, m.RequiresPerCallConnection(config), + "precondition: this client must be sticky for the test to exercise the live-connection branch") + require.Equal(t, []string{"sticky-echo"}, refreshToolMapKeys(m, config.ID), + "precondition: the initial tool must be cached after AddClient") + + addTool("ping", "Ping tool added after the client already connected") + + count, err := m.RefreshClientTools(context.Background(), config.ID) + require.NoError(t, err) + require.Equal(t, 2, count, "refresh should report the freshly discovered tool count") + require.Equal(t, []string{"sticky-echo", "sticky-ping"}, refreshToolMapKeys(m, config.ID), + "refresh should have picked up the tool added upstream") +} + +// TestRefreshClientTools_FiresToolsChangeCallback pins the seam the transport +// layer persists through: a refresh that genuinely changes the tool set must +// reach the tools-change callback, which is what writes the new set to the DB +// and re-syncs the hosted /mcp surface. Gated on real change, exactly like +// every other discovery path. +func TestRefreshClientTools_FiresToolsChangeCallback(t *testing.T) { + // Both discovery branches, because they reach the callback by different + // routes: the per-call one through writeBackDiscoveredTools after an + // ephemeral discovery, the sticky one through the same write-back after a + // tools/list over the live connection. + for _, sticky := range []bool{false, true} { + name := "per_call" + if sticky { + name = "sticky" + } + t.Run(name, func(t *testing.T) { + ts, addTool := buildRefreshMCPServer(t) + + m := NewMCPManager(context.Background(), schemas.MCPConfig{}, nil, &MockLogger{}, nil) + t.Cleanup(func() { _ = m.Cleanup() }) + + fired := make(chan map[string]schemas.ChatTool, 4) + m.SetToolsChangeCallback(func(_, _ string, tools map[string]schemas.ChatTool, _ map[string]string) { + fired <- tools + }) + + config := newRefreshConfig("refresh-callback-"+name, "cb", ts.URL, sticky) + require.NoError(t, m.AddClient(context.Background(), config)) + drainToolsChange(fired) + + // A refresh that rediscovers the identical set is not a change, + // so it must not churn the DB or the hosted /mcp surface. + _, err := m.RefreshClientTools(context.Background(), config.ID) + require.NoError(t, err) + require.Empty(t, fired, "an unchanged rediscovery must not fire the tools-change callback") + + addTool("ping", "Ping tool added after the client was registered") + _, err = m.RefreshClientTools(context.Background(), config.ID) + require.NoError(t, err) + + require.Len(t, fired, 1, "a genuine tool-set change must fire the tools-change callback exactly once") + require.Contains(t, <-fired, "cb-ping") + }) + } +} + +// TestRefreshClientTools_UnknownClient_Errors keeps the handler's 404 mapping honest. +func TestRefreshClientTools_UnknownClient_Errors(t *testing.T) { + m := NewMCPManager(context.Background(), schemas.MCPConfig{}, nil, &MockLogger{}, nil) + t.Cleanup(func() { _ = m.Cleanup() }) + + _, err := m.RefreshClientTools(context.Background(), "no-such-client") + require.ErrorIs(t, err, schemas.ErrMCPClientNotFound) +} + +// drainToolsChange empties any callback firings from client setup so a test +// only observes what its own refresh produced. +func drainToolsChange(ch chan map[string]schemas.ChatTool) { + for { + select { + case <-ch: + default: + return + } + } +} + +// TestRefreshClientTools_AwaitingAdminVerification_Refuses covers the case an +// on-demand refresh must not quietly resolve. A client parked in +// pending_verification is waiting on a human, and for token_exchange the +// refresh would not merely fail — it would succeed, because that auth type +// resolves its own client-credentials token with no admin involved. Since +// awaitsAdminVerification reads DiscoveredTools == nil as the pending signal +// there, persisting a discovered set would promote the client out of +// pending_verification permanently and drop the Verify CTA with it. +func TestRefreshClientTools_AwaitingAdminVerification_Refuses(t *testing.T) { + cases := []struct { + name string + authType schemas.MCPAuthType + pending func(*schemas.MCPClientConfig) + }{ + { + name: "oauth_with_unauthorized_inline_config", + authType: schemas.MCPAuthTypeOauth, + pending: func(c *schemas.MCPClientConfig) { c.PendingOAuthConfig = &schemas.OAuth2Config{} }, + }, + { + name: "per_user_oauth_with_unauthorized_inline_config", + authType: schemas.MCPAuthTypePerUserOauth, + pending: func(c *schemas.MCPClientConfig) { c.PendingOAuthConfig = &schemas.OAuth2Config{} }, + }, + { + name: "token_exchange_never_verified", + authType: schemas.MCPAuthTypeTokenExchange, + pending: func(c *schemas.MCPClientConfig) { c.DiscoveredTools = nil }, + }, + { + name: "per_user_headers_never_verified", + authType: schemas.MCPAuthTypePerUserHeaders, + pending: func(c *schemas.MCPClientConfig) { c.DiscoveredTools = nil }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + m := NewMCPManager(context.Background(), schemas.MCPConfig{}, nil, &MockLogger{}, nil) + t.Cleanup(func() { m.checkerManager.StopAll() }) + + config := &schemas.MCPClientConfig{ + ID: "refresh-pending", + Name: "pending", + AuthType: tc.authType, + ConnectionType: schemas.MCPConnectionTypeHTTP, + ConnectionString: schemas.NewSecretVar("http://127.0.0.1:0/mcp"), + } + tc.pending(config) + require.True(t, awaitsAdminVerification(config), + "precondition: this config must be one AddClient parks in pending_verification") + + m.mu.Lock() + m.clientMap[config.ID] = &schemas.MCPClientState{ + Name: config.Name, + ExecutionConfig: config, + State: schemas.MCPConnectionStatePendingVerification, + ToolMap: map[string]schemas.ChatTool{}, + ToolNameMapping: map[string]string{}, + ConnectionInfo: &schemas.MCPClientConnectionInfo{Type: config.ConnectionType}, + } + m.mu.Unlock() + + _, err := m.RefreshClientTools(context.Background(), config.ID) + require.ErrorIs(t, err, schemas.ErrMCPRefreshNotApplicable) + + m.mu.RLock() + state := m.clientMap[config.ID].State + tools := len(m.clientMap[config.ID].ToolMap) + m.mu.RUnlock() + assert.Equal(t, schemas.MCPConnectionStatePendingVerification, state, + "a refused refresh must leave the client awaiting its admin") + assert.Zero(t, tools, "nothing may be discovered onto a client still awaiting verification") + }) + } +} + +// TestRefreshClientTools_DisabledAndNeedsReauth_Refuse pins the other two +// states where discovery is meaningless, and that both are reported as the +// same not-applicable class so the handler can map them to 400 rather than +// surfacing them as discovery failures. +func TestRefreshClientTools_DisabledAndNeedsReauth_Refuse(t *testing.T) { + for _, state := range []schemas.MCPConnectionState{ + schemas.MCPConnectionStateDisabled, + schemas.MCPConnectionStateNeedsReauth, + } { + t.Run(string(state), func(t *testing.T) { + m := NewMCPManager(context.Background(), schemas.MCPConfig{}, nil, &MockLogger{}, nil) + t.Cleanup(func() { m.checkerManager.StopAll() }) + + config := newRefreshConfig("refresh-"+string(state), "svc", "http://127.0.0.1:0/mcp", false) + m.mu.Lock() + m.clientMap[config.ID] = &schemas.MCPClientState{ + Name: config.Name, + ExecutionConfig: config, + State: state, + ToolMap: map[string]schemas.ChatTool{}, + ToolNameMapping: map[string]string{}, + ConnectionInfo: &schemas.MCPClientConnectionInfo{Type: config.ConnectionType}, + } + m.mu.Unlock() + + _, err := m.RefreshClientTools(context.Background(), config.ID) + require.ErrorIs(t, err, schemas.ErrMCPRefreshNotApplicable) + }) + } +} + +// buildRefreshMCPServerWithHook is buildRefreshMCPServer plus a hook fired on +// every inbound tools/list, which lets a test mutate manager state at the exact +// moment a discovery is in flight. +func buildRefreshMCPServerWithHook(t *testing.T, onToolsList func()) (ts *httptest.Server, addTool func(name, description string)) { + t.Helper() + + s := server.NewMCPServer("refresh-tools", "1.0.0", server.WithToolCapabilities(true)) + echoTool := mcpgo.NewTool("echo", + mcpgo.WithDescription("Echo tool"), + mcpgo.WithString("message", mcpgo.Required(), mcpgo.Description("message")), + ) + s.AddTool(echoTool, func(_ context.Context, req mcpgo.CallToolRequest) (*mcpgo.CallToolResult, error) { + msg, _ := req.GetArguments()["message"].(string) + return mcpgo.NewToolResultText(msg), nil + }) + + streamable := server.NewStreamableHTTPServer(s) + ts = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == http.MethodPost && onToolsList != nil { + body, err := io.ReadAll(r.Body) + if err == nil { + r.Body = io.NopCloser(bytes.NewReader(body)) + if bytes.Contains(body, []byte(`"tools/list"`)) { + onToolsList() + } + } + } + streamable.ServeHTTP(w, r) + })) + t.Cleanup(func() { + ts.CloseClientConnections() + ts.Close() + }) + + return ts, func(name, description string) { + s.AddTool( + mcpgo.NewTool(name, mcpgo.WithDescription(description)), + func(context.Context, mcpgo.CallToolRequest) (*mcpgo.CallToolResult, error) { + return mcpgo.NewToolResultText("ok"), nil + }, + ) + } +} + +// TestRefreshClientTools_StaleWriteBack_ReportsInstalledCount pins tool_count to +// what the client actually serves. writeBackDiscoveredTools drops a discovery +// whose ConnGeneration no longer matches — a reconnect swapped the connection +// while the tools/list was in flight — and leaves the existing ToolMap in place. +// Reporting the discarded discovery's length would then describe a tool set the +// client is not serving, contradicting the endpoint's documented contract. +func TestRefreshClientTools_StaleWriteBack_ReportsInstalledCount(t *testing.T) { + var m *MCPManager + const clientID = "refresh-stale-writeback" + + // Bump the generation from inside the in-flight tools/list, which is + // exactly the window writeBackDiscoveredTools guards against. + bumped := false + ts, addTool := buildRefreshMCPServerWithHook(t, func() { + if m == nil || bumped { + return + } + m.mu.Lock() + if cs, ok := m.clientMap[clientID]; ok { + cs.ConnGeneration++ + bumped = true + } + m.mu.Unlock() + }) + + m = NewMCPManager(context.Background(), schemas.MCPConfig{}, nil, &MockLogger{}, nil) + t.Cleanup(func() { _ = m.Cleanup() }) + + config := newRefreshConfig(clientID, "stale", ts.URL, false) + require.NoError(t, m.AddClient(context.Background(), config)) + require.Equal(t, []string{"stale-echo"}, refreshToolMapKeys(m, clientID), + "precondition: the initial tool must be cached after AddClient") + + // The upstream now serves two tools, so a write-back that landed would + // install 2 and a discarded one leaves the original 1 in place. + addTool("ping", "Ping tool added after the client was registered") + bumped = false + + count, err := m.RefreshClientTools(context.Background(), clientID) + require.NoError(t, err) + require.True(t, bumped, "precondition: the generation must have been bumped mid-discovery") + + installed := refreshToolMapKeys(m, clientID) + require.Equal(t, []string{"stale-echo"}, installed, + "precondition: a stale write-back must leave the existing tool map untouched") + assert.Equal(t, len(installed), count, + "tool_count must describe the tools the client actually serves, not a discarded discovery") +} + +// TestRefreshClientTools_EditedVerifiedClient_StillRefreshes reproduces a +// reported refusal: "mcp client ... is awaiting admin verification, complete +// that instead of refreshing its tools", returned for a per_user_headers server +// the UI showed as healthy. The sequence that produced it was refresh (worked), +// edit the server, refresh again (refused, and refused from then on). +// +// The refusal reads awaitsAdminVerification off the manager's in-memory config, +// which for this auth type is "DiscoveredTools is nil". UpdateClient rebuilt +// that config without DiscoveredTools, so an edit left a verified client +// reading as never verified, while State and the tool map, which the edit does +// not rebuild, went on saying healthy. The fix lives in UpdateClient; this pins +// what it bought on the route that reported it. +func TestRefreshClientTools_EditedVerifiedClient_StillRefreshes(t *testing.T) { + ts, addTool := buildRefreshMCPServer(t) + + m := NewMCPManager(context.Background(), schemas.MCPConfig{}, nil, &MockLogger{}, nil) + t.Cleanup(func() { _ = m.Cleanup() }) + // The admin's stored sample values, which per_user_headers discovery runs on. + m.credStore = &fakeAdminCredStore{headers: http.Header{"X-Api-Key": []string{"admin-sample"}}} + + config := newRefreshConfig("refresh-edited", "edited", ts.URL, false) + config.AuthType = schemas.MCPAuthTypePerUserHeaders + config.PerUserHeaderKeys = []string{"X-Api-Key"} + // Already verified: exactly what a restart loads for a server whose admin + // verification ran in an earlier process. + config.DiscoveredTools = map[string]schemas.ChatTool{"edited-echo": {}} + config.DiscoveredToolNameMapping = map[string]string{} + require.NoError(t, m.AddClient(context.Background(), config)) + + _, err := m.RefreshClientTools(context.Background(), config.ID) + require.NoError(t, err, "precondition: a verified client refreshes before it is edited") + + // An ordinary edit, shaped the way the update handler shapes it: the + // editable fields, and no DiscoveredTools. + edit := *config + edit.DiscoveredTools = nil + edit.DiscoveredToolNameMapping = nil + edit.ToolsToAutoExecute = []string{"echo"} + require.NoError(t, m.UpdateClient(config.ID, &edit)) + + addTool("ping", "Ping tool added after the client was edited") + + count, err := m.RefreshClientTools(context.Background(), config.ID) + require.NoError(t, err, "editing a verified client must not make it read as awaiting admin verification") + require.Equal(t, 2, count) + require.Equal(t, []string{"edited-echo", "edited-ping"}, refreshToolMapKeys(m, config.ID), + "and the refresh picks up what changed upstream") +} diff --git a/core/mcp/toolshash_test.go b/core/mcp/toolshash_test.go index 5979634ade2..0a4fd49d4c1 100644 --- a/core/mcp/toolshash_test.go +++ b/core/mcp/toolshash_test.go @@ -82,17 +82,16 @@ func TestWriteBackTools_UnchangedTools_DoesNotFireCallback(t *testing.T) { callCount++ }) - checker := NewClientConnectionChecker(manager, config.ID, 0, false, &MockLogger{}) tools := map[string]schemas.ChatTool{"echo": {Type: "function"}} mapping := map[string]string{"echo": "echo-server"} - checker.writeBackTools(0, tools, mapping) + manager.writeBackDiscoveredTools(config.ID, 0, tools, mapping) require.Equal(t, 1, callCount) - checker.writeBackTools(0, map[string]schemas.ChatTool{"echo": {Type: "function"}}, map[string]string{"echo": "echo-server"}) + manager.writeBackDiscoveredTools(config.ID, 0, map[string]schemas.ChatTool{"echo": {Type: "function"}}, map[string]string{"echo": "echo-server"}) assert.Equal(t, 1, callCount, "an unchanged tick must not re-fire") - checker.writeBackTools(0, map[string]schemas.ChatTool{}, map[string]string{}) + manager.writeBackDiscoveredTools(config.ID, 0, map[string]schemas.ChatTool{}, map[string]string{}) assert.Equal(t, 2, callCount, "the server legitimately losing all its tools is still a genuine change") } diff --git a/core/schemas/mcp.go b/core/schemas/mcp.go index 5e559c59865..49668b7b769 100644 --- a/core/schemas/mcp.go +++ b/core/schemas/mcp.go @@ -45,6 +45,16 @@ var ( // each user manages their own auth and there is no shared upstream // connection to "reconnect". Distinct from "not implemented". ErrMCPReconnectNotApplicable = errors.New("reconnect is not applicable for this client type") + // ErrMCPClientNotFound signals that no MCP client is registered under the + // given ID. Callers that surface this to an operator should map it to a + // 404 rather than a generic failure. + ErrMCPClientNotFound = errors.New("mcp client not found") + // ErrMCPRefreshNotApplicable signals that an on-demand tool refresh is not + // meaningful for this client right now — it is disabled, its credential is + // confirmed dead, or it is still awaiting the one-time admin verification + // flow. The request is well-formed; the client is just not in a state + // where discovery means anything. Distinct from a discovery failure. + ErrMCPRefreshNotApplicable = errors.New("tool refresh is not applicable for this client's current state") ) // MCPAuthRequiredKind discriminates the kind of inline-401 auth flow surfaced diff --git a/docs/docs.json b/docs/docs.json index 2b836d703c7..f4e7fb16a57 100644 --- a/docs/docs.json +++ b/docs/docs.json @@ -1247,6 +1247,7 @@ "PUT /api/mcp/client/{id}", "DELETE /api/mcp/client/{id}", "POST /api/mcp/client/{id}/reconnect", + "POST /api/mcp/client/{id}/refresh-tools", "POST /api/mcp/client/{id}/complete-oauth", "POST /api/mcp/client/{id}/initiate-verification", "POST /api/mcp/client/{id}/reauthorize", diff --git a/docs/openapi/openapi.yaml b/docs/openapi/openapi.yaml index 4371d727a6d..1d7e3765d43 100644 --- a/docs/openapi/openapi.yaml +++ b/docs/openapi/openapi.yaml @@ -819,6 +819,8 @@ paths: $ref: "./paths/management/mcp.yaml#/client-by-id" /api/mcp/client/{id}/reconnect: $ref: "./paths/management/mcp.yaml#/client-reconnect" + /api/mcp/client/{id}/refresh-tools: + $ref: "./paths/management/mcp.yaml#/client-refresh-tools" /api/mcp/client/{id}/complete-oauth: $ref: "./paths/management/mcp.yaml#/client-complete-oauth" /api/mcp/client/{id}/initiate-verification: diff --git a/docs/openapi/paths/management/mcp.yaml b/docs/openapi/paths/management/mcp.yaml index acaa39658dc..be115407884 100644 --- a/docs/openapi/paths/management/mcp.yaml +++ b/docs/openapi/paths/management/mcp.yaml @@ -330,6 +330,70 @@ client-reconnect: '500': $ref: '../../openapi.yaml#/components/responses/InternalError' +client-refresh-tools: + post: + operationId: refreshMCPClientTools + summary: Refresh MCP client tools + description: | + Re-discovers an MCP client's tools from its upstream server immediately, + instead of waiting for the periodic connection checker's next tick (the + client's tool_sync_interval, or the global mcp_tool_sync_interval — + 10 minutes by default). + + Use it after adding, removing, or re-describing a tool on the upstream + MCP server. Unlike reconnect, this applies to every client type, + including per-call clients (any per-user auth type, and any shared http + client with needs_session_stickiness false/omitted), which hold no + persistent connection and so have no other on-demand refresh path. + + How the refresh runs depends on the client: a sticky client with a live + connection is re-listed over it; a per-call client goes through an + ephemeral connect-discover-close cycle; a sticky client whose connection + is currently down is reconnected, which re-discovers as part of the dial. + + The freshly discovered set is persisted and the hosted /mcp surface is + re-synced automatically, exactly as for every other discovery path. + Rejected with 400 when the client is not in a state where discovery + means anything: disabled (enable it first), needs_reauth (reauthorize + it first), or still pending_verification (complete the one-time admin + verification instead — refreshing a client awaiting that flow would + bypass it). 404 if no client is registered under the given ID. + tags: + - MCP + parameters: + - name: id + in: path + required: true + description: MCP client ID + schema: + type: string + security: + - ManagementBearerAuth: [] + responses: + '200': + description: MCP client tools refreshed successfully + content: + application/json: + schema: + type: object + properties: + status: + type: string + example: success + message: + type: string + example: MCP client tools refreshed successfully + tool_count: + type: integer + description: Number of tools the client is serving after the refresh + example: 12 + '400': + $ref: '../../openapi.yaml#/components/responses/BadRequest' + '404': + $ref: '../../openapi.yaml#/components/responses/NotFound' + '500': + $ref: '../../openapi.yaml#/components/responses/InternalError' + client-complete-oauth: post: operationId: completeMCPClientOAuth diff --git a/transports/bifrost-http/handlers/mcp.go b/transports/bifrost-http/handlers/mcp.go index 93e65b93ac3..52be7d7380c 100644 --- a/transports/bifrost-http/handlers/mcp.go +++ b/transports/bifrost-http/handlers/mcp.go @@ -38,6 +38,12 @@ type MCPManager interface { // UpdateMCPClientCredentials reconnects an existing MCP client using updated headers UpdateMCPClientCredentials(ctx context.Context, id string, newConfig *schemas.MCPClientConfig) error ReconnectMCPClient(ctx context.Context, id string) error + // RefreshMCPClientTools re-discovers a client's tools from its upstream + // server on demand and reports how many it serves afterwards. Unlike + // ReconnectMCPClient it applies to per-call clients too, which have no + // persistent connection and so no other way to pick up an upstream + // tool-set change before the periodic checker's next tick. + RefreshMCPClientTools(ctx context.Context, id string) (int, error) // CloseAndMarkNeedsReauth closes a shared client's live upstream // connection and flips it to needs_reauth, without attempting a new // dial. Used after OAuth credential rotation. @@ -107,6 +113,7 @@ func (h *MCPHandler) RegisterRoutes(r *router.Router, middlewares ...schemas.Bif r.PUT("/api/mcp/client/{id}", lib.ChainMiddlewares(h.updateMCPClient, middlewares...)) r.DELETE("/api/mcp/client/{id}", lib.ChainMiddlewares(h.deleteMCPClient, middlewares...)) r.POST("/api/mcp/client/{id}/reconnect", lib.ChainMiddlewares(h.reconnectMCPClient, middlewares...)) + r.POST("/api/mcp/client/{id}/refresh-tools", lib.ChainMiddlewares(h.refreshMCPClientTools, middlewares...)) r.POST("/api/mcp/client/{id}/complete-oauth", lib.ChainMiddlewares(h.completeMCPClientOAuth, middlewares...)) r.POST("/api/mcp/client/{id}/initiate-verification", lib.ChainMiddlewares(h.initiateMCPClientVerification, middlewares...)) r.POST("/api/mcp/client/{id}/reauthorize", lib.ChainMiddlewares(h.reauthorizeMCPClient, middlewares...)) @@ -1569,6 +1576,62 @@ func (h *MCPHandler) reconnectMCPClient(ctx *fasthttp.RequestCtx) { }) } +// refreshMCPClientTools re-discovers one client's tools from its upstream MCP +// server right now. It exists because the periodic connection checker is +// otherwise the only thing that revisits a client's tool list, and its +// steady-state cadence is the tool sync interval — 10 minutes by default — so +// an operator who has just added or removed a tool upstream had nothing to +// reach for. Unlike reconnect, this applies to per-call clients as well, which +// hold no persistent connection and previously had no refresh path at all. +func (h *MCPHandler) refreshMCPClientTools(ctx *fasthttp.RequestCtx) { + if h.store.ConfigStore == nil { + SendError(ctx, fasthttp.StatusServiceUnavailable, "MCP operations unavailable: config store is disabled") + return + } + id, err := getIDFromCtx(ctx) + if err != nil { + SendError(ctx, fasthttp.StatusBadRequest, fmt.Sprintf("Invalid id: %v", err)) + return + } + // Reject a disabled client the same way reconnect does: it holds no + // connection and runs no workers, so discovered tools would have nothing + // to execute them. + if h.store.MCPConfig != nil { + for _, client := range h.store.MCPConfig.ClientConfigs { + if client.ID == id { + if client.Disabled { + SendError(ctx, fasthttp.StatusBadRequest, "cannot refresh tools for a disabled MCP client: enable the client first") + return + } + break + } + } + } + count, err := h.mcpManager.RefreshMCPClientTools(ctx, id) + if err != nil { + // An unknown client is the caller naming something that does not + // exist, not a discovery failure. + if errors.Is(err, schemas.ErrMCPClientNotFound) { + SendError(ctx, fasthttp.StatusNotFound, err.Error()) + return + } + // A client that is disabled, needs reauthorization, or is still + // awaiting admin verification is a 400: the request is well-formed, + // the client just is not in a state where discovery means anything. + if errors.Is(err, schemas.ErrMCPRefreshNotApplicable) { + SendError(ctx, fasthttp.StatusBadRequest, err.Error()) + return + } + SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("Failed to refresh MCP client tools: %v", err)) + return + } + SendJSON(ctx, map[string]any{ + "status": "success", + "message": "MCP client tools refreshed successfully", + "tool_count": count, + }) +} + // OAuthConfigRequest represents OAuth configuration in the request type OAuthConfigRequest struct { ClientID *schemas.SecretVar `json:"client_id"` diff --git a/transports/bifrost-http/handlers/mcp_disabled_to_enabled_verifyheaders_test.go b/transports/bifrost-http/handlers/mcp_disabled_to_enabled_verifyheaders_test.go index 9e7914074ae..7f94c594361 100644 --- a/transports/bifrost-http/handlers/mcp_disabled_to_enabled_verifyheaders_test.go +++ b/transports/bifrost-http/handlers/mcp_disabled_to_enabled_verifyheaders_test.go @@ -33,6 +33,9 @@ func (m *fakeMCPManagerVerifyOnly) UpdateMCPClientCredentials(_ context.Context, return nil } func (m *fakeMCPManagerVerifyOnly) ReconnectMCPClient(_ context.Context, _ string) error { return nil } +func (m *fakeMCPManagerVerifyOnly) RefreshMCPClientTools(_ context.Context, _ string) (int, error) { + return 0, nil +} func (m *fakeMCPManagerVerifyOnly) CloseAndMarkNeedsReauth(_ context.Context, _ string) error { return nil } diff --git a/transports/bifrost-http/handlers/mcp_updateclientcredentials_retry_test.go b/transports/bifrost-http/handlers/mcp_updateclientcredentials_retry_test.go index f1e45ee8376..e4151ef6b64 100644 --- a/transports/bifrost-http/handlers/mcp_updateclientcredentials_retry_test.go +++ b/transports/bifrost-http/handlers/mcp_updateclientcredentials_retry_test.go @@ -31,6 +31,9 @@ func (m *fakeUpdateCredsMCPManager) UpdateMCPClient(_ context.Context, _ string, return nil } func (m *fakeUpdateCredsMCPManager) ReconnectMCPClient(_ context.Context, _ string) error { return nil } +func (m *fakeUpdateCredsMCPManager) RefreshMCPClientTools(_ context.Context, _ string) (int, error) { + return 0, nil +} func (m *fakeUpdateCredsMCPManager) CloseAndMarkNeedsReauth(_ context.Context, _ string) error { return nil } diff --git a/transports/bifrost-http/server/server.go b/transports/bifrost-http/server/server.go index a8cb5853f9d..79eabea9643 100644 --- a/transports/bifrost-http/server/server.go +++ b/transports/bifrost-http/server/server.go @@ -165,6 +165,10 @@ type ServerCallbacks interface { // together. RequiresPerCallConnection(config *schemas.MCPClientConfig) bool ReconnectMCPClient(ctx context.Context, id string) error + // RefreshMCPClientTools re-discovers a client's tools from its upstream + // server on demand and reports how many it serves afterwards. Unlike + // ReconnectMCPClient it applies to per-call clients too. + RefreshMCPClientTools(ctx context.Context, id string) (int, error) // CloseAndMarkNeedsReauth closes a shared client's live upstream // connection and flips it to needs_reauth, without attempting a new // dial. Used after OAuth credential rotation. @@ -388,6 +392,19 @@ func (s *BifrostHTTPServer) ReconnectMCPClient(ctx context.Context, id string) e return nil } +// RefreshMCPClientTools re-discovers an MCP client's tools from its upstream +// server on demand, so an operator who has just changed that server does not +// have to wait out the connection checker's tool-sync interval (10 minutes by +// default) or restart the gateway. Applies to every client type, including +// the per-call ones ReconnectMCPClient rejects. +// +// The discovery itself persists the new tool set and re-syncs the hosted +// /mcp surface through the tools-change callback, exactly like every other +// discovery path, so there is nothing to sync here. +func (s *BifrostHTTPServer) RefreshMCPClientTools(ctx context.Context, id string) (int, error) { + return s.Client.RefreshMCPClientTools(ctx, id) +} + // UpdateMCPClient updates an MCP client in the in-memory store func (s *BifrostHTTPServer) UpdateMCPClient(ctx context.Context, id string, updatedConfig *schemas.MCPClientConfig) error { if err := s.Config.UpdateMCPClient(ctx, id, updatedConfig); err != nil {