Skip to content

Commit 04c7f19

Browse files
Simplify reload
* fix issue with generating too many tool list change notifications
1 parent 0456157 commit 04c7f19

File tree

4 files changed

+398
-308
lines changed

4 files changed

+398
-308
lines changed

pkg/gateway/capabilitites.go

Lines changed: 40 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,42 @@ type ResourceTemplateRegistration struct {
4646
Handler mcp.ResourceHandler
4747
}
4848

49+
func (caps *Capabilities) getToolByName(toolName string) (ToolRegistration, error) {
50+
for _, tool := range caps.Tools {
51+
if tool.Tool.Name == toolName {
52+
return tool, nil
53+
}
54+
}
55+
return ToolRegistration{}, fmt.Errorf("unable to find tool")
56+
}
57+
58+
func (caps *Capabilities) getPromptByName(promptName string) (PromptRegistration, error) {
59+
for _, prompt := range caps.Prompts {
60+
if prompt.Prompt.Name == promptName {
61+
return prompt, nil
62+
}
63+
}
64+
return PromptRegistration{}, fmt.Errorf("unable to find prompt")
65+
}
66+
67+
func (caps *Capabilities) getResourceByURI(resourceURI string) (ResourceRegistration, error) {
68+
for _, resource := range caps.Resources {
69+
if resource.Resource.URI == resourceURI {
70+
return resource, nil
71+
}
72+
}
73+
return ResourceRegistration{}, fmt.Errorf("unable to find resource")
74+
}
75+
76+
func (caps *Capabilities) getResourceTemplateByURITemplate(resource string) (ResourceTemplateRegistration, error) {
77+
for _, template := range caps.ResourceTemplates {
78+
if template.ResourceTemplate.URITemplate == resource {
79+
return template, nil
80+
}
81+
}
82+
return ResourceTemplateRegistration{}, fmt.Errorf("unable to find resource template")
83+
}
84+
4985
func (g *Gateway) listCapabilities(ctx context.Context, configuration Configuration, serverNames []string, clientConfig *clientConfig) (*Capabilities, error) {
5086
var (
5187
lock sync.Mutex
@@ -222,17 +258,17 @@ func (g *Gateway) listCapabilities(ctx context.Context, configuration Configurat
222258
}, nil
223259
}
224260

225-
func (c *Capabilities) ToolNames() []string {
261+
func (caps *Capabilities) ToolNames() []string {
226262
var names []string
227-
for _, tool := range c.Tools {
263+
for _, tool := range caps.Tools {
228264
names = append(names, tool.Tool.Name)
229265
}
230266
return names
231267
}
232268

233-
func (c *Capabilities) PromptNames() []string {
269+
func (caps *Capabilities) PromptNames() []string {
234270
var names []string
235-
for _, prompt := range c.Prompts {
271+
for _, prompt := range caps.Prompts {
236272
names = append(names, prompt.Prompt.Name)
237273
}
238274
return names

pkg/gateway/dynamic_mcps.go

Lines changed: 17 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -208,7 +208,7 @@ type ServerMatch struct {
208208
}
209209

210210
// mcpAddTool implements a tool for adding new servers to the registry
211-
func (g *Gateway) createMcpAddTool(configuration Configuration, clientConfig *clientConfig) *ToolRegistration {
211+
func (g *Gateway) createMcpAddTool(clientConfig *clientConfig) *ToolRegistration {
212212
tool := &mcp.Tool{
213213
Name: "mcp-add",
214214
Description: "Add a new MCP server to the session. The server must exist in the catalog.",
@@ -250,7 +250,7 @@ func (g *Gateway) createMcpAddTool(configuration Configuration, clientConfig *cl
250250
serverName := strings.TrimSpace(params.Name)
251251

252252
// Check if server exists in catalog
253-
_, _, found := configuration.Find(serverName)
253+
_, _, found := g.configuration.Find(serverName)
254254
if !found {
255255
return &mcp.CallToolResult{
256256
Content: []mcp.Content{&mcp.TextContent{
@@ -261,31 +261,29 @@ func (g *Gateway) createMcpAddTool(configuration Configuration, clientConfig *cl
261261

262262
// Append the new server to the current serverNames if not already present
263263
found = false
264-
for _, existing := range configuration.serverNames {
264+
for _, existing := range g.configuration.serverNames {
265265
if existing == serverName {
266266
found = true
267267
break
268268
}
269269
}
270270
if !found {
271-
configuration.serverNames = append(configuration.serverNames, serverName)
271+
g.configuration.serverNames = append(g.configuration.serverNames, serverName)
272272
}
273273

274274
// Fetch updated secrets for the new server list
275275
if g.configurator != nil {
276276
if fbc, ok := g.configurator.(*FileBasedConfiguration); ok {
277-
updatedSecrets, err := fbc.readDockerDesktopSecrets(ctx, configuration.servers, configuration.serverNames)
277+
updatedSecrets, err := fbc.readDockerDesktopSecrets(ctx, g.configuration.servers, g.configuration.serverNames)
278278
if err == nil {
279-
configuration.secrets = updatedSecrets
279+
g.configuration.secrets = updatedSecrets
280280
} else {
281281
log("Warning: Failed to update secrets:", err)
282282
}
283283
}
284284
}
285285

286-
// Update the current configuration state
287-
updatedServerNames := configuration.serverNames
288-
if err := g.reloadConfiguration(ctx, configuration, updatedServerNames, clientConfig); err != nil {
286+
if err := g.reloadServerConfiguration(ctx, serverName, clientConfig); err != nil {
289287
return nil, fmt.Errorf("failed to reload configuration: %w", err)
290288
}
291289

@@ -303,7 +301,7 @@ func (g *Gateway) createMcpAddTool(configuration Configuration, clientConfig *cl
303301
}
304302

305303
// mcpRemoveTool implements a tool for removing servers from the registry
306-
func (g *Gateway) createMcpRemoveTool(_ Configuration, clientConfig *clientConfig) *ToolRegistration {
304+
func (g *Gateway) createMcpRemoveTool() *ToolRegistration {
307305
tool := &mcp.Tool{
308306
Name: "mcp-remove",
309307
Description: "Remove an MCP server from the registry and reload the configuration. This will disable the server.",
@@ -352,8 +350,8 @@ func (g *Gateway) createMcpRemoveTool(_ Configuration, clientConfig *clientConfi
352350
// Update the current configuration state
353351
g.configuration.serverNames = updatedServerNames
354352

355-
if err := g.reloadConfiguration(ctx, g.configuration, updatedServerNames, clientConfig); err != nil {
356-
return nil, fmt.Errorf("failed to reload configuration: %w", err)
353+
if err := g.removeServerConfiguration(ctx, serverName); err != nil {
354+
return nil, fmt.Errorf("failed to remove server configuration: %w", err)
357355
}
358356

359357
return &mcp.CallToolResult{
@@ -558,7 +556,7 @@ type configValue struct {
558556
}
559557

560558
// mcpConfigSetTool implements a tool for setting configuration values for MCP servers
561-
func (g *Gateway) createMcpConfigSetTool(configuration Configuration, clientConfig *clientConfig) *ToolRegistration {
559+
func (g *Gateway) createMcpConfigSetTool(clientConfig *clientConfig) *ToolRegistration {
562560
tool := &mcp.Tool{
563561
Name: "mcp-config-set",
564562
Description: "Set configuration values for MCP servers. Creates or updates server configuration with the specified key-value pairs.",
@@ -610,22 +608,22 @@ func (g *Gateway) createMcpConfigSetTool(configuration Configuration, clientConf
610608
configKey := strings.TrimSpace(params.Key)
611609

612610
// Check if server exists in catalog (optional check - we can configure servers that don't exist yet)
613-
_, _, serverExists := configuration.Find(serverName)
611+
_, _, serverExists := g.configuration.Find(serverName)
614612

615613
// Initialize the server's config map if it doesn't exist
616-
if configuration.config[serverName] == nil {
617-
configuration.config[serverName] = make(map[string]any)
614+
if g.configuration.config[serverName] == nil {
615+
g.configuration.config[serverName] = make(map[string]any)
618616
}
619617

620618
// Set the configuration value
621-
oldValue := configuration.config[serverName][configKey]
622-
configuration.config[serverName][configKey] = params.Value
619+
oldValue := g.configuration.config[serverName][configKey]
620+
g.configuration.config[serverName][configKey] = params.Value
623621

624622
// Log the configuration change
625623
log(fmt.Sprintf(" - Set config for server '%s': %s = %v", serverName, configKey, params.Value))
626624

627625
// Reload configuration with current server list to apply changes
628-
if err := g.reloadConfiguration(ctx, configuration, configuration.serverNames, clientConfig); err != nil {
626+
if err := g.reloadServerConfiguration(ctx, serverName, clientConfig); err != nil {
629627
return nil, fmt.Errorf("failed to reload configuration: %w", err)
630628
}
631629

0 commit comments

Comments
 (0)