patterm/internal/mcp/mcp.go

// Package mcp is patterm's in-process MCP server and the stdio proxy
// subcommand that spawned children connect through. SPEC §7 + §10.
//
// v1 stubs out the server: it listens on the per-PID socket, accepts
// connections from `patterm mcp-stdio` proxies, and returns a "not
// implemented" error for every tool call. The plumbing is in place so
// later milestones (suggested build order §15 step 4 onwards) can fill
// in real tool handlers without touching the lifecycle code.
package mcp

import (
	"bufio"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"sync"
)

// Server is patterm's in-process MCP server. SPEC §10: bound to a
// per-PID unix socket under $XDG_RUNTIME_DIR/patterm/<pid>.sock.
type Server struct {
	socket   string
	listener net.Listener

	mu     sync.Mutex
	closed bool
	host   ToolHost
}

// SocketPath returns the per-PID socket path with the standard fallback.
// SPEC §3.
func SocketPath() (string, error) {
	pid := strconv.Itoa(os.Getpid())
	if runtime := os.Getenv("XDG_RUNTIME_DIR"); runtime != "" {
		dir := filepath.Join(runtime, "patterm")
		if err := os.MkdirAll(dir, 0o700); err != nil {
			return "", fmt.Errorf("mcp: mkdir %s: %w", dir, err)
		}
		return filepath.Join(dir, pid+".sock"), nil
	}
	return filepath.Join("/tmp", "patterm-"+pid+".sock"), nil
}

// Start opens the per-PID socket and serves JSON-RPC over it. The
// returned Server can be Close()d on shutdown.
func Start() (*Server, error) {
	path, err := SocketPath()
	if err != nil {
		return nil, err
	}
	_ = os.Remove(path)
	ln, err := net.Listen("unix", path)
	if err != nil {
		return nil, fmt.Errorf("mcp: listen %s: %w", path, err)
	}
	_ = os.Chmod(path, 0o600)
	s := &Server{socket: path, listener: ln}
	go s.acceptLoop()
	return s, nil
}

func (s *Server) Socket() string { return s.socket }

func (s *Server) Close() error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return nil
	}
	s.closed = true
	_ = s.listener.Close()
	_ = os.Remove(s.socket)
	return nil
}

func (s *Server) acceptLoop() {
	for {
		conn, err := s.listener.Accept()
		if err != nil {
			if errors.Is(err, net.ErrClosed) {
				return
			}
			continue
		}
		go s.handleConn(conn)
	}
}

// handleConn reads newline-delimited JSON-RPC requests from a connected
// child and dispatches them. The first line carries the per-spawn
// identity token (SPEC §10); we resolve it to a child id and stash that
// as the caller for every subsequent tool call.
func (s *Server) handleConn(conn net.Conn) {
	var writeMu sync.Mutex
	var wg sync.WaitGroup
	defer func() {
		wg.Wait()
		_ = conn.Close()
	}()
	r := bufio.NewReader(conn)

	var callerID string
	writeResp := func(resp []byte) bool {
		if resp == nil {
			return true
		}
		resp = append(resp, '\n')
		writeMu.Lock()
		defer writeMu.Unlock()
		for len(resp) > 0 {
			n, err := conn.Write(resp)
			if err != nil {
				return false
			}
			if n == 0 {
				return false
			}
			resp = resp[n:]
		}
		return true
	}

	greeting, err := r.ReadBytes('\n')
	if err != nil {
		return
	}
	if tok := greetingIdentity(greeting); tok != "" {
		s.mu.Lock()
		host := s.host
		s.mu.Unlock()
		if host != nil {
			callerID = host.ResolveCallerIdentity(tok)
		}
	} else {
		// Treat as a real request from an unknown caller.
		resp := s.dispatch("", greeting)
		if !writeResp(resp) {
			return
		}
	}

	for {
		line, err := r.ReadBytes('\n')
		if len(line) > 0 {
			req := append([]byte(nil), line...)
			wg.Add(1)
			go func() {
				defer wg.Done()
				resp := s.dispatch(callerID, req)
				_ = writeResp(resp)
			}()
		}
		if err != nil {
			return
		}
	}
}

func greetingIdentity(b []byte) string {
	var probe struct {
		Identity string `json:"patterm_identity"`
	}
	if err := json.Unmarshal(b, &probe); err != nil {
		return ""
	}
	return probe.Identity
}

// RunStdioProxy is the entry point for `patterm mcp-stdio`. It opens
// the per-PID socket and shuttles bytes between os.Stdin/os.Stdout and
// the socket. SPEC §10: the vendor CLI thinks it's launching a normal
// stdio MCP server; this proxy forwards JSON-RPC to the running
// patterm process.
func RunStdioProxy(socket, identity string) error {
	conn, err := net.Dial("unix", socket)
	if err != nil {
		return fmt.Errorf("dial %s: %w", socket, err)
	}
	defer conn.Close()

	// Send a one-line greeting carrying the identity so the server
	// knows which child it's talking to. Format: {"patterm_identity":
	// "<token>"} + newline. Real protocol handshake is a later
	// milestone.
	greeting := map[string]string{"patterm_identity": identity}
	if key := os.Getenv("PATTERM_PROJECT_KEY"); key != "" {
		greeting["project_key"] = key
	}
	gb, _ := json.Marshal(greeting)
	gb = append(gb, '\n')
	if _, err := conn.Write(gb); err != nil {
		return fmt.Errorf("greeting: %w", err)
	}

	errCh := make(chan error, 2)
	go func() {
		_, err := io.Copy(conn, os.Stdin)
		errCh <- err
	}()
	go func() {
		_, err := io.Copy(os.Stdout, conn)
		errCh <- err
	}()
	<-errCh
	return nil
}