1// server.go — MCP server process lifecycle management.
2//
3// Each configured MCP server is wrapped in a managedServer that tracks:
4// - Connection state (stopped / starting / ready)
5// - Last-access time (for idle timeout)
6// - Discovered tools
7//
8// Smart lazy: servers are spawned on startup to discover tools, then
9// killed after an idle period. On the next tool call, they're respawned.
10package main
11
12import (
13 "bufio"
14 "context"
15 "encoding/json"
16 "fmt"
17 "io"
18 "log"
19 "strings"
20 "sync"
21 "time"
22
23 "github.com/mark3labs/mcp-go/client"
24 "github.com/mark3labs/mcp-go/client/transport"
25 "github.com/mark3labs/mcp-go/mcp"
26)
27
28// serverState tracks the lifecycle of one MCP server.
29type serverState int
30
31const (
32 stateStopped serverState = iota // not running
33 stateStarting // spawning / initializing
34 stateReady // connected, tools discovered
35 stateError // failed to start
36)
37
38func (s serverState) String() string {
39 switch s {
40 case stateStopped:
41 return "stopped"
42 case stateStarting:
43 return "starting"
44 case stateReady:
45 return "ready"
46 case stateError:
47 return "error"
48 default:
49 return "unknown"
50 }
51}
52
53// managedServer wraps one MCP server with lifecycle management.
54type managedServer struct {
55 name string
56 config ServerConfig
57 logger *log.Logger
58
59 mu sync.Mutex
60 state serverState
61 client *client.Client
62 tools []mcp.Tool
63 lastUsed time.Time
64 startErr error
65 gen uint64 // bumped by stop(); a start attempt only commits if unchanged
66}
67
68// newManagedServer creates a new server wrapper.
69func newManagedServer(name string, cfg ServerConfig, logger *log.Logger) *managedServer {
70 return &managedServer{
71 name: name,
72 config: cfg,
73 logger: logger,
74 state: stateStopped,
75 lastUsed: time.Now(),
76 }
77}
78
79// start spawns the MCP server process and discovers its tools.
80// Safe to call concurrently; a stop() issued while a start is in
81// flight wins — the late result is discarded and its client closed.
82func (s *managedServer) start(ctx context.Context) error {
83 ctx, cancel := context.WithTimeout(ctx, time.Duration(s.config.ConnectTimeout)*time.Second)
84 defer cancel()
85
86 s.mu.Lock()
87 if s.state == stateReady {
88 s.lastUsed = time.Now()
89 s.mu.Unlock()
90 return nil
91 }
92 if s.state == stateStarting {
93 s.mu.Unlock()
94 return s.waitForReady(ctx)
95 }
96 s.state = stateStarting
97 s.startErr = nil
98 gen := s.gen
99 s.mu.Unlock()
100
101 c, tools, err := s.doStart(ctx)
102 return s.finishStart(gen, c, tools, err)
103}
104
105// finishStart commits the outcome of a start attempt, unless stop()
106// bumped the generation in the meantime.
107func (s *managedServer) finishStart(gen uint64, c *client.Client, tools []mcp.Tool, err error) error {
108 s.mu.Lock()
109 defer s.mu.Unlock()
110
111 if s.gen != gen {
112 if c != nil {
113 c.Close()
114 }
115 return fmt.Errorf("[%s] stopped while starting", s.name)
116 }
117 if err != nil {
118 s.state = stateError
119 s.startErr = err
120 s.logger.Printf("[%s] start failed: %v", s.name, err)
121 return err
122 }
123 s.client = c
124 s.tools = tools
125 s.state = stateReady
126 s.lastUsed = time.Now()
127 s.logger.Printf("[%s] ready with %d tools", s.name, len(s.tools))
128 return nil
129}
130
131// doStart performs the actual spawn + initialize + tools/list.
132func (s *managedServer) doStart(ctx context.Context) (*client.Client, []mcp.Tool, error) {
133 var c *client.Client
134 var err error
135
136 switch s.config.Transport {
137 case "stdio", "":
138 c, err = s.startStdio()
139 case "streamable-http":
140 c, err = s.startStreamableHTTP(ctx)
141 case "sse":
142 c, err = s.startSSE(ctx)
143 default:
144 return nil, nil, fmt.Errorf("unknown transport: %s", s.config.Transport)
145 }
146
147 if err != nil {
148 return nil, nil, err
149 }
150
151 // Initialize the MCP session
152 _, err = c.Initialize(ctx, mcp.InitializeRequest{
153 Params: mcp.InitializeParams{
154 ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION,
155 ClientInfo: mcp.Implementation{
156 Name: "zot-mcp-bridge",
157 Version: "1.0.0",
158 },
159 Capabilities: mcp.ClientCapabilities{},
160 },
161 })
162 if err != nil {
163 c.Close()
164 return nil, nil, fmt.Errorf("initialize: %w", err)
165 }
166
167 // Discover tools
168 result, err := c.ListTools(ctx, mcp.ListToolsRequest{})
169 if err != nil {
170 c.Close()
171 return nil, nil, fmt.Errorf("list tools: %w", err)
172 }
173
174 return c, result.Tools, nil
175}
176
177// startStdio spawns a stdio-based MCP server.
178func (s *managedServer) startStdio() (*client.Client, error) {
179 if s.config.Command == "" {
180 return nil, fmt.Errorf("stdio transport requires 'command' field")
181 }
182
183 // mcp-go merges this slice with os.Environ() before spawning the subprocess.
184 env := make([]string, 0, len(s.config.Env))
185 for k, v := range s.config.Env {
186 env = append(env, fmt.Sprintf("%s=%s", k, v))
187 }
188
189 // Create stdio client (this spawns the subprocess)
190 c, err := client.NewStdioMCPClient(s.config.Command, env, s.config.Args...)
191 if err != nil {
192 return nil, fmt.Errorf("spawn %s: %w", s.config.Command, err)
193 }
194
195 // Capture stderr for debugging
196 if stderr, ok := client.GetStderr(c); ok {
197 go s.pipeStderr(stderr)
198 }
199
200 return c, nil
201}
202
203// startStreamableHTTP connects to an HTTP-based MCP server.
204func (s *managedServer) startStreamableHTTP(ctx context.Context) (*client.Client, error) {
205 if s.config.URL == "" {
206 return nil, fmt.Errorf("streamable-http transport requires 'url' field")
207 }
208
209 // Build HTTP options
210 opts := []transport.StreamableHTTPCOption{}
211
212 // Streamable HTTP requires Accept: application/json, text/event-stream.
213 // This is a protocol detail, so set it automatically. User-provided
214 // headers are merged on top for auth/customization.
215 headers := map[string]string{
216 "Accept": "application/json, text/event-stream",
217 }
218 for k, v := range s.config.Headers {
219 headers[k] = v
220 }
221 opts = append(opts, transport.WithHTTPHeaders(headers))
222
223 // Create streamable HTTP client
224 c, err := client.NewStreamableHttpClient(s.config.URL, opts...)
225 if err != nil {
226 return nil, fmt.Errorf("create HTTP client: %w", err)
227 }
228
229 // Start the transport
230 if err := c.GetTransport().Start(ctx); err != nil {
231 return nil, fmt.Errorf("start HTTP transport: %w", err)
232 }
233
234 s.logger.Printf("[%s] connected to %s", s.name, s.config.URL)
235 return c, nil
236}
237
238// startSSE connects to an SSE-based MCP server (legacy transport).
239func (s *managedServer) startSSE(ctx context.Context) (*client.Client, error) {
240 if s.config.URL == "" {
241 return nil, fmt.Errorf("sse transport requires 'url' field")
242 }
243
244 // Create SSE client
245 c, err := client.NewSSEMCPClient(s.config.URL)
246 if err != nil {
247 return nil, fmt.Errorf("create SSE client: %w", err)
248 }
249
250 // Start the transport
251 if err := c.GetTransport().Start(ctx); err != nil {
252 return nil, fmt.Errorf("start SSE transport: %w", err)
253 }
254
255 s.logger.Printf("[%s] connected to SSE server at %s", s.name, s.config.URL)
256 return c, nil
257}
258
259// pipeStderr reads server stderr line by line and logs it.
260func (s *managedServer) pipeStderr(r io.Reader) {
261 sc := bufio.NewScanner(r)
262 sc.Buffer(make([]byte, 0, 64*1024), 1024*1024)
263 for sc.Scan() {
264 line := strings.TrimSpace(sc.Text())
265 if line != "" {
266 s.logger.Printf("[%s:stderr] %s", s.name, line)
267 }
268 }
269}
270
271// stop gracefully shuts down the server.
272func (s *managedServer) stop() {
273 s.mu.Lock()
274 defer s.mu.Unlock()
275
276 s.gen++ // invalidate any start attempt still in flight
277 if s.client != nil {
278 s.client.Close()
279 s.client = nil
280 }
281 if s.state != stateError {
282 s.state = stateStopped
283 }
284}
285
286// callTool forwards a tool call to the MCP server.
287// If the server is not running, it starts it first.
288func (s *managedServer) callTool(ctx context.Context, toolName string, args json.RawMessage) (*mcp.CallToolResult, error) {
289 s.mu.Lock()
290 c := s.client
291 st := s.state
292 s.mu.Unlock()
293
294 // If not ready, start the server
295 if st != stateReady || c == nil {
296 if err := s.start(ctx); err != nil {
297 return nil, err
298 }
299 s.mu.Lock()
300 c = s.client
301 s.mu.Unlock()
302 if c == nil {
303 return nil, fmt.Errorf("[%s] stopped before the call could run", s.name)
304 }
305 }
306
307 // Parse args into map[string]any
308 var argsMap map[string]any
309 if len(args) > 0 {
310 if err := json.Unmarshal(args, &argsMap); err != nil {
311 return nil, fmt.Errorf("invalid args: %w", err)
312 }
313 }
314
315 // Call the tool
316 callCtx, cancel := context.WithTimeout(ctx, time.Duration(s.config.RequestTimeout)*time.Second)
317 defer cancel()
318
319 result, err := c.CallTool(callCtx, mcp.CallToolRequest{
320 Params: mcp.CallToolParams{
321 Name: toolName,
322 Arguments: argsMap,
323 },
324 })
325 if err != nil {
326 return nil, err
327 }
328
329 // Update last-used time
330 s.mu.Lock()
331 s.lastUsed = time.Now()
332 s.mu.Unlock()
333
334 return result, nil
335}
336
337// waitForReady blocks until the server is ready or the context
338// expires. Polling is a deliberate simplification for this example;
339// a channel closed by finishStart would avoid the 100ms latency.
340func (s *managedServer) waitForReady(ctx context.Context) error {
341 ticker := time.NewTicker(100 * time.Millisecond)
342 defer ticker.Stop()
343 for {
344 select {
345 case <-ctx.Done():
346 return ctx.Err()
347 case <-ticker.C:
348 s.mu.Lock()
349 st := s.state
350 err := s.startErr
351 s.mu.Unlock()
352 if st == stateReady {
353 return nil
354 }
355 if st == stateError {
356 return err
357 }
358 }
359 }
360}
361
362// isIdle returns true if the server hasn't been used within the timeout.
363func (s *managedServer) isIdle(timeout time.Duration) bool {
364 s.mu.Lock()
365 defer s.mu.Unlock()
366 if s.state != stateReady {
367 return false
368 }
369 return time.Since(s.lastUsed) > timeout
370}
371
372// status returns a compact human-readable status string.
373func (s *managedServer) status() string {
374 s.mu.Lock()
375 defer s.mu.Unlock()
376
377 switch s.state {
378 case stateReady:
379 return compactToolStatus(s.name, len(s.tools))
380 case stateError:
381 if s.startErr != nil {
382 return fmt.Sprintf("%s (%s)", s.name, compactErr(s.startErr.Error()))
383 }
384 return fmt.Sprintf("%s (error)", s.name)
385 case stateStopped:
386 if len(s.tools) > 0 {
387 return fmt.Sprintf("%s (sleeping, %s)", s.name, toolCountText(len(s.tools)))
388 }
389 return fmt.Sprintf("%s (sleeping)", s.name)
390 case stateStarting:
391 return fmt.Sprintf("%s (starting)", s.name)
392 default:
393 return fmt.Sprintf("%s (%s)", s.name, s.state)
394 }
395}
396
397// detailStatus returns a multi-line status for one server.
398func (s *managedServer) detailStatus() string {
399 s.mu.Lock()
400 name := s.name
401 transport := s.config.Transport
402 url := s.config.URL
403 command := s.config.Command
404 args := append([]string(nil), s.config.Args...)
405 state := s.state
406 toolCount := len(s.tools)
407 idleTimeout := s.config.IdleTimeout
408 requestTimeout := s.config.RequestTimeout
409 connectTimeout := s.config.ConnectTimeout
410 lastUsed := s.lastUsed
411 startErr := s.startErr
412 s.mu.Unlock()
413
414 if transport == "" {
415 transport = "stdio"
416 }
417
418 var sb strings.Builder
419 sb.WriteString("MCP server: ")
420 sb.WriteString(name)
421 sb.WriteByte('\n')
422 sb.WriteString(" status: ")
423 sb.WriteString(state.String())
424 if state == stateError && startErr != nil {
425 sb.WriteString(" (")
426 sb.WriteString(compactErr(startErr.Error()))
427 sb.WriteByte(')')
428 }
429 sb.WriteByte('\n')
430 sb.WriteString(" tools: ")
431 sb.WriteString(toolCountText(toolCount))
432 sb.WriteByte('\n')
433 sb.WriteString(" transport: ")
434 sb.WriteString(transport)
435 sb.WriteByte('\n')
436 if url != "" {
437 sb.WriteString(" url: ")
438 sb.WriteString(url)
439 sb.WriteByte('\n')
440 }
441 if command != "" {
442 sb.WriteString(" command: ")
443 sb.WriteString(command)
444 if len(args) > 0 {
445 sb.WriteByte(' ')
446 sb.WriteString(strings.Join(args, " "))
447 }
448 sb.WriteByte('\n')
449 }
450 sb.WriteString(fmt.Sprintf(" timeouts: connect=%ds request=%ds idle=%ds\n", connectTimeout, requestTimeout, idleTimeout))
451 if !lastUsed.IsZero() {
452 sb.WriteString(" last used: ")
453 sb.WriteString(time.Since(lastUsed).Round(time.Second).String())
454 sb.WriteString(" ago\n")
455 }
456 if state == stateError && startErr != nil {
457 sb.WriteString(" error: ")
458 sb.WriteString(firstLine(startErr.Error()))
459 sb.WriteByte('\n')
460 }
461 return strings.TrimRight(sb.String(), "\n")
462}
463
464func compactToolStatus(name string, n int) string {
465 if n <= 0 {
466 return name
467 }
468 return fmt.Sprintf("%s (%s)", name, toolCountText(n))
469}
470
471func toolCountText(n int) string {
472 if n == 1 {
473 return "1 tool"
474 }
475 return fmt.Sprintf("%d tools", n)
476}
477
478func firstLine(s string) string {
479 s = strings.TrimSpace(s)
480 if i := strings.IndexAny(s, "\n\r"); i >= 0 {
481 s = s[:i]
482 }
483 return s
484}
485
486func compactErr(s string) string {
487 s = firstLine(s)
488 if s == "" {
489 return "error"
490 }
491 lower := strings.ToLower(s)
492 switch {
493 case strings.Contains(lower, "authorization") || strings.Contains(lower, "unauthorized") || strings.Contains(lower, "401") || strings.Contains(lower, "403"):
494 return "auth failed"
495 case strings.Contains(lower, "transport") || strings.Contains(lower, "connection") || strings.Contains(lower, "connect"):
496 return "connection failed"
497 case strings.Contains(lower, "context deadline exceeded") || strings.Contains(lower, "timeout") || strings.Contains(lower, "timed out"):
498 return "timeout"
499 case strings.Contains(lower, "no such file") || strings.Contains(lower, "executable file not found"):
500 return "command not found"
501 case strings.HasPrefix(lower, "initialize:"):
502 return "initialize failed"
503 case strings.HasPrefix(lower, "list tools:"):
504 return "tool discovery failed"
505 }
506 const max = 44
507 if r := []rune(s); len(r) > max {
508 s = string(r[:max-1]) + "…"
509 }
510 return s
511}