Maven/internal/mcp/stdio.go

package mcp

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"sync"
)

// maxLine bounds one JSON-RPC frame from a subprocess. A tool result bigger
// than this is a misbehaving server, not something to buffer.
//
// The bound is enforced by bufio.Scanner's own buffer limit, not by measuring
// the line after it was assembled. Measuring afterwards is not a bound: a
// server that emits 500 MB with no newline would have all 500 MB in mavend's
// heap before the check could reject it, which on the deploy target is an OOM
// kill of the core daemon.
const maxLine = 1 << 20 // 1 MiB

// stdioTransport speaks newline-delimited JSON-RPC to a child process. This is
// the local transport: the server runs on this box, under this user, and gets
// no network guard because it never touches the network on our behalf.
//
// Args are argv, never a shell string — the same discipline internal/tool
// keeps, for the same reason.
//
// Reading happens on its own goroutine, feeding frames down a channel. That is
// what makes a call abandonable: bufio never observes a context, so a server
// that accepts a request and then writes nothing at all would otherwise block
// the reader forever with the transport lock held, and every other server in
// the manager behind it.
type stdioTransport struct {
	cmd   *exec.Cmd
	in    io.WriteCloser
	lines chan []byte
	stop  chan struct{} // closed by Close, so the reader can give up

	// callMu serialises whole calls, so two callers cannot consume each
	// other's frames off the shared channel. It is deliberately NOT the lock
	// alive() takes: a hung call must not make the manager's health check
	// block on it.
	callMu sync.Mutex

	mu      sync.Mutex
	dead    bool
	readErr error
}

func newStdioTransport(ctx context.Context, argv []string, env []string, dir string) (*stdioTransport, error) {
	if len(argv) == 0 {
		return nil, errors.New("mcp: stdio server needs a command")
	}
	cmd := exec.Command(argv[0], argv[1:]...)
	cmd.Dir = dir
	if len(env) > 0 {
		cmd.Env = append(os.Environ(), env...)
	}
	cmd.Stderr = os.Stderr
	in, err := cmd.StdinPipe()
	if err != nil {
		return nil, fmt.Errorf("mcp: stdin pipe: %w", err)
	}
	out, err := cmd.StdoutPipe()
	if err != nil {
		return nil, fmt.Errorf("mcp: stdout pipe: %w", err)
	}
	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("mcp: start %q: %w", argv[0], err)
	}
	t := &stdioTransport{
		cmd:   cmd,
		in:    in,
		lines: make(chan []byte),
		stop:  make(chan struct{}),
	}
	go t.readLoop(out)
	return t, nil
}

// readLoop pushes one frame per line onto t.lines until the pipe ends. The
// scanner's own buffer limit is the frame bound: a line longer than maxLine
// ends the scan with bufio.ErrTooLong having buffered at most maxLine, rather
// than assembling the whole thing first and rejecting it afterwards.
func (t *stdioTransport) readLoop(out io.Reader) {
	defer close(t.lines)
	sc := bufio.NewScanner(out)
	sc.Buffer(make([]byte, 0, 64<<10), maxLine)
	for sc.Scan() {
		line := bytes.TrimSpace(sc.Bytes())
		if len(line) == 0 {
			continue
		}
		frame := append([]byte(nil), line...)
		select {
		case t.lines <- frame:
		case <-t.stop:
			return
		}
	}
	err := sc.Err()
	switch {
	case errors.Is(err, bufio.ErrTooLong):
		err = fmt.Errorf("mcp: frame exceeds %d bytes", maxLine)
	case err == nil:
		err = io.EOF
	}
	t.mu.Lock()
	t.readErr = err
	t.mu.Unlock()
}

func (t *stdioTransport) Call(ctx context.Context, req *rpcRequest) (*rpcResponse, error) {
	t.callMu.Lock()
	defer t.callMu.Unlock()
	if !t.alive() {
		return nil, ErrClosed
	}
	t.mu.Lock()
	err := t.write(req)
	t.mu.Unlock()
	if err != nil {
		_ = t.Close()
		return nil, err
	}
	// Read until the frame with our id turns up; anything else on the pipe is
	// a notification or a server-initiated request we do not answer.
	for {
		select {
		case <-ctx.Done():
			// A server that took the request and answered nothing is not a
			// server this connection can be reused with: the next call would
			// read into a pipe whose state we no longer know. Drop it and let
			// the manager re-dial.
			_ = t.Close()
			return nil, ctx.Err()
		case line, ok := <-t.lines:
			if !ok {
				t.mu.Lock()
				rerr := t.readErr
				t.mu.Unlock()
				_ = t.Close()
				if rerr == nil {
					rerr = ErrClosed
				}
				return nil, fmt.Errorf("mcp: read: %w", rerr)
			}
			var resp rpcResponse
			if err := json.Unmarshal(line, &resp); err != nil {
				continue // not a response frame; ignore rather than break the turn
			}
			if resp.ID == nil || *resp.ID != req.ID {
				continue
			}
			return &resp, nil
		}
	}
}

func (t *stdioTransport) Notify(ctx context.Context, method string, params any) error {
	t.mu.Lock()
	defer t.mu.Unlock()
	if t.dead {
		return ErrClosed
	}
	return t.write(&rpcRequest{JSONRPC: "2.0", Method: method, Params: params})
}

// write must be called with t.mu held.
func (t *stdioTransport) write(req *rpcRequest) error {
	req.JSONRPC = "2.0"
	raw, err := json.Marshal(req)
	if err != nil {
		return err
	}
	if _, err := t.in.Write(append(raw, '\n')); err != nil {
		return fmt.Errorf("mcp: write %s: %w", req.Method, err)
	}
	return nil
}

// Close is idempotent: a call that abandoned a silent pipe calls it, and so
// does the manager.
func (t *stdioTransport) Close() error {
	t.mu.Lock()
	if t.dead {
		t.mu.Unlock()
		return nil
	}
	t.dead = true
	close(t.stop)
	if t.in != nil {
		_ = t.in.Close()
	}
	t.mu.Unlock()
	if t.cmd.Process != nil {
		_ = t.cmd.Process.Kill()
		_ = t.cmd.Wait()
	}
	return nil
}

// alive reports whether the transport can still carry a call. The manager uses
// it to decide on a reconnect instead of retrying into a dead pipe.
func (t *stdioTransport) alive() bool {
	t.mu.Lock()
	defer t.mu.Unlock()
	if t.dead {
		return false
	}
	return t.readErr == nil
}