1package mcp
2
3import (
4 "bufio"
5 "context"
6 "encoding/json"
7 "errors"
8 "fmt"
9 "io"
10 "os"
11 "os/exec"
12 "path/filepath"
13 "runtime"
14 "strings"
15 "sync"
16 "sync/atomic"
17 "time"
18 "unicode/utf8"
19)
20
21// maxStdioFrame bounds one newline-delimited JSON-RPC message from a server.
22const maxStdioFrame = 64 << 20
23
24const (
25 // stderrTailLines is how many stderr lines are kept for diagnostics.
26 stderrTailLines = 20
27 // maxStderrLine truncates one retained stderr line; the rest of a longer
28 // line is still drained from the pipe, just not kept.
29 maxStderrLine = 4096
30 // maxStderrOverlap caps how far past maxStderrLine raw bytes are kept so
31 // a straddling secret can still be redacted.
32 maxStderrOverlap = 64 << 10
33 // writeQueueSize bounds frames waiting for the writer goroutine.
34 writeQueueSize = 64
35 // responseWriteTimeout bounds answering one server-initiated request.
36 responseWriteTimeout = 30 * time.Second
37)
38
39// defaultEnvVars are inherited by stdio servers. Everything else from the host
40// environment, including provider API keys, is withheld unless listed in
41// envVars or set explicitly in env.
42var defaultEnvVars = func() []string {
43 if runtime.GOOS == "windows" {
44 return []string{
45 "APPDATA", "COMPUTERNAME", "COMSPEC", "HOMEDRIVE", "HOMEPATH", "LOCALAPPDATA",
46 "NUMBER_OF_PROCESSORS", "OS", "PATH", "PATHEXT", "PROCESSOR_ARCHITECTURE",
47 "PROGRAMDATA", "PROGRAMFILES", "PROGRAMFILES(X86)", "SYSTEMDRIVE", "SYSTEMROOT",
48 "TEMP", "TMP", "USERDOMAIN", "USERNAME", "USERPROFILE", "WINDIR",
49 }
50 }
51 return []string{
52 "HOME", "LOGNAME", "PATH", "SHELL", "USER", "LANG", "LC_ALL", "LC_CTYPE",
53 "TERM", "TMPDIR", "TZ", "XDG_RUNTIME_DIR", "XDG_CONFIG_HOME", "XDG_CACHE_HOME", "XDG_DATA_HOME",
54 "SSL_CERT_FILE", "SSL_CERT_DIR", "NODE_EXTRA_CA_CERTS",
55 "HTTP_PROXY", "HTTPS_PROXY", "NO_PROXY", "http_proxy", "https_proxy", "no_proxy",
56 }
57}()
58
59// beforeReapHook, when set (tests), runs after the group was cleaned and
60// before the leader is reaped.
61var beforeReapHook atomic.Pointer[func(pid int)]
62
63// writeReq is one frame queued for the writer goroutine.
64type writeReq struct {
65 ctx context.Context
66 frame []byte
67 res chan error
68 state atomic.Int32 // reqQueued, reqWriting, or reqSkipped
69}
70
71// A queued frame is claimed exactly once: by the writer (it will be written)
72// or by its caller (it will never be written).
73const (
74 reqQueued int32 = iota
75 reqWriting
76 reqSkipped
77)
78
79type stdioTransport struct {
80 name string
81 cmd *exec.Cmd
82 guard *processGuard
83 red *redactor
84
85 // ctx spans the transport's lifetime; it is cancelled when the
86 // connection fails or closes and bounds every background send.
87 ctx context.Context
88 cancel context.CancelFunc
89
90 stdin io.WriteCloser
91 outR *os.File
92 errR *os.File
93 writeQ chan *writeReq
94
95 mu sync.Mutex
96 pending map[string]chan message
97 handler incomingHandler
98 err error
99
100 stderrMu sync.Mutex
101 stderr []string
102
103 doneCh chan struct{}
104 closeOnce sync.Once
105 exited chan struct{}
106 readDone chan struct{}
107 stderrDone chan struct{}
108
109 // cancelsSent counts notifications/cancelled sends started (diagnostics
110 // and tests).
111 cancelsSent atomic.Int64
112 // Test hooks for the writer (nil in production): hookDequeued runs when the
113 // writer takes a frame off the queue, hookSkipped after it claimed a frame
114 // whose context had ended and before it reports that to the caller.
115 hookDequeued func()
116 hookSkipped func()
117 // writesStarted counts frames the writer has begun writing (diagnostics
118 // and tests).
119 writesStarted atomic.Int64
120}
121
122// buildStdioEnv returns the restricted child environment.
123func buildStdioEnv(cfg ServerConfig) []string {
124 env := map[string]string{}
125 for _, key := range append(append([]string(nil), defaultEnvVars...), cfg.EnvVars...) {
126 if v, ok := os.LookupEnv(key); ok {
127 env[key] = v
128 }
129 }
130 for k, v := range cfg.Env {
131 env[k] = v
132 }
133 out := make([]string, 0, len(env))
134 for _, k := range sortedNames(env) {
135 out = append(out, k+"="+env[k])
136 }
137 return out
138}
139
140func startStdio(name string, cfg ServerConfig, cwd string) (*stdioTransport, error) {
141 cmd := exec.Command(cfg.Command, cfg.Args...)
142 cmd.Env = buildStdioEnv(cfg)
143 dir := cwd
144 if cfg.CWD != "" {
145 dir = cfg.CWD
146 if !filepath.IsAbs(dir) && cwd != "" {
147 dir = filepath.Join(cwd, dir)
148 }
149 }
150 cmd.Dir = dir
151 isolateProcess(cmd)
152 stdin, err := cmd.StdinPipe()
153 if err != nil {
154 return nil, err
155 }
156 // Plain OS pipes instead of StdoutPipe/StderrPipe: Wait closes those as
157 // soon as the child exits, which can discard output the child wrote just
158 // before exiting (a final response, a fatal stderr line).
159 outR, outW, err := os.Pipe()
160 if err != nil {
161 return nil, err
162 }
163 errR, errW, err := os.Pipe()
164 if err != nil {
165 outR.Close()
166 outW.Close()
167 return nil, err
168 }
169 cmd.Stdout = outW
170 cmd.Stderr = errW
171 red := newRedactor(cfg)
172 if err := cmd.Start(); err != nil {
173 outR.Close()
174 outW.Close()
175 errR.Close()
176 errW.Close()
177 return nil, fmt.Errorf("start %s: %w", cfg.Command, red.scrubError(err))
178 }
179 outW.Close()
180 errW.Close()
181 guard, err := newProcessGuard(cmd)
182 if err != nil {
183 // Fail closed: never run a server that cannot be contained.
184 _ = stdin.Close()
185 _ = cmd.Process.Kill()
186 _ = cmd.Wait()
187 outR.Close()
188 errR.Close()
189 return nil, fmt.Errorf("start %s: %w", cfg.Command, red.scrubError(err))
190 }
191 ctx, cancel := context.WithCancel(context.Background())
192 t := &stdioTransport{
193 name: name,
194 cmd: cmd,
195 guard: guard,
196 red: red,
197 ctx: ctx,
198 cancel: cancel,
199 stdin: stdin,
200 outR: outR,
201 errR: errR,
202 writeQ: make(chan *writeReq, writeQueueSize),
203 pending: map[string]chan message{},
204 doneCh: make(chan struct{}),
205 exited: make(chan struct{}),
206 readDone: make(chan struct{}),
207 stderrDone: make(chan struct{}),
208 }
209 go t.writeLoop()
210 go t.readStderr(errR)
211 go t.readLoop(outR)
212 go func() {
213 // Wait for the leader to exit WITHOUT reaping it, clean the group
214 // while the leader is a zombie (its pid, the pgid, cannot be reused
215 // yet), and only then reap it.
216 leaderGone := t.guard.awaitLeaderExit()
217 if leaderGone {
218 t.guard.cleanup()
219 if h := beforeReapHook.Load(); h != nil {
220 (*h)(cmd.Process.Pid)
221 }
222 }
223 waitErr := cmd.Wait()
224 close(t.exited)
225 if waitErr == nil {
226 waitErr = errors.New("process exited")
227 }
228 // Let the reader hand over whatever the server wrote before exiting,
229 // so a response sitting in the pipe wins over the failure.
230 select {
231 case <-t.readDone:
232 case <-time.After(250 * time.Millisecond):
233 }
234 t.fail(fmt.Errorf("server process ended: %w", waitErr))
235 if !leaderGone {
236 // No wait-without-reap on this platform: clean up after the fact.
237 t.guard.cleanup()
238 }
239 }()
240 return t, nil
241}
242
243// readStderr drains the child's stderr for its whole life, keeping a bounded
244// tail of bounded lines. It never stops early: a child blocked on a full
245// stderr pipe would hang.
246//
247// Redaction happens on the complete line before it is shortened for display.
248// For a line longer than the display limit the raw prefix is kept
249// maxSecretLen bytes past the cut, so a secret that straddles the cut is
250// matched whole; what remains of one anyway is trimmed afterwards.
251func (t *stdioTransport) readStderr(r io.Reader) {
252 defer close(t.stderrDone)
253 reader := bufio.NewReaderSize(r, 4096)
254 keep := maxStderrLine + min(t.red.maxLen, maxStderrOverlap)
255 var line []byte
256 truncated := false
257 flush := func() {
258 text := t.red.scrub(strings.TrimRight(string(line), "\r"))
259 if len(text) > maxStderrLine || truncated {
260 if len(text) > maxStderrLine {
261 cut := maxStderrLine
262 for cut > 0 && !utf8.RuneStart(text[cut]) {
263 cut--
264 }
265 text = text[:cut]
266 }
267 text = t.red.trimPartialSecret(text) + "…[line truncated]"
268 }
269 line, truncated = nil, false
270 t.stderrMu.Lock()
271 t.stderr = append(t.stderr, text)
272 if len(t.stderr) > stderrTailLines {
273 t.stderr = t.stderr[len(t.stderr)-stderrTailLines:]
274 }
275 t.stderrMu.Unlock()
276 }
277 for {
278 chunk, isPrefix, err := reader.ReadLine()
279 if room := keep - len(line); len(chunk) > room {
280 if room > 0 {
281 line = append(line, chunk[:room]...)
282 }
283 truncated = true
284 } else {
285 line = append(line, chunk...)
286 }
287 if err != nil {
288 if len(line) > 0 || truncated {
289 flush()
290 }
291 return
292 }
293 if !isPrefix {
294 flush()
295 }
296 }
297}
298
299// stderrTail returns recent stderr output, scrubbed for display.
300func (t *stdioTransport) stderrTail() string {
301 t.stderrMu.Lock()
302 lines := append([]string(nil), t.stderr...)
303 t.stderrMu.Unlock()
304 return t.red.scrub(strings.Join(lines, "\n"))
305}
306
307func (t *stdioTransport) readLoop(r io.Reader) {
308 defer close(t.readDone)
309 reader := bufio.NewReaderSize(r, 64*1024)
310 for {
311 line, err := readLine(reader, maxStdioFrame)
312 if len(trimSpace(line)) > 0 {
313 frames, perr := parseFrames(line)
314 if perr == nil {
315 for _, msg := range frames {
316 t.dispatch(msg)
317 }
318 }
319 // Non-JSON stdout lines (stray logging) are ignored.
320 }
321 if err != nil {
322 if errors.Is(err, io.EOF) || errors.Is(err, os.ErrClosed) {
323 // Output closed. If the process is gone the wait goroutine
324 // reports its exit status; otherwise give it a moment.
325 select {
326 case <-t.exited:
327 return
328 case <-t.doneCh:
329 return
330 case <-time.After(time.Second):
331 }
332 }
333 t.fail(fmt.Errorf("read server output: %w", err))
334 return
335 }
336 }
337}
338
339func readLine(r *bufio.Reader, limit int) ([]byte, error) {
340 var buf []byte
341 for {
342 chunk, isPrefix, err := r.ReadLine()
343 buf = append(buf, chunk...)
344 if len(buf) > limit {
345 return nil, fmt.Errorf("message exceeds %d bytes", limit)
346 }
347 if err != nil {
348 return buf, err
349 }
350 if !isPrefix {
351 return buf, nil
352 }
353 }
354}
355
356func (t *stdioTransport) dispatch(msg message) {
357 if msg.isResponse() {
358 t.mu.Lock()
359 ch := t.pending[idKey(msg.ID)]
360 delete(t.pending, idKey(msg.ID))
361 t.mu.Unlock()
362 if ch != nil {
363 ch <- msg
364 }
365 return
366 }
367 t.mu.Lock()
368 handler := t.handler
369 t.mu.Unlock()
370 if msg.isRequest() {
371 go func() {
372 var result any
373 rpcErr := &RPCError{Code: codeMethodNotFound, Message: "method not found: " + msg.Method}
374 if handler != nil {
375 result, rpcErr = handler(msg)
376 }
377 if frame, err := encodeResponse(msg.ID, result, rpcErr); err == nil {
378 ctx, cancel := context.WithTimeout(t.ctx, responseWriteTimeout)
379 defer cancel()
380 _ = t.write(ctx, frame)
381 }
382 }()
383 return
384 }
385 if msg.Method != "" && handler != nil {
386 go handler(msg)
387 }
388}
389
390// writeLoop owns all writes to the child's stdin. Callers never block on the
391// pipe themselves: they queue a frame and wait for the result or their
392// context, so a child that stops reading cannot wedge a request, a timeout, or
393// close().
394func (t *stdioTransport) writeLoop() {
395 for {
396 select {
397 case <-t.doneCh:
398 return
399 case req := <-t.writeQ:
400 if t.hookDequeued != nil {
401 t.hookDequeued()
402 }
403 if err := req.ctx.Err(); err != nil && req.state.CompareAndSwap(reqQueued, reqSkipped) {
404 if t.hookSkipped != nil {
405 t.hookSkipped()
406 }
407 req.res <- err // never started: nothing partial on the wire
408 continue
409 }
410 if !req.state.CompareAndSwap(reqQueued, reqWriting) {
411 continue // the caller gave up before we started
412 }
413 t.writesStarted.Add(1)
414 _, err := t.stdin.Write(req.frame)
415 req.res <- err
416 if err != nil {
417 // A failed write may have left half a frame behind.
418 t.fail(fmt.Errorf("write to server: %w", err))
419 return
420 }
421 }
422 }
423}
424
425// write queues one frame (newline appended) and waits until it has been
426// written, ctx ends, or the connection fails.
427func (t *stdioTransport) write(ctx context.Context, frame []byte) error {
428 _, err := t.writeFrame(ctx, frame)
429 return err
430}
431
432// writeFrame is write that also reports whether any byte of the frame may
433// have reached the server. When ctx ends:
434// - a frame the writer has not started is withdrawn and never written, and
435// the transport stays usable;
436// - a frame whose write is already in progress cannot be withdrawn (part of
437// it may be on the wire and the rest could still be flushed later), so the
438// transport is failed and stdin closed. The manager reconnects on the next
439// call and the abandoned frame can never complete.
440func (t *stdioTransport) writeFrame(ctx context.Context, frame []byte) (sent bool, err error) {
441 req := &writeReq{ctx: ctx, frame: append(append(make([]byte, 0, len(frame)+1), frame...), '\n'), res: make(chan error, 1)}
442 select {
443 case t.writeQ <- req:
444 case <-ctx.Done():
445 return false, ctx.Err()
446 case <-t.doneCh:
447 return false, t.failure()
448 }
449 select {
450 case err := <-req.res:
451 return req.state.Load() == reqWriting, err
452 case <-ctx.Done():
453 if req.state.CompareAndSwap(reqQueued, reqSkipped) {
454 return false, ctx.Err()
455 }
456 if req.state.Load() == reqSkipped {
457 // The writer already withdrew the frame itself (its context had
458 // ended when it was dequeued): nothing was written, so the
459 // transport is untouched.
460 return false, ctx.Err()
461 }
462 select {
463 case err := <-req.res:
464 return req.state.Load() == reqWriting, err
465 default:
466 }
467 t.fail(fmt.Errorf("write to server interrupted: %w", ctx.Err()))
468 _ = t.stdin.Close() // unblocks the stuck write
469 return true, ctx.Err()
470 case <-t.doneCh:
471 select {
472 case err := <-req.res:
473 return req.state.Load() == reqWriting, err
474 default:
475 }
476 return req.state.Load() == reqWriting, t.failure()
477 }
478}
479
480func (t *stdioTransport) request(ctx context.Context, id int64, method string, params any) (json.RawMessage, error) {
481 frame, err := encodeRequest(id, method, params)
482 if err != nil {
483 return nil, err
484 }
485 key := fmt.Sprint(id)
486 ch := make(chan message, 1)
487 t.mu.Lock()
488 if t.err != nil {
489 err := t.err
490 t.mu.Unlock()
491 return nil, err
492 }
493 t.pending[key] = ch
494 t.mu.Unlock()
495 cleanup := func() {
496 t.mu.Lock()
497 delete(t.pending, key)
498 t.mu.Unlock()
499 }
500 if sent, err := t.writeFrame(ctx, frame); err != nil {
501 cleanup()
502 if ctx.Err() != nil {
503 if sent {
504 t.cancelRequest(id, method, ctx.Err().Error())
505 }
506 return nil, ctx.Err()
507 }
508 return nil, err
509 }
510 result := func(msg message) (json.RawMessage, error) {
511 if msg.Error != nil {
512 return nil, t.red.scrubRPCError(msg.Error)
513 }
514 return msg.Result, nil
515 }
516 select {
517 case msg := <-ch:
518 return result(msg)
519 case <-ctx.Done():
520 cleanup()
521 t.cancelRequest(id, method, ctx.Err().Error())
522 return nil, ctx.Err()
523 case <-t.doneCh:
524 // A response that arrived together with the failure still counts.
525 select {
526 case msg := <-ch:
527 return result(msg)
528 default:
529 }
530 cleanup()
531 return nil, t.failure()
532 }
533}
534
535// cancelRequest announces an abandoned request on a best-effort basis. It
536// never blocks the caller, is bounded by cancelNotifyTimeout, and dies with
537// the transport.
538func (t *stdioTransport) cancelRequest(id int64, method, reason string) {
539 if !shouldSendCancel(method) {
540 return
541 }
542 t.cancelsSent.Add(1)
543 go func() {
544 ctx, cancel := context.WithTimeout(t.ctx, cancelNotifyTimeout)
545 defer cancel()
546 _ = t.notify(ctx, "notifications/cancelled", map[string]any{"requestId": id, "reason": reason})
547 }()
548}
549
550func (t *stdioTransport) notify(ctx context.Context, method string, params any) error {
551 frame, err := encodeNotification(method, params)
552 if err != nil {
553 return err
554 }
555 return t.write(ctx, frame)
556}
557
558func (t *stdioTransport) setHandler(h incomingHandler) {
559 t.mu.Lock()
560 t.handler = h
561 t.mu.Unlock()
562}
563
564func (t *stdioTransport) setProtocolVersion(string) {}
565
566func (t *stdioTransport) done() <-chan struct{} { return t.doneCh }
567
568func (t *stdioTransport) failure() error {
569 t.mu.Lock()
570 defer t.mu.Unlock()
571 if t.err == nil {
572 return errors.New("connection closed")
573 }
574 return t.err
575}
576
577func (t *stdioTransport) fail(err error) {
578 t.mu.Lock()
579 if t.err == nil {
580 t.err = err
581 }
582 t.mu.Unlock()
583 t.closeOnce.Do(func() {
584 t.cancel()
585 close(t.doneCh)
586 })
587}
588
589func (t *stdioTransport) waitExit(d time.Duration) bool {
590 select {
591 case <-t.exited:
592 return true
593 case <-time.After(d):
594 return false
595 }
596}
597
598// close shuts the server down: close stdin (the MCP stdio shutdown signal)
599// without waiting behind any blocked write, wait briefly, then terminate the
600// process group. Group cleanup is one-shot and also runs when the direct
601// child exits on its own, so repeated or late calls never signal a stale id.
602func (t *stdioTransport) close() error {
603 t.fail(errors.New("connection closed"))
604 _ = t.stdin.Close() // also unblocks a writer stuck on a full pipe
605 if !t.waitExit(500 * time.Millisecond) {
606 t.guard.terminate(false)
607 if !t.waitExit(1500 * time.Millisecond) {
608 t.guard.terminate(true)
609 t.waitExit(time.Second)
610 }
611 }
612 t.guard.cleanup()
613 // With the child gone its stderr hits EOF once buffered data is read; give
614 // the reader a moment so diagnostics are complete. A descendant holding the
615 // pipe open only costs the timeout.
616 select {
617 case <-t.exited:
618 select {
619 case <-t.stderrDone:
620 case <-time.After(500 * time.Millisecond):
621 }
622 default:
623 }
624 _ = t.outR.Close()
625 _ = t.errR.Close()
626 return nil
627}