Maven/internal/mcp/stdio.go

package mcp

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

// maxLine bounds one JSON-RPC frame from a subprocess. A tool result bigger
// than this is a misbehaving server, not something to buffer.
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.
type stdioTransport struct {
	mu   sync.Mutex
	cmd  *exec.Cmd
	in   io.WriteCloser
	out  *bufio.Reader
	dead bool
}

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)
	}
	return &stdioTransport{cmd: cmd, in: in, out: bufio.NewReaderSize(out, 64<<10)}, nil
}

func (t *stdioTransport) Call(ctx context.Context, req *rpcRequest) (*rpcResponse, error) {
	t.mu.Lock()
	defer t.mu.Unlock()
	if t.dead {
		return nil, ErrClosed
	}
	if err := t.write(req); err != nil {
		t.dead = true
		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 {
		if err := ctx.Err(); err != nil {
			return nil, err
		}
		line, err := t.readLine()
		if err != nil {
			t.dead = true
			return nil, err
		}
		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})
}

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
}

func (t *stdioTransport) readLine() ([]byte, error) {
	for {
		line, err := t.out.ReadString('\n')
		if err != nil {
			if len(strings.TrimSpace(line)) == 0 {
				return nil, fmt.Errorf("mcp: read: %w", err)
			}
			return []byte(line), nil
		}
		if len(line) > maxLine {
			return nil, fmt.Errorf("mcp: frame exceeds %d bytes", maxLine)
		}
		if s := strings.TrimSpace(line); s != "" {
			return []byte(s), nil
		}
	}
}

func (t *stdioTransport) Close() error {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.dead = true
	if t.in != nil {
		_ = t.in.Close()
	}
	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()
	return !t.dead
}