diff --git a/internal/broker/upstream/manager.go b/internal/broker/upstream/manager.go index 4300f47dd..7738ce4fe 100644 --- a/internal/broker/upstream/manager.go +++ b/internal/broker/upstream/manager.go @@ -8,7 +8,6 @@ import ( "errors" "fmt" "log/slog" - "reflect" "slices" "sync" "time" @@ -171,20 +170,12 @@ type MCPManager struct { // baseBackoff stores the initial backoff configuration for resets baseBackoff wait.Backoff gatewayServer ToolsAdderDeleter - // serverTools is an internal copy that contains the managed MCP's tools with prefixed names. It is these that are externally available via the gateway - serverTools []GatewayTool - // tools is the original set from MCP server with no prefix - tools []mcp.Tool - toolsMap map[string]*mcp.Tool - servedToolsMap map[string]*mcp.Tool + // tools registry manages the lifecycle of tools from this upstream server + tools *entityRegistry[mcp.Tool, GatewayTool] promptsServer PromptsAdderDeleter - // serverPrompts is an internal copy with prefixed names - serverPrompts []GatewayPrompt - // prompts is the original set from MCP server with no prefix - prompts []mcp.Prompt - promptsMap map[string]*mcp.Prompt - servedPromptsMap map[string]*mcp.Prompt + // prompts registry manages the lifecycle of prompts from this upstream server + prompts *entityRegistry[mcp.Prompt, GatewayPrompt] // toolsLock protects tools, serverTools, prompts, serverPrompts toolsLock sync.RWMutex @@ -250,12 +241,46 @@ func NewUpstreamMCPManager(upstream MCP, gatewayServer ToolsAdderDeleter, prompt toolEvents: make(chan struct{}, 1), promptEvents: make(chan struct{}, 1), done: make(chan struct{}), - toolsMap: map[string]*mcp.Tool{}, - servedToolsMap: map[string]*mcp.Tool{}, - serverTools: []GatewayTool{}, - promptsMap: map[string]*mcp.Prompt{}, - servedPromptsMap: map[string]*mcp.Prompt{}, - serverPrompts: []GatewayPrompt{}, + tools: newToolRegistry( + upstream.GetPrefix(), + string(upstream.ID()), + gatewayServer, + func(ctx context.Context) ([]mcp.Tool, error) { + res, err := upstream.ListTools(ctx) + if err != nil { + return nil, err + } + fetched := make([]mcp.Tool, 0, len(res.Tools)) + for _, t := range res.Tools { + if t == nil { + continue + } + fetched = append(fetched, *t) + } + return fetched, nil + }, + logger, + ), + prompts: newPromptRegistry( + upstream.GetPrefix(), + string(upstream.ID()), + promptsServer, + func(ctx context.Context) ([]mcp.Prompt, error) { + res, err := upstream.ListPrompts(ctx) + if err != nil { + return nil, err + } + fetched := make([]mcp.Prompt, 0, len(res.Prompts)) + for _, p := range res.Prompts { + if p == nil { + continue + } + fetched = append(fetched, *p) + } + return fetched, nil + }, + logger, + ), }, nil } @@ -280,8 +305,8 @@ func (man *MCPManager) Start(ctx context.Context) ActiveMCPServer { if err := man.mcp.Disconnect(); err != nil { man.logger.Error("failed to disconnect during stop", "upstream mcp server", man.mcp.ID(), "error", err) } - man.removeAllPrompts() - man.removeAllTools() + man.prompts.removeAll(&man.toolsLock) + man.tools.removeAll(&man.toolsLock) close(man.done) man.logger.Debug("manager stopped", "upstream mcp server", man.mcp.ID()) @@ -373,25 +398,28 @@ func (man *MCPManager) manage(ctx context.Context, event eventType) { // The manager stays alive and will check again on the next ticker tick. if !man.mcp.IsEnabled() { man.logger.Debug("server is not enabled, removing tools and prompts", "upstream mcp server", man.mcp.ID()) - man.removeAllTools() - man.removeAllPrompts() + man.tools.removeAll(&man.toolsLock) + man.prompts.removeAll(&man.toolsLock) _ = man.mcp.Disconnect() man.consecutiveFailures = 0 man.setStatus(fmt.Errorf("server is disabled"), 0, 0, nil, nil) return } - numberOfTools := len(man.tools) - numberOfPrompts := len(man.prompts) + numberOfTools := len(man.tools.items) + numberOfPrompts := len(man.prompts.items) // during connect the client will validate the protocol. So we don't have a separate validate requirement currently. If a client already exists it will be re-used. man.logger.DebugContext(ctx, "attempting to connect", "upstream mcp server", man.mcp.ID()) if err := man.mcp.Connect(ctx, man.registerCallbacks()); err != nil { - man.handleConnectionFailure(ctx, span, fmt.Errorf("failed to connect to upstream mcp %s : %w", man.mcp.ID(), err), numberOfTools, numberOfPrompts) + err = fmt.Errorf("failed to connect to upstream mcp %s removing tools : %w", man.mcp.ID(), err) + man.handleConnectionFailure(ctx, span, err, numberOfTools, numberOfPrompts) return } // there may be an active client so we also ping if err := man.mcp.Ping(ctx); err != nil { - man.handleConnectionFailure(ctx, span, fmt.Errorf("upstream mcp failed to ping server %s : %w", man.mcp.ID(), err), numberOfTools, numberOfPrompts) + // if we fail to ping we disconnect to ensure a fresh connection next time around + err = fmt.Errorf("upstream mcp failed to ping server %s removing tools : %w", man.mcp.ID(), err) + man.handleConnectionFailure(ctx, span, err, numberOfTools, numberOfPrompts) return } man.consecutiveFailures = 0 @@ -413,7 +441,7 @@ func (man *MCPManager) manage(ctx context.Context, event eventType) { man.logger.DebugContext(ctx, "not fetching tools", "event", event, "upstream mcp server", man.mcp.ID(), "waiting for notification", notificationToolsListChanged) } else { man.logger.DebugContext(ctx, "fetching tools", "upstream mcp server", man.mcp.ID()) - current, fetched, err := man.getTools(ctx) + current, fetched, err := man.tools.get(ctx) if err != nil { toolErr = fmt.Errorf("upstream mcp failed to list tools server %s : %w", man.mcp.ID(), err) man.recordBackendError(span, toolErr) @@ -430,7 +458,7 @@ func (man *MCPManager) manage(ctx context.Context, event eventType) { if man.invalidToolPolicy == InvalidToolPolicyRejectServer { toolErr = fmt.Errorf("upstream mcp %s rejected: %d invalid tools found", man.mcp.ID(), len(invalids)) man.recordBackendError(span, toolErr) - man.removeAllTools() + man.tools.removeAll(&man.toolsLock) } else { fetched = validTools } @@ -438,26 +466,26 @@ func (man *MCPManager) manage(ctx context.Context, event eventType) { if toolErr == nil { // always compare the tools without prefix - toAdd, toRemove := man.diffTools(current, fetched) - if conflictErr := man.findToolConflicts(toAdd); conflictErr != nil { + toAdd, toRemove := man.tools.diff(current, fetched) + if conflictErr := man.tools.findConflicts(toAdd); conflictErr != nil { toolErr = fmt.Errorf("upstream mcp failed to add tools to gateway %s : %w", man.mcp.ID(), conflictErr) man.recordBackendError(span, toolErr) man.logger.ErrorContext(ctx, "tool conflict detected", "upstream mcp server", man.mcp.ID(), "error", toolErr) } else { man.toolsLock.Lock() - man.tools = fetched + man.tools.items = fetched numberOfTools = len(fetched) - man.toolsMap = make(map[string]*mcp.Tool, len(fetched)) - man.servedToolsMap = make(map[string]*mcp.Tool, len(fetched)) + man.tools.byName = make(map[string]*mcp.Tool, len(fetched)) + man.tools.byServedName = make(map[string]*mcp.Tool, len(fetched)) for i := range fetched { - man.toolsMap[fetched[i].Name] = &fetched[i] + man.tools.byName[fetched[i].Name] = &fetched[i] toolName := prefixedName(man.mcp.GetPrefix(), fetched[i].Name) - man.servedToolsMap[toolName] = &fetched[i] + man.tools.byServedName[toolName] = &fetched[i] } - man.serverTools = slices.DeleteFunc(man.serverTools, func(tool GatewayTool) bool { + man.tools.serverItems = slices.DeleteFunc(man.tools.serverItems, func(tool GatewayTool) bool { return slices.Contains(toRemove, tool.Tool.Name) }) - man.serverTools = append(man.serverTools, toAdd...) + man.tools.serverItems = append(man.tools.serverItems, toAdd...) man.toolsLock.Unlock() man.logger.DebugContext(ctx, "updating gateway tools", "upstream mcp server", man.mcp.ID(), "adding", len(toAdd), "removing", len(toRemove)) @@ -467,7 +495,7 @@ func (man *MCPManager) manage(ctx context.Context, event eventType) { if len(toAdd) > 0 { man.gatewayServer.AddTools(toAdd...) } - man.logger.DebugContext(ctx, "internal tools", "upstream mcp server", man.mcp.ID(), "total", len(man.serverTools)) + man.logger.DebugContext(ctx, "internal tools", "upstream mcp server", man.mcp.ID(), "total", len(man.tools.serverItems)) } } } @@ -476,7 +504,7 @@ func (man *MCPManager) manage(ctx context.Context, event eventType) { var promptErr error var invalidPrompts []InvalidPromptInfo if man.promptsServer != nil && man.mcp.SupportsPrompts() && man.shouldFetchPrompts(event) { - currentPrompts, fetchedPrompts, listErr := man.getPrompts(ctx) + currentPrompts, fetchedPrompts, listErr := man.prompts.get(ctx) if listErr != nil { promptErr = fmt.Errorf("upstream mcp failed to list prompts server %s : %w", man.mcp.ID(), listErr) man.recordBackendError(span, promptErr) @@ -492,26 +520,26 @@ func (man *MCPManager) manage(ctx context.Context, event eventType) { fetchedPrompts = validPrompts } - toAddPrompts, toRemovePrompts := man.diffPrompts(currentPrompts, fetchedPrompts) - if conflictErr := man.findPromptConflicts(toAddPrompts); conflictErr != nil { + toAddPrompts, toRemovePrompts := man.prompts.diff(currentPrompts, fetchedPrompts) + if conflictErr := man.prompts.findConflicts(toAddPrompts); conflictErr != nil { promptErr = fmt.Errorf("upstream mcp failed to add prompts to gateway %s : %w", man.mcp.ID(), conflictErr) man.recordBackendError(span, promptErr) man.logger.ErrorContext(ctx, "prompt conflict detected", "upstream mcp server", man.mcp.ID(), "error", promptErr) } else { man.toolsLock.Lock() - man.prompts = fetchedPrompts + man.prompts.items = fetchedPrompts numberOfPrompts = len(fetchedPrompts) - man.promptsMap = make(map[string]*mcp.Prompt, len(fetchedPrompts)) - man.servedPromptsMap = make(map[string]*mcp.Prompt, len(fetchedPrompts)) + man.prompts.byName = make(map[string]*mcp.Prompt, len(fetchedPrompts)) + man.prompts.byServedName = make(map[string]*mcp.Prompt, len(fetchedPrompts)) for i := range fetchedPrompts { - man.promptsMap[fetchedPrompts[i].Name] = &fetchedPrompts[i] + man.prompts.byName[fetchedPrompts[i].Name] = &fetchedPrompts[i] promptName := prefixedName(man.mcp.GetPrefix(), fetchedPrompts[i].Name) - man.servedPromptsMap[promptName] = &fetchedPrompts[i] + man.prompts.byServedName[promptName] = &fetchedPrompts[i] } - man.serverPrompts = slices.DeleteFunc(man.serverPrompts, func(prompt GatewayPrompt) bool { + man.prompts.serverItems = slices.DeleteFunc(man.prompts.serverItems, func(prompt GatewayPrompt) bool { return slices.Contains(toRemovePrompts, prompt.Prompt.Name) }) - man.serverPrompts = append(man.serverPrompts, toAddPrompts...) + man.prompts.serverItems = append(man.prompts.serverItems, toAddPrompts...) man.toolsLock.Unlock() man.logger.DebugContext(ctx, "updating gateway prompts", "upstream mcp server", man.mcp.ID(), "adding", len(toAddPrompts), "removing", len(toRemovePrompts)) @@ -540,12 +568,24 @@ func (man *MCPManager) manage(ctx context.Context, event eventType) { // upstreams without the stream and events sent while it reconnects are // lost unless the upstream replays them. func (man *MCPManager) shouldFetchTools(event eventType) bool { - return event == eventTypeTimer || event == eventTypeToolNotification + // fetch if no support for tools list change notifications + if !man.mcp.SupportsToolsListChanged() { + return true + } + if event == eventTypeToolNotification { + return true + } + return event == eventTypeTimer } -// shouldFetchPrompts mirrors shouldFetchTools for prompts. func (man *MCPManager) shouldFetchPrompts(event eventType) bool { - return event == eventTypeTimer || event == eventTypePromptNotification + if !man.mcp.SupportsPromptsListChanged() { + return true + } + if event == eventTypePromptNotification { + return true + } + return event == eventTypeTimer } // GetStatus returns the current status of the MCP Server @@ -590,8 +630,8 @@ func (man *MCPManager) handleConnectionFailure(ctx context.Context, span trace.S man.logger.ErrorContext(ctx, "upstream connection failure", "upstream mcp server", man.mcp.ID(), "error", err, "consecutive failures", man.consecutiveFailures) if man.consecutiveFailures >= maxConsecutiveFailures { man.logger.ErrorContext(ctx, "failure threshold reached, removing tools and prompts", "upstream mcp server", man.mcp.ID(), "threshold", maxConsecutiveFailures) - man.removeAllTools() - man.removeAllPrompts() + man.tools.removeAll(&man.toolsLock) + man.prompts.removeAll(&man.toolsLock) } _ = man.mcp.Disconnect() man.setStatus(err, numberOfTools, numberOfPrompts, nil, nil) @@ -629,68 +669,32 @@ func (man *MCPManager) recordBackendError(span trace.Span, err error) { ) } -func (man *MCPManager) findToolConflicts(mcpTools []GatewayTool) error { - gatewayServerTools := man.gatewayServer.ListTools() - var conflictingToolNames []string - for _, tool := range mcpTools { - for existingToolName, existingToolInfo := range gatewayServerTools { - existingTool := existingToolInfo.Tool - if existingTool.Meta == nil { - man.logger.Debug("skipping conflict check, tool meta is nil", "upstream mcp server", man.mcp.ID(), "tool", existingToolName) - continue - } - existingToolID, ok := existingTool.Meta[gatewayServerID] - if !ok { - man.logger.Debug("skipping conflict check, tool id is missing", "upstream mcp server", man.mcp.ID(), "tool", existingToolName) - continue - } - toolID, is := existingToolID.(string) - if !is { - man.logger.Debug("skipping conflict check, tool id is not a string", "upstream mcp server", man.mcp.ID(), "type", reflect.TypeOf(existingToolID)) - continue - } - - if existingToolName == tool.Tool.Name && toolID != string(man.mcp.ID()) { - man.logger.Debug("tool name conflict found", "upstream mcp server", man.mcp.ID(), "existing", existingToolName, "new", tool.Tool.Name, "conflicting server", toolID) - conflictingToolNames = append(conflictingToolNames, tool.Tool.Name) - } - - } - } - if len(conflictingToolNames) > 0 { - return fmt.Errorf("conflicting tools discovered. conflicting tool names %v", conflictingToolNames) - } +// SetStatusForTesting sets the status directly for testing purposes. +// This bypasses the normal status update flow and should only be used in tests. +func (man *MCPManager) SetStatusForTesting(status ServerValidationStatus) { + man.status = status +} - return nil +// NewActiveForTesting wraps a manager as an ActiveMCPServer without starting +// the event loop. Stop is a no-op. Only for use in tests that need a static +// manager with pre-seeded tools/status. +func NewActiveForTesting(man *MCPManager) ActiveMCPServer { + return &activeMCP{manager: man, cancel: func() {}} } -// getTools return the existing, and new tools. Must only be called from the Start() event loop. -func (man *MCPManager) getTools(ctx context.Context) ([]mcp.Tool, []mcp.Tool, error) { - tools := make([]mcp.Tool, len(man.tools)) - copy(tools, man.tools) - res, err := man.mcp.ListTools(ctx) - if err != nil { - return tools, tools, fmt.Errorf("failed to get tools: %w", err) - } - fetched := make([]mcp.Tool, 0, len(res.Tools)) - for _, t := range res.Tools { - if t == nil { - man.logger.Debug("skipping nil tool from upstream", "server", man.mcp.ID()) - continue - } - fetched = append(fetched, *t) +func prefixedName(prefix, tool string) string { + if prefix == "" { + return tool } - return tools, fetched, nil + return fmt.Sprintf("%s%s", prefix, tool) } // GetManagedTools returns a copy of all tools discovered from the upstream server. // The returned tools have their original names without the gateway prefix. func (man *MCPManager) GetManagedTools() []mcp.Tool { man.toolsLock.RLock() - result := make([]mcp.Tool, len(man.tools)) - copy(result, man.tools) - man.toolsLock.RUnlock() - return result + defer man.toolsLock.RUnlock() + return man.tools.getManagedCopy() } // GetServedManagedTool will return the tool if present that is actually being served by the gateway. @@ -699,7 +703,7 @@ func (man *MCPManager) GetManagedTools() []mcp.Tool { func (man *MCPManager) GetServedManagedTool(toolName string) *mcp.Tool { man.toolsLock.RLock() defer man.toolsLock.RUnlock() - return man.servedToolsMap[toolName] + return man.tools.getServed(toolName) } // SetToolsForTesting sets the tools directly for testing purposes. @@ -708,224 +712,26 @@ func (man *MCPManager) GetServedManagedTool(toolName string) *mcp.Tool { func (man *MCPManager) SetToolsForTesting(tools []mcp.Tool) { man.toolsLock.Lock() defer man.toolsLock.Unlock() - man.tools = tools - // set a tools map for quick look up by other functions - for i := range tools { - man.toolsMap[tools[i].Name] = &tools[i] - man.servedToolsMap[prefixedName(man.mcp.GetPrefix(), tools[i].Name)] = &tools[i] - } -} - -// SetStatusForTesting sets the status directly for testing purposes. -// This bypasses the normal status update flow and should only be used in tests. -func (man *MCPManager) SetStatusForTesting(status ServerValidationStatus) { - man.status = status -} - -// NewActiveForTesting wraps a manager as an ActiveMCPServer without starting -// the event loop. Stop is a no-op. Only for use in tests that need a static -// manager with pre-seeded tools/status. -func NewActiveForTesting(man *MCPManager) ActiveMCPServer { - return &activeMCP{manager: man, cancel: func() {}} -} - -func (man *MCPManager) removeAllTools() { - man.toolsLock.Lock() - toolsToRemove := make([]string, 0, len(man.serverTools)) - man.logger.Debug("removing tools from gateway", "upstream mcp server", man.mcp.ID(), "total", len(man.serverTools)) - for _, tool := range man.serverTools { - man.logger.Debug("removing tool from server ", "upstream mcp server", man.mcp.ID(), "tool", tool.Tool.Name) - toolsToRemove = append(toolsToRemove, tool.Tool.Name) - } - man.serverTools = []GatewayTool{} - man.tools = []mcp.Tool{} - man.toolsMap = map[string]*mcp.Tool{} - man.servedToolsMap = map[string]*mcp.Tool{} - man.toolsLock.Unlock() - man.gatewayServer.DeleteTools(toolsToRemove...) - man.logger.Debug("removed all tools", "upstream mcp server", man.mcp.ID(), "count", len(toolsToRemove)) -} - -func (man *MCPManager) removeAllPrompts() { - if man.promptsServer == nil { - return - } - man.toolsLock.Lock() - promptsToRemove := make([]string, 0, len(man.serverPrompts)) - for _, prompt := range man.serverPrompts { - promptsToRemove = append(promptsToRemove, prompt.Prompt.Name) - } - man.serverPrompts = []GatewayPrompt{} - man.prompts = []mcp.Prompt{} - man.promptsMap = map[string]*mcp.Prompt{} - man.servedPromptsMap = map[string]*mcp.Prompt{} - man.toolsLock.Unlock() - man.promptsServer.DeletePrompts(promptsToRemove...) -} - -func (man *MCPManager) toolToServerTool(newTool mcp.Tool) GatewayTool { - newTool.Name = prefixedName(man.mcp.GetPrefix(), newTool.Name) - newTool.Meta = mcp.Meta{ - gatewayServerID: string(man.mcp.ID()), - } - return GatewayTool{ - Tool: newTool, - Handler: NoopToolHandler, - } -} - -func (man *MCPManager) diffTools(oldTools, newTools []mcp.Tool) ([]GatewayTool, []string) { - oldToolMap := make(map[string]mcp.Tool) - for _, oldTool := range oldTools { - oldToolMap[oldTool.Name] = oldTool - } - - newToolMap := make(map[string]mcp.Tool) - for _, newTool := range newTools { - newToolMap[newTool.Name] = newTool - } - - addedTools := make([]GatewayTool, 0) - for _, newTool := range newToolMap { - _, ok := oldToolMap[newTool.Name] - if !ok { - addedTools = append(addedTools, man.toolToServerTool(newTool)) - } - } - - removedTools := make([]string, 0) - for _, oldTool := range oldToolMap { - _, ok := newToolMap[oldTool.Name] - if !ok { - removedTools = append(removedTools, prefixedName(man.mcp.GetPrefix(), oldTool.Name)) - } - } - - return addedTools, removedTools -} - -func prefixedName(prefix, tool string) string { - if prefix == "" { - return tool - } - return fmt.Sprintf("%s%s", prefix, tool) -} - -func (man *MCPManager) promptToServerPrompt(newPrompt mcp.Prompt) GatewayPrompt { - newPrompt.Name = prefixedName(man.mcp.GetPrefix(), newPrompt.Name) - newPrompt.Meta = mcp.Meta{ - gatewayServerID: string(man.mcp.ID()), - } - return GatewayPrompt{ - Prompt: newPrompt, - Handler: NoopPromptHandler, - } -} - -func (man *MCPManager) diffPrompts(oldPrompts, newPrompts []mcp.Prompt) ([]GatewayPrompt, []string) { - oldPromptMap := make(map[string]mcp.Prompt) - for _, p := range oldPrompts { - oldPromptMap[p.Name] = p - } - - newPromptMap := make(map[string]mcp.Prompt) - for _, p := range newPrompts { - newPromptMap[p.Name] = p - } - - addedPrompts := make([]GatewayPrompt, 0) - for _, newPrompt := range newPromptMap { - if _, ok := oldPromptMap[newPrompt.Name]; !ok { - addedPrompts = append(addedPrompts, man.promptToServerPrompt(newPrompt)) - } - } - - removedPrompts := make([]string, 0) - for _, oldPrompt := range oldPromptMap { - if _, ok := newPromptMap[oldPrompt.Name]; !ok { - removedPrompts = append(removedPrompts, prefixedName(man.mcp.GetPrefix(), oldPrompt.Name)) - } - } - - return addedPrompts, removedPrompts -} - -// getPrompts returns the existing and new prompts. Must only be called from the Start() event loop. -func (man *MCPManager) getPrompts(ctx context.Context) ([]mcp.Prompt, []mcp.Prompt, error) { - prompts := make([]mcp.Prompt, len(man.prompts)) - copy(prompts, man.prompts) - res, err := man.mcp.ListPrompts(ctx) - if err != nil { - return prompts, prompts, fmt.Errorf("failed to get prompts: %w", err) - } - fetched := make([]mcp.Prompt, 0, len(res.Prompts)) - for _, p := range res.Prompts { - if p == nil { - man.logger.Debug("skipping nil prompt from upstream", "server", man.mcp.ID()) - continue - } - fetched = append(fetched, *p) - } - return prompts, fetched, nil -} - -func (man *MCPManager) findPromptConflicts(mcpPrompts []GatewayPrompt) error { - if man.promptsServer == nil { - return nil - } - gatewayServerPrompts := man.promptsServer.ListPrompts() - var conflictingPromptNames []string - for _, prompt := range mcpPrompts { - for existingPromptName, existingPromptInfo := range gatewayServerPrompts { - existingPrompt := existingPromptInfo.Prompt - if existingPrompt.Meta == nil { - man.logger.Debug("skipping conflict check, prompt meta is nil", "upstream mcp server", man.mcp.ID(), "prompt", existingPromptName) - continue - } - existingPromptID, ok := existingPrompt.Meta[gatewayServerID] - if !ok { - man.logger.Debug("skipping conflict check, prompt id is missing", "upstream mcp server", man.mcp.ID(), "prompt", existingPromptName) - continue - } - promptID, is := existingPromptID.(string) - if !is { - man.logger.Debug("skipping conflict check, prompt id is not a string", "upstream mcp server", man.mcp.ID(), "type", reflect.TypeOf(existingPromptID)) - continue - } - if existingPromptName == prompt.Prompt.Name && promptID != string(man.mcp.ID()) { - conflictingPromptNames = append(conflictingPromptNames, prompt.Prompt.Name) - } - } - } - if len(conflictingPromptNames) > 0 { - return fmt.Errorf("conflicting prompts discovered. conflicting prompt names %v", conflictingPromptNames) - } - return nil + man.tools.setForTesting(tools, man.mcp.GetPrefix()) } // GetManagedPrompts returns a copy of all prompts discovered from the upstream server. func (man *MCPManager) GetManagedPrompts() []mcp.Prompt { man.toolsLock.RLock() - result := make([]mcp.Prompt, len(man.prompts)) - copy(result, man.prompts) - man.toolsLock.RUnlock() - return result + defer man.toolsLock.RUnlock() + return man.prompts.getManagedCopy() } // GetServedManagedPrompt returns the prompt if present that is being served by the gateway. func (man *MCPManager) GetServedManagedPrompt(promptName string) *mcp.Prompt { man.toolsLock.RLock() defer man.toolsLock.RUnlock() - return man.servedPromptsMap[promptName] + return man.prompts.getServed(promptName) } // SetPromptsForTesting sets prompts directly for testing purposes. func (man *MCPManager) SetPromptsForTesting(prompts []mcp.Prompt) { man.toolsLock.Lock() defer man.toolsLock.Unlock() - man.prompts = prompts - for i := range prompts { - man.promptsMap[prompts[i].Name] = &prompts[i] - man.servedPromptsMap[prefixedName(man.mcp.GetPrefix(), prompts[i].Name)] = &prompts[i] - } + man.prompts.setForTesting(prompts, man.mcp.GetPrefix()) } diff --git a/internal/broker/upstream/manager_test.go b/internal/broker/upstream/manager_test.go index 0b73b194d..77ac8f06d 100644 --- a/internal/broker/upstream/manager_test.go +++ b/internal/broker/upstream/manager_test.go @@ -387,7 +387,7 @@ func TestMCPManager_setStatus(t *testing.T) { mock := newMockMCP("test-server", "test_") manager, err := NewUpstreamMCPManager(mock, newMockToolsAdderDeleter(), nil, logger, 0, InvalidToolPolicyFilterOut) require.NoError(t, err) - manager.serverTools = make([]GatewayTool, tc.numServerTools) + manager.tools.serverItems = make([]GatewayTool, tc.numServerTools) manager.setStatus(tc.err, tc.totalTools, 0, nil, nil) @@ -474,7 +474,7 @@ func TestMCPManager_toolToServerTool(t *testing.T) { Description: "A test tool", } - serverTool := manager.toolToServerTool(tool) + serverTool := manager.tools.toServer(tool) assert.Equal(t, "prefix_mytool", serverTool.Tool.Name) assert.Equal(t, "A test tool", serverTool.Tool.Description) @@ -780,7 +780,7 @@ func TestDiffTools(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - added, removed := manager.diffTools(tt.oldTools, tt.newTools) + added, removed := manager.tools.diff(tt.oldTools, tt.newTools) assert.Len(t, added, tt.expectedAdded, "unexpected number of added tools") assert.Len(t, removed, tt.expectedRemoved, "unexpected number of removed tools") @@ -907,7 +907,8 @@ func TestMCPManager_shouldFetchTools(t *testing.T) { require.NoError(t, err) if tt.hasExistingTools { - manager.serverTools = []GatewayTool{{Tool: mcp.Tool{Name: "existing_tool"}}} + // set serverTools directly since shouldFetchTools checks this field + manager.tools.serverItems = []GatewayTool{{Tool: mcp.Tool{Name: "existing_tool"}}} } result := manager.shouldFetchTools(tt.eventType) @@ -1107,9 +1108,9 @@ func TestServerToolsManagement(t *testing.T) { // Verify serverTools manager.toolsLock.RLock() - serverToolNames := make([]string, len(manager.serverTools)) - for i, st := range manager.serverTools { - serverToolNames[i] = st.Tool.Name + serverToolNames := make([]string, len(manager.tools.serverItems)) + for i := range manager.tools.serverItems { + serverToolNames[i] = manager.tools.serverItems[i].Tool.Name } manager.toolsLock.RUnlock() @@ -1395,7 +1396,7 @@ func TestMCPManager_promptToServerPrompt(t *testing.T) { manager, err := NewUpstreamMCPManager(mock, newMockToolsAdderDeleter(), nil, logger, 0, InvalidToolPolicyFilterOut) require.NoError(t, err) - serverPrompt := manager.promptToServerPrompt(mcp.Prompt{Name: tt.promptName, Description: tt.promptDesc}) + serverPrompt := manager.prompts.toServer(mcp.Prompt{Name: tt.promptName, Description: tt.promptDesc}) assert.Equal(t, tt.expectedName, serverPrompt.Prompt.Name) assert.Equal(t, tt.promptDesc, serverPrompt.Prompt.Description) @@ -1457,7 +1458,7 @@ func TestMCPManager_diffPrompts(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - added, removed := manager.diffPrompts(tt.oldPrompts, tt.newPrompts) + added, removed := manager.prompts.diff(tt.oldPrompts, tt.newPrompts) assert.Len(t, added, tt.expectedAdded) assert.Len(t, removed, tt.expectedRemoved) }) diff --git a/internal/broker/upstream/registry.go b/internal/broker/upstream/registry.go new file mode 100644 index 000000000..e9b34da23 --- /dev/null +++ b/internal/broker/upstream/registry.go @@ -0,0 +1,302 @@ +package upstream + +import ( + "context" + "fmt" + "log/slog" + "reflect" + + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +// entityRegistry consolidates the duplicated tool/prompt logic into a single +// generic struct. The type-specific bits (how to convert an entity to its +// server representation, how to interact with the gateway, how to fetch from +// upstream) are injected via function pointers at construction time. +// +// Callers are responsible for holding the appropriate lock around mutations. +type entityRegistry[E any, S any] struct { + // items holds the original entities from the upstream server (no prefix) + items []E + // byName maps unprefixed name → entity pointer for fast lookup + byName map[string]*E + // byServedName maps prefixed name → entity pointer for fast lookup + byServedName map[string]*E + // serverItems holds the gateway-facing representations (with prefixed names) + serverItems []S + + // getName returns the unprefixed name of an entity + getName func(E) string + // getServerName returns the prefixed name of a server-side entity + getServerName func(S) string + // toServer converts an upstream entity to its gateway-facing representation. + // This replaces toolToServerTool / promptToServerPrompt. + toServer func(E) S + // getMetaFields extracts the Meta map from a server-side entity + getMetaFields func(S) map[string]any + + // addToGateway registers server items with the gateway + addToGateway func(...S) + // deleteFromGateway removes named items from the gateway + deleteFromGateway func(...string) + // listFromGateway returns all items currently registered with the gateway + listFromGateway func() map[string]*S + // listFromMCP fetches the current items from the upstream MCP server + listFromMCP func(context.Context) ([]E, error) + + logger *slog.Logger + serverID string + // entityTag is used in log and error messages e.g. "tool" or "prompt" + entityTag string + // prefix is the gateway-facing name prefix applied to served entities + prefix string +} + +// newEntityRegistry constructs an entityRegistry with all function pointers injected. +func newEntityRegistry[E any, S any]( + getName func(E) string, + getServerName func(S) string, + toServer func(E) S, + getMetaFields func(S) map[string]any, + addToGateway func(...S), + deleteFromGateway func(...string), + listFromGateway func() map[string]*S, + listFromMCP func(context.Context) ([]E, error), + logger *slog.Logger, + serverID string, + entityTag string, + prefix string, +) *entityRegistry[E, S] { + return &entityRegistry[E, S]{ + items: []E{}, + byName: make(map[string]*E), + byServedName: make(map[string]*E), + serverItems: []S{}, + getName: getName, + getServerName: getServerName, + toServer: toServer, + getMetaFields: getMetaFields, + addToGateway: addToGateway, + deleteFromGateway: deleteFromGateway, + listFromGateway: listFromGateway, + listFromMCP: listFromMCP, + logger: logger, + serverID: serverID, + entityTag: entityTag, + prefix: prefix, + } +} + +// diff computes what needs to be added and removed when transitioning from +// oldItems to newItems. Replaces diffTools / diffPrompts. +func (r *entityRegistry[E, S]) diff(oldItems, newItems []E) (toAdd []S, toRemove []string) { + oldMap := make(map[string]E, len(oldItems)) + for _, item := range oldItems { + oldMap[r.getName(item)] = item + } + newMap := make(map[string]E, len(newItems)) + for _, item := range newItems { + newMap[r.getName(item)] = item + } + for name, item := range newMap { + if _, exists := oldMap[name]; !exists { + toAdd = append(toAdd, r.toServer(item)) + } + } + for name := range oldMap { + if _, exists := newMap[name]; !exists { + toRemove = append(toRemove, prefixedName(r.prefix, name)) + } + } + return toAdd, toRemove +} + +// findConflicts checks whether any of toAdd conflicts with items already +// registered in the gateway that belong to a different upstream server. +// Replaces findToolConflicts / findPromptConflicts. +func (r *entityRegistry[E, S]) findConflicts(toAdd []S) error { + if r.listFromGateway == nil { + return nil + } + gatewayItems := r.listFromGateway() + var conflicting []string + for _, candidate := range toAdd { + candidateName := r.getServerName(candidate) + for existingName, existingPtr := range gatewayItems { + if existingPtr == nil { + r.logger.Debug("skipping conflict check, nil entry in gateway items", + "upstream mcp server", r.serverID, r.entityTag, existingName) + continue + } + meta := r.getMetaFields(*existingPtr) + if meta == nil { + r.logger.Debug("skipping conflict check, meta is nil", + "upstream mcp server", r.serverID, r.entityTag, existingName) + continue + } + existingID, ok := meta[gatewayServerID] + if !ok { + r.logger.Debug("skipping conflict check, id is missing", + "upstream mcp server", r.serverID, r.entityTag, existingName) + continue + } + idStr, ok := existingID.(string) + if !ok { + r.logger.Debug("skipping conflict check, id is not a string", + "upstream mcp server", r.serverID, r.entityTag, existingName, "type", reflect.TypeOf(existingID)) + continue + } + if existingName == candidateName && idStr != r.serverID { + r.logger.Debug("conflict found", + "upstream mcp server", r.serverID, + "existing", existingName, "new", candidateName, + "conflicting server", idStr) + conflicting = append(conflicting, candidateName) + } + } + } + if len(conflicting) > 0 { + return fmt.Errorf("conflicting %ss discovered. conflicting %s names %v", r.entityTag, r.entityTag, conflicting) + } + return nil +} + +// removeAll removes all registered items from the gateway and clears internal +// state. The caller must hold the write lock before calling, and must NOT hold +// it when the gateway delete call happens — matching the existing pattern in +// removeAllTools / removeAllPrompts. +// Replaces removeAllTools / removeAllPrompts. +func (r *entityRegistry[E, S]) removeAll(lock interface { + Lock() + Unlock() +}) { + lock.Lock() + toRemove := make([]string, 0, len(r.serverItems)) + for _, item := range r.serverItems { + toRemove = append(toRemove, r.getServerName(item)) + } + r.serverItems = []S{} + r.items = []E{} + r.byName = make(map[string]*E) + r.byServedName = make(map[string]*E) + lock.Unlock() + + if len(toRemove) > 0 { + r.deleteFromGateway(toRemove...) + } +} + +// get fetches the current items from the upstream MCP server and returns +// (existing, fetched, error). Replaces getTools / getPrompts. +func (r *entityRegistry[E, S]) get(ctx context.Context) (existing []E, fetched []E, err error) { + existing = make([]E, len(r.items)) + copy(existing, r.items) + fetched, err = r.listFromMCP(ctx) + if err != nil { + return existing, existing, fmt.Errorf("failed to get %ss: %w", r.entityTag, err) + } + return existing, fetched, nil +} + +// getManagedCopy returns a copy of the current items slice. +// Replaces GetManagedTools / GetManagedPrompts. +func (r *entityRegistry[E, S]) getManagedCopy() []E { + result := make([]E, len(r.items)) + copy(result, r.items) + return result +} + +// getServed returns the entity pointer for the given prefixed name, or nil. +// Replaces GetServedManagedTool / GetServedManagedPrompt. +func (r *entityRegistry[E, S]) getServed(servedName string) *E { + return r.byServedName[servedName] +} + +// setForTesting sets items directly, bypassing the normal fetch flow. +// Replaces SetToolsForTesting / SetPromptsForTesting. +func (r *entityRegistry[E, S]) setForTesting(items []E, prefix string) { + r.items = items + r.byName = make(map[string]*E, len(items)) + r.byServedName = make(map[string]*E, len(items)) + for i := range items { + name := r.getName(items[i]) + r.byName[name] = &items[i] + served := prefixedName(prefix, name) + r.byServedName[served] = &items[i] + } +} + +// newToolRegistry constructs the tools entityRegistry for an MCPManager. +func newToolRegistry(prefix, serverID string, gatewayServer ToolsAdderDeleter, listFromMCP func(context.Context) ([]mcp.Tool, error), logger *slog.Logger) *entityRegistry[mcp.Tool, GatewayTool] { + return newEntityRegistry[mcp.Tool, GatewayTool]( + func(t mcp.Tool) string { return t.Name }, + func(gt GatewayTool) string { return gt.Tool.Name }, + func(t mcp.Tool) GatewayTool { + t.Name = prefixedName(prefix, t.Name) + t.Meta = mcp.Meta{ + gatewayServerID: serverID, + } + return GatewayTool{ + Tool: t, + Handler: func(_ context.Context, _ *mcp.CallToolRequest) (*mcp.CallToolResult, error) { + return NewToolResultError("Kagenti MCP Broker doesn't forward tool calls"), nil + }, + } + }, + func(gt GatewayTool) map[string]any { + return gt.Tool.Meta + }, + gatewayServer.AddTools, + gatewayServer.DeleteTools, + func() map[string]*GatewayTool { return gatewayServer.ListTools() }, + listFromMCP, + logger, + serverID, + "tool", + prefix, + ) +} + +// newPromptRegistry constructs the prompts entityRegistry for an MCPManager. +func newPromptRegistry(prefix, serverID string, promptsServer PromptsAdderDeleter, listFromMCP func(context.Context) ([]mcp.Prompt, error), logger *slog.Logger) *entityRegistry[mcp.Prompt, GatewayPrompt] { + return newEntityRegistry[mcp.Prompt, GatewayPrompt]( + func(p mcp.Prompt) string { return p.Name }, + func(gp GatewayPrompt) string { return gp.Prompt.Name }, + func(p mcp.Prompt) GatewayPrompt { + p.Name = prefixedName(prefix, p.Name) + p.Meta = mcp.Meta{ + gatewayServerID: serverID, + } + return GatewayPrompt{ + Prompt: p, + Handler: func(_ context.Context, _ *mcp.GetPromptRequest) (*mcp.GetPromptResult, error) { + return &mcp.GetPromptResult{}, nil + }, + } + }, + func(gp GatewayPrompt) map[string]any { + return gp.Prompt.Meta + }, + func(gp ...GatewayPrompt) { + if promptsServer != nil { + promptsServer.AddPrompts(gp...) + } + }, + func(names ...string) { + if promptsServer != nil { + promptsServer.DeletePrompts(names...) + } + }, + func() map[string]*GatewayPrompt { + if promptsServer == nil { + return nil + } + return promptsServer.ListPrompts() + }, + listFromMCP, + logger, + serverID, + "prompt", + prefix, + ) +}