Cyrene_v2/internal/capability/mcp/stdio.go

package mcp

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"sync"
)

// stdioTransport runs an MCP server as a child process and speaks
// newline-delimited JSON-RPC 2.0 to its stdin/stdout.
type stdioTransport struct {
	cfg            MCPConfig
	log            *slog.Logger
	line           *lineRPC
	cmd            *exec.Cmd
	stdin          io.WriteCloser
	closeOnce      sync.Once
	closed         bool
	onNotification func(method string)
}

// setOnNotification records the handler invoked when the server sends a
// notification frame (a JSON-RPC message with no id). It is wired by the client
// before initialize so server notifications are not lost.
func (s *stdioTransport) setOnNotification(fn func(method string)) {
	s.onNotification = fn
	if s.line != nil {
		s.line.onNotification = fn
	}
}

// newStdioTransport builds a stdio transport for cfg.
func newStdioTransport(cfg MCPConfig, log *slog.Logger) *stdioTransport {
	return &stdioTransport{cfg: cfg, log: log}
}

// Start spawns the subprocess and wires up the request/response streams.
func (s *stdioTransport) Start(ctx context.Context) error {
	if err := ctx.Err(); err != nil {
		return fmt.Errorf("mcp stdio start: %w", err)
	}
	if s.cfg.Command == "" {
		return fmt.Errorf("mcp stdio: %w", ErrNoTransport)
	}
	cmd := exec.Command(s.cfg.Command, s.cfg.Args...)
	cmd.Env = envWith(os.Environ(), s.cfg.Env)
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return fmt.Errorf("mcp stdio stdin: %w", err)
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return fmt.Errorf("mcp stdio stdout: %w", err)
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return fmt.Errorf("mcp stdio stderr: %w", err)
	}
	if err := cmd.Start(); err != nil {
		return fmt.Errorf("mcp stdio start %q: %w", s.cfg.Command, err)
	}
	s.cmd = cmd
	s.stdin = stdin
	s.line = &lineRPC{in: stdin, out: bufio.NewReader(stdout), log: s.log, onNotification: s.onNotification}
	go s.drainStderr(stderr)
	return nil
}

// Call sends a request over stdio and decodes the response.
func (s *stdioTransport) Call(ctx context.Context, method string, params any, result any) error {
	if s.line == nil {
		return fmt.Errorf("mcp stdio %s: %w", method, ErrNotConnected)
	}
	return s.line.call(ctx, method, params, result)
}

// Close terminates the child process and closes its streams.
func (s *stdioTransport) Close() error {
	s.closeOnce.Do(func() {
		s.closed = true
		if s.stdin != nil {
			_ = s.stdin.Close()
		}
		if s.cmd != nil && s.cmd.Process != nil {
			_ = s.cmd.Process.Kill()
		}
		if s.cmd != nil {
			_ = s.cmd.Wait()
		}
	})
	return nil
}

// drainStderr reads and logs the child's standard error at debug level so server
// diagnostics are visible without blocking the process.
func (s *stdioTransport) drainStderr(r io.Reader) {
	buf := make([]byte, 4096)
	for {
		n, err := r.Read(buf)
		if n > 0 {
			s.log.Debug("mcp server stderr", "name", s.cfg.Name, "msg", strings.TrimSpace(string(buf[:n])))
		}
		if err != nil {
			return
		}
	}
}

// lineRPC carries newline-delimited JSON-RPC 2.0 over a stream pair. Calls are
// serialized so at most one request is in flight at a time, which keeps id
// matching simple and deterministic.
type lineRPC struct {
	in             io.Writer
	out            *bufio.Reader
	log            *slog.Logger
	mu             sync.Mutex
	nextID         int
	onNotification func(method string)
}

// call writes a request and reads until the matching response arrives.
func (l *lineRPC) call(ctx context.Context, method string, params any, result any) error {
	l.mu.Lock()
	defer l.mu.Unlock()
	if err := ctx.Err(); err != nil {
		return fmt.Errorf("mcp %s: %w", method, err)
	}
	l.nextID++
	id := strconv.Itoa(l.nextID)
	req, err := encodeRequest(id, method, params)
	if err != nil {
		return fmt.Errorf("mcp %s: %w", method, err)
	}
	req = append(req, '\n')
	if _, err := l.in.Write(req); err != nil {
		return fmt.Errorf("mcp %s write: %w", method, err)
	}
	return l.readResponse(id, method, result)
}

// readResponse reads lines until a response matching id arrives, skipping blank
// lines, malformed frames and notifications (no id).
func (l *lineRPC) readResponse(id, method string, result any) error {
	for {
		line, err := l.out.ReadBytes('\n')
		if err != nil {
			return fmt.Errorf("mcp %s read: %w", method, err)
		}
		line = bytes.TrimSpace(line)
		if len(line) == 0 {
			continue
		}
		var resp rpcResponse
		if err := json.Unmarshal(line, &resp); err != nil {
			l.log.Debug("mcp ignoring malformed frame", "line", string(line))
			continue
		}
		if resp.ID == nil {
			// A server-initiated notification (no id); hand it to the client's
			// handler if one is installed. The handler must not re-enter this
			// transport synchronously (see the serverNotifier contract).
			if resp.Method != "" && l.onNotification != nil {
				l.onNotification(resp.Method)
			}
			continue
		}
		if !decodeIDString(resp, id) {
			continue
		}
		if resp.Error != nil {
			return fmt.Errorf("mcp %s: %w: %w", method, ErrRPC, resp.Error)
		}
		if result != nil {
			if err := json.Unmarshal(resp.Result, result); err != nil {
				return fmt.Errorf("mcp %s decode: %w", method, err)
			}
		}
		return nil
	}
}

// notify writes a fire-and-forget JSON-RPC notification (no id).
func (l *lineRPC) notify(ctx context.Context, method string, params any) error {
	l.mu.Lock()
	defer l.mu.Unlock()
	req, err := encodeRequest(nil, method, params)
	if err != nil {
		return fmt.Errorf("mcp notify %s: %w", method, err)
	}
	req = append(req, '\n')
	if _, err := l.in.Write(req); err != nil {
		return fmt.Errorf("mcp notify %s write: %w", method, err)
	}
	return nil
}

// envWith returns base with extra merged over it (an extra key overrides a
// same-named base entry). A nil/empty extra returns base unchanged.
func envWith(base []string, extra map[string]string) []string {
	if len(extra) == 0 {
		return base
	}
	out := make([]string, 0, len(base)+len(extra))
	for _, kv := range base {
		if k, _, ok := strings.Cut(kv, "="); ok {
			if _, dup := extra[k]; dup {
				continue
			}
		}
		out = append(out, kv)
	}
	for k, v := range extra {
		out = append(out, k+"="+v)
	}
	return out
}