6
6
"net"
7
7
"os"
8
8
"strings"
9
- "sync"
10
9
"time"
11
10
12
11
"github.com/mark3labs/mcp-go/server"
@@ -90,34 +89,14 @@ func (g *Gateway) Run(ctx context.Context) error {
90
89
}
91
90
toolCallbacks := interceptors .Callbacks (g .LogCalls , g .BlockSecrets , customInterceptors )
92
91
93
- // TODO: cleanup stopped servers. That happens in stdio over TCP mode.
94
- var (
95
- lock sync. Mutex
96
- changeListeners [] func ( * Capabilities )
92
+ mcpServer := server . NewMCPServer (
93
+ "Docker AI MCP Gateway" ,
94
+ "2.0.1" ,
95
+ server . WithToolHandlerMiddleware ( toolCallbacks ),
97
96
)
98
97
99
- capabilities , err := g .listServersAndCapabilities (ctx , configuration )
100
- if err != nil {
101
- return fmt .Errorf ("listing capabilities: %w" , err )
102
- }
103
-
104
- newMCPServer := func () * server.MCPServer {
105
- mcpServer := server .NewMCPServer (
106
- "Docker AI MCP Gateway" ,
107
- "2.0.1" ,
108
- server .WithToolHandlerMiddleware (toolCallbacks ),
109
- )
110
-
111
- // TODO: This will create a new server instance with an outdated vision of the capabilities.
112
- refreshCapabilities (mcpServer , capabilities )
113
-
114
- lock .Lock ()
115
- changeListeners = append (changeListeners , func (newCapabilities * Capabilities ) {
116
- refreshCapabilities (mcpServer , newCapabilities )
117
- })
118
- lock .Unlock ()
119
-
120
- return mcpServer
98
+ if err := g .reloadConfiguration (ctx , mcpServer , configuration ); err != nil {
99
+ return fmt .Errorf ("loading configuration: %w" , err )
121
100
}
122
101
123
102
// Optionally watch for configuration updates.
@@ -137,19 +116,10 @@ func (g *Gateway) Run(ctx context.Context) error {
137
116
continue
138
117
}
139
118
140
- capabilities , err := g .listServersAndCapabilities (ctx , configuration )
141
- if err != nil {
119
+ if err := g .reloadConfiguration (ctx , mcpServer , configuration ); err != nil {
142
120
logf ("> Unable to list capabilities: %s" , err )
143
121
continue
144
122
}
145
-
146
- g .health .SetUnhealthy ()
147
- lock .Lock ()
148
- for _ , listener := range changeListeners {
149
- listener (capabilities )
150
- }
151
- lock .Unlock ()
152
- g .health .SetHealthy ()
153
123
}
154
124
}
155
125
}()
@@ -162,31 +132,25 @@ func (g *Gateway) Run(ctx context.Context) error {
162
132
}
163
133
164
134
// Start the server
165
- g .health .SetHealthy ()
166
135
switch strings .ToLower (g .Transport ) {
167
136
case "stdio" :
168
- if g .Port == 0 {
169
- log ("> Start stdio server" )
170
- return g .startStdioServer (ctx , newMCPServer , os .Stdin , os .Stdout )
171
- }
172
-
173
- log ("> Start stdio over TCP server on port" , g .Port )
174
- return g .startStdioOverTCPServer (ctx , newMCPServer , ln )
137
+ log ("> Start stdio server" )
138
+ return g .startStdioServer (ctx , mcpServer , os .Stdin , os .Stdout )
175
139
176
140
case "sse" :
177
141
log ("> Start sse server on port" , g .Port )
178
- return g .startSseServer (ctx , newMCPServer , ln )
142
+ return g .startSseServer (ctx , mcpServer , ln )
179
143
180
144
case "streaming" :
181
145
log ("> Start streaming server on port" , g .Port )
182
- return g .startStreamingServer (ctx , newMCPServer , ln )
146
+ return g .startStreamingServer (ctx , mcpServer , ln )
183
147
184
148
default :
185
149
return fmt .Errorf ("unknown transport %q, expected 'stdio', 'sse' or 'streaming" , g .Transport )
186
150
}
187
151
}
188
152
189
- func (g * Gateway ) listServersAndCapabilities (ctx context.Context , configuration Configuration ) ( * Capabilities , error ) {
153
+ func (g * Gateway ) reloadConfiguration (ctx context.Context , mcpServer * server. MCPServer , configuration Configuration ) error {
190
154
// Which servers are enabled in the registry.yaml?
191
155
serverNames := configuration .ServerNames ()
192
156
if len (serverNames ) == 0 {
@@ -200,19 +164,20 @@ func (g *Gateway) listServersAndCapabilities(ctx context.Context, configuration
200
164
log ("- Listing MCP tools..." )
201
165
capabilities , err := g .listCapabilities (ctx , configuration , serverNames )
202
166
if err != nil {
203
- return nil , fmt .Errorf ("listing resources: %w" , err )
167
+ return fmt .Errorf ("listing resources: %w" , err )
204
168
}
205
169
log (">" , len (capabilities .Tools ), "tools listed in" , time .Since (startList ))
206
170
207
- return capabilities , nil
208
- }
209
-
210
- func refreshCapabilities (s * server.MCPServer , c * Capabilities ) {
211
- s .SetTools (c .Tools ... )
212
- s .SetPrompts (c .Prompts ... )
213
- s .SetResources (c .Resources ... )
214
- s .RemoveAllResourceTemplates ()
215
- for _ , v := range c .ResourceTemplates {
216
- s .AddResourceTemplate (v .ResourceTemplate , v .Handler )
171
+ // Update the server's capabilities.
172
+ g .health .SetUnhealthy ()
173
+ mcpServer .SetTools (capabilities .Tools ... )
174
+ mcpServer .SetPrompts (capabilities .Prompts ... )
175
+ mcpServer .SetResources (capabilities .Resources ... )
176
+ mcpServer .RemoveAllResourceTemplates ()
177
+ for _ , v := range capabilities .ResourceTemplates {
178
+ mcpServer .AddResourceTemplate (v .ResourceTemplate , v .Handler )
217
179
}
180
+ g .health .SetHealthy ()
181
+
182
+ return nil
218
183
}
0 commit comments