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