1package mcp
2
3import (
4 "context"
5 "errors"
6 "fmt"
7 "net/http"
8 "sort"
9 "strings"
10 "sync"
11)
12
13// ServerStatus is a snapshot of one configured server.
14type ServerStatus struct {
15 Name string
16 Transport string
17 Connected bool
18 Tools int
19 Resources bool
20 Err string
21 Info ServerInfo
22}
23
24type serverEntry struct {
25 name string
26 cfg ServerConfig
27
28 // connectMu serializes connection attempts; mu guards the fields below
29 // and is never held across network or process I/O, so Status stays fast.
30 connectMu sync.Mutex
31 mu sync.Mutex
32 client *Client
33 tools []Tool
34 err error
35 refresher *toolRefresher
36}
37
38// Manager owns the MCP connections of one zut process. It is safe for
39// concurrent use; resident subagents share the primary's connections.
40type Manager struct {
41 cwd string
42 info ClientInfo
43 httpClient *http.Client
44
45 servers map[string]*serverEntry
46 order []string
47
48 mu sync.Mutex
49 closed bool
50}
51
52// NewManager builds a manager for an already loaded configuration.
53func NewManager(cfg Config, cwd string, info ClientInfo, httpClient *http.Client) *Manager {
54 m := &Manager{cwd: cwd, info: info, httpClient: httpClient, servers: map[string]*serverEntry{}}
55 for _, name := range sortedNames(cfg.Servers) {
56 m.servers[name] = &serverEntry{name: name, cfg: cfg.Servers[name]}
57 m.order = append(m.order, name)
58 }
59 return m
60}
61
62// Empty reports whether no servers are configured.
63func (m *Manager) Empty() bool { return m == nil || len(m.order) == 0 }
64
65// Start connects every server concurrently and lists its tools. Each server
66// is bounded by its own startup timeout. Failures are returned per server and
67// leave the server disconnected; a later tool call retries the connection.
68// The returned error is non-nil only when a server marked required failed.
69func (m *Manager) Start(ctx context.Context) (failures []error, requiredErr error) {
70 if m.Empty() {
71 return nil, nil
72 }
73 var wg sync.WaitGroup
74 errs := make([]error, len(m.order))
75 for i, name := range m.order {
76 wg.Add(1)
77 go func(i int, s *serverEntry) {
78 defer wg.Done()
79 if _, err := m.connect(ctx, s); err != nil {
80 errs[i] = err
81 }
82 }(i, m.servers[name])
83 }
84 wg.Wait()
85 var required []error
86 for i, err := range errs {
87 if err == nil {
88 continue
89 }
90 failures = append(failures, err)
91 if m.servers[m.order[i]].cfg.Required {
92 required = append(required, err)
93 }
94 }
95 return failures, errors.Join(required...)
96}
97
98// connect returns a live client for s, connecting (or reconnecting after a
99// dropped connection) when needed.
100func (m *Manager) connect(ctx context.Context, s *serverEntry) (*Client, error) {
101 if c := s.live(); c != nil {
102 return c, nil
103 }
104 s.connectMu.Lock()
105 defer s.connectMu.Unlock()
106 if c := s.live(); c != nil {
107 return c, nil // another caller reconnected while we waited
108 }
109 s.mu.Lock()
110 stale := s.client
111 s.client = nil
112 staleRefresher := s.refresher
113 s.refresher = nil
114 s.mu.Unlock()
115 if staleRefresher != nil {
116 staleRefresher.stop()
117 }
118 if stale != nil {
119 _ = stale.Close()
120 }
121 m.mu.Lock()
122 closed := m.closed
123 m.mu.Unlock()
124 if closed {
125 return nil, fmt.Errorf("mcp server %q: manager closed", s.name)
126 }
127 startCtx, cancel := withTimeout(ctx, s.cfg.StartupTimeout())
128 defer cancel()
129 // The refresher exists before the handshake so a list_changed notification
130 // that arrives while the first listing is in flight is remembered, and
131 // the listing is repeated once the client is published.
132 refresher := newToolRefresher()
133 client, err := connectWith(startCtx, s.name, s.cfg, m.cwd, m.info, m.httpClient, refresher.notify)
134 var tools []Tool
135 if err == nil {
136 tools, err = client.ListTools(startCtx)
137 if err != nil {
138 _ = client.Close()
139 }
140 }
141 if err != nil {
142 refresher.stop()
143 if errors.Is(err, context.DeadlineExceeded) && ctx.Err() == nil {
144 err = fmt.Errorf("timed out after %s; raise startupTimeoutSec if the server is slow to start", s.cfg.StartupTimeout())
145 }
146 err = fmt.Errorf("mcp server %q: %w", s.name, err)
147 s.mu.Lock()
148 s.err = err
149 s.mu.Unlock()
150 return nil, err
151 }
152 // Publish under m.mu so Close either sees this client or we see closed.
153 m.mu.Lock()
154 if m.closed {
155 m.mu.Unlock()
156 // Close ran during the handshake; do not leak the new process.
157 refresher.stop()
158 _ = client.Close()
159 return nil, fmt.Errorf("mcp server %q: manager closed", s.name)
160 }
161 s.mu.Lock()
162 s.client = client
163 s.tools = filterTools(s.cfg, tools)
164 s.err = nil
165 old := s.refresher
166 s.refresher = refresher
167 s.mu.Unlock()
168 m.mu.Unlock()
169 if old != nil {
170 old.stop()
171 }
172 go func() {
173 <-client.Done()
174 refresher.stop()
175 }()
176 refresher.start(func(ctx context.Context) { m.refreshTools(ctx, s, client) })
177 return client, nil
178}
179
180// live returns the current client when its connection is still usable.
181func (s *serverEntry) live() *Client {
182 s.mu.Lock()
183 defer s.mu.Unlock()
184 if s.client == nil {
185 return nil
186 }
187 select {
188 case <-s.client.Done():
189 return nil
190 default:
191 return s.client
192 }
193}
194
195// refreshTools re-lists tools after notifications/tools/list_changed. New
196// definitions reach the model the next time zut builds a tool registry. It is
197// only called from the server's refresher, which serializes and coalesces
198// calls.
199func (m *Manager) refreshTools(ctx context.Context, s *serverEntry, client *Client) {
200 ctx, cancel := withTimeout(ctx, s.cfg.StartupTimeout())
201 defer cancel()
202 tools, err := client.ListTools(ctx)
203 if err != nil {
204 return
205 }
206 s.mu.Lock()
207 if s.client == client {
208 s.tools = filterTools(s.cfg, tools)
209 }
210 s.mu.Unlock()
211}
212
213// toolRefresher coalesces list_changed notifications for one connection into
214// serialized re-listings. Any number of notifications that arrive while a
215// listing runs cause exactly one more listing afterwards. It does nothing
216// until start installs the listing function, so notifications that arrive
217// during the initial handshake are remembered rather than raced.
218type toolRefresher struct {
219 mu sync.Mutex
220 dirty bool
221 running bool
222 stopped bool
223 list func(ctx context.Context)
224 ctx context.Context
225 cancel context.CancelFunc
226}
227
228func newToolRefresher() *toolRefresher {
229 ctx, cancel := context.WithCancel(context.Background())
230 return &toolRefresher{ctx: ctx, cancel: cancel}
231}
232
233// notify records that the server's tool list changed.
234func (r *toolRefresher) notify() {
235 r.mu.Lock()
236 r.dirty = true
237 r.spawnLocked()
238 r.mu.Unlock()
239}
240
241// start enables listing with fn and runs it once if a notification arrived
242// earlier.
243func (r *toolRefresher) start(fn func(ctx context.Context)) {
244 r.mu.Lock()
245 r.list = fn
246 r.spawnLocked()
247 r.mu.Unlock()
248}
249
250// stop cancels any running listing and ignores later notifications.
251func (r *toolRefresher) stop() {
252 r.mu.Lock()
253 r.stopped = true
254 r.dirty = false
255 r.mu.Unlock()
256 r.cancel()
257}
258
259func (r *toolRefresher) spawnLocked() {
260 if r.stopped || r.running || !r.dirty || r.list == nil {
261 return
262 }
263 r.running = true
264 go r.run()
265}
266
267func (r *toolRefresher) run() {
268 for {
269 r.mu.Lock()
270 if r.stopped || !r.dirty {
271 r.running = false
272 r.mu.Unlock()
273 return
274 }
275 r.dirty = false
276 list := r.list
277 r.mu.Unlock()
278 list(r.ctx)
279 }
280}
281
282func filterTools(cfg ServerConfig, tools []Tool) []Tool {
283 out := make([]Tool, 0, len(tools))
284 seen := map[string]bool{}
285 for _, tool := range tools {
286 if tool.Name == "" || seen[tool.Name] || !cfg.ToolAllowed(tool.Name) {
287 continue
288 }
289 seen[tool.Name] = true
290 out = append(out, tool)
291 }
292 return out
293}
294
295// client returns a live client for server, reconnecting when necessary.
296func (m *Manager) client(ctx context.Context, server string) (*Client, ServerConfig, error) {
297 if m == nil {
298 return nil, ServerConfig{}, errors.New("mcp is not enabled")
299 }
300 s, ok := m.servers[server]
301 if !ok {
302 return nil, ServerConfig{}, fmt.Errorf("unknown mcp server %q", server)
303 }
304 c, err := m.connect(ctx, s)
305 return c, s.cfg, err
306}
307
308// Servers returns configured server names in stable order.
309func (m *Manager) Servers() []string {
310 if m == nil {
311 return nil
312 }
313 return append([]string(nil), m.order...)
314}
315
316// Status returns a snapshot of every server.
317func (m *Manager) Status() []ServerStatus {
318 if m == nil {
319 return nil
320 }
321 out := make([]ServerStatus, 0, len(m.order))
322 for _, name := range m.order {
323 s := m.servers[name]
324 s.mu.Lock()
325 st := ServerStatus{Name: name, Transport: s.cfg.Kind(), Tools: len(s.tools)}
326 if s.client != nil {
327 select {
328 case <-s.client.Done():
329 default:
330 st.Connected = true
331 }
332 st.Info = s.client.Info()
333 st.Resources = st.Info.HasResources
334 }
335 if s.err != nil {
336 st.Err = s.err.Error()
337 }
338 s.mu.Unlock()
339 out = append(out, st)
340 }
341 return out
342}
343
344// ServerTools returns the known tool definitions of one server.
345func (m *Manager) ServerTools(server string) []Tool {
346 if m == nil {
347 return nil
348 }
349 s, ok := m.servers[server]
350 if !ok {
351 return nil
352 }
353 s.mu.Lock()
354 defer s.mu.Unlock()
355 return append([]Tool(nil), s.tools...)
356}
357
358// hasResources reports whether any connected server advertises resources.
359func (m *Manager) hasResources() bool {
360 for _, st := range m.Status() {
361 if st.Resources {
362 return true
363 }
364 }
365 return false
366}
367
368// Close stops every connection. It is idempotent.
369func (m *Manager) Close() error {
370 if m == nil {
371 return nil
372 }
373 m.mu.Lock()
374 if m.closed {
375 m.mu.Unlock()
376 return nil
377 }
378 m.closed = true
379 m.mu.Unlock()
380 var wg sync.WaitGroup
381 for _, name := range m.order {
382 s := m.servers[name]
383 wg.Add(1)
384 go func() {
385 defer wg.Done()
386 s.mu.Lock()
387 client := s.client
388 s.client = nil
389 refresher := s.refresher
390 s.refresher = nil
391 s.mu.Unlock()
392 if refresher != nil {
393 refresher.stop()
394 }
395 if client != nil {
396 _ = client.Close()
397 }
398 }()
399 }
400 wg.Wait()
401 return nil
402}
403
404// FormatStatus renders a short human-readable status table. Server-derived
405// text (names, errors) is stripped of terminal control sequences.
406func FormatStatus(statuses []ServerStatus) string {
407 if len(statuses) == 0 {
408 return "no MCP servers configured"
409 }
410 sort.SliceStable(statuses, func(i, j int) bool { return statuses[i].Name < statuses[j].Name })
411 var b strings.Builder
412 for _, st := range statuses {
413 state := "disconnected"
414 if st.Connected {
415 state = "connected"
416 }
417 fmt.Fprintf(&b, "%s (%s): %s, %d tools", sanitizeText(st.Name, false), sanitizeText(st.Transport, false), state, st.Tools)
418 if st.Resources {
419 b.WriteString(", resources")
420 }
421 if st.Err != "" {
422 b.WriteString("\n error: " + sanitizeText(firstLine(st.Err), false))
423 }
424 b.WriteByte('\n')
425 }
426 return strings.TrimRight(b.String(), "\n")
427}
428
429func firstLine(s string) string {
430 line, _, _ := strings.Cut(s, "\n")
431 return line
432}