patterm/internal/mcp/tools.go

package mcp

import (
	"encoding/json"
	"errors"
	"fmt"
	"syscall"

	"github.com/hjbdev/patterm/internal/scratchpad"
)

// JSON-RPC error codes used by the patterm MCP surface. -32700..-32600
// are the standard parse/invalid-request codes; the SPEC-defined error
// names live in the -32000 range with a structured `data.kind` so the
// caller can branch on the error type rather than parsing strings.
const (
	ErrorKindInvalidArgs   = "invalid_args"
	ErrorKindInvalidKind   = "invalid_kind"
	ErrorKindNeedsTrust    = "needs_trust"
	ErrorKindRoleForbidden = "role_forbidden"
	ErrorKindNotRelated    = "not_related"
	ErrorKindNotFound      = "not_found"
	ErrorKindWrongKind     = "wrong_kind"
	ErrorKindUnknownAgent  = "unknown_agent"

	codeParseError     = -32700
	codeInvalidRequest = -32600
	codeMethodNotFound = -32601
	codeInvalidParams  = -32602
	codeInternal       = -32000
	codeNeedsTrust     = -32010
	codeRoleForbidden  = -32011
	codeNotRelated     = -32012
	codeNotFound       = -32013
	codeWrongKind      = -32014
	codeUnknownAgent   = -32015
)

// CallerRole is one of "orchestrator" or "sub-agent". SPEC §7 caller
// role: orchestrators sit at the root of a session tree; sub-agents
// were spawned by another agent.
type CallerRole string

const (
	RoleOrchestrator CallerRole = "orchestrator"
	RoleSubAgent     CallerRole = "sub-agent"
)

// ToolHost is the interface the in-process server uses to reach the
// running patterm process's state. internal/app implements this; the
// split keeps internal/mcp free of an internal/app import (which would
// be cyclic).
type ToolHost interface {
	// Identity resolution. The mcp-stdio greeting carries a per-spawn
	// token; the server resolves it to a process_id before dispatching
	// the rest of the connection's calls.
	ResolveCallerIdentity(identity string) string

	// CallerRole returns the role for the given process_id. Unknown
	// callers default to RoleOrchestrator (treated as a top-level peer)
	// so they don't get silently denied.
	CallerRole(processID string) CallerRole

	// Lifecycle (SPEC §7).
	SpawnAgent(callerID string, args SpawnAgentArgs) (ProcessInfo, error)
	SpawnProcess(callerID string, args SpawnProcessArgs) (ProcessInfo, error)
	StartProcess(callerID, processID string) (ProcessInfo, error)
	RestartProcess(callerID, processID string, signal syscall.Signal) (ProcessInfo, error)
	StopProcess(callerID, processID string, signal syscall.Signal) (ProcessInfo, error)
	CloseProcess(callerID, processID string) error
	RenameProcess(callerID, processID, name string) error
	SelectProcess(callerID, processID string) error

	// Inspection.
	ListProcesses(callerID, kindFilter string) []ProcessInfo
	GetProcessStatus(callerID, processID string) (ProcessStatus, error)
	GetProjectStatus(callerID string) (ProjectStatus, error)
	GetProcessOutput(callerID, processID, mode string, sinceOffset int64) (ProcessOutput, error)
	GetProcessRawOutput(callerID, processID string, sinceOffset int64) (RawOutput, error)
	SearchOutput(callerID, processID, pattern, kind string, limit int) (SearchResult, error)
	WaitForPattern(callerID, processID, pattern string, timeoutSeconds float64, scope string) (matched bool, snippet string, err error)
	GetProcessPorts(callerID, processID string) ([]PortSighting, error)

	// I/O.
	SendInput(callerID string, args SendInputArgs) (SendInputResult, error)

	// Coordination.
	SendMessage(callerID, targetID, message string) error
	RequestHumanAttention(callerID, processID, reason string) error
	TimerWait(callerID string, seconds float64, label string) (string, error)

	// Scratchpads.
	ScratchpadList() ([]scratchpad.Entry, error)
	ScratchpadRead(name string) (content string, revision string, err error)
	ScratchpadWrite(name, content, expectedRevision string) (revision string, err error)
	ScratchpadAppend(name, content string) error

	// Meta.
	WhoAmI(callerID string) WhoAmI
	Help(callerID, topic string) HelpResponse
}

// ProcessInfo is the entry shape returned by list_processes, spawn_*,
// stop_process, restart_process, start_process. SPEC §7.
type ProcessInfo struct {
	ID              string `json:"process_id"`
	Name            string `json:"name"`
	Kind            string `json:"kind"`
	Status          string `json:"status"`
	ParentProcessID string `json:"parent_process_id,omitempty"`
	ExitCode        *int   `json:"exit_code,omitempty"`
	IdleMS          int64  `json:"idle_ms,omitempty"`
	Trusted         *bool  `json:"trusted,omitempty"`
}

// ProcessStatus is what get_process_status returns. Richer than
// ProcessInfo: includes pane geometry, cursor, and active screen.
type ProcessStatus struct {
	ProcessInfo
	WorkingDir    string   `json:"working_dir,omitempty"`
	Argv          []string `json:"argv,omitempty"`
	StartedAt     string   `json:"started_at,omitempty"`
	ActiveScreen  string   `json:"active_screen,omitempty"`
	Rows          int      `json:"rows,omitempty"`
	Cols          int      `json:"cols,omitempty"`
	Cursor        Cursor   `json:"cursor"`
	ScreenVersion int64    `json:"screen_version,omitempty"`
}

// Cursor matches SPEC §7's `{x, y}` payload.
type Cursor struct {
	X int `json:"x"`
	Y int `json:"y"`
}

// ProjectStatus is what get_project_status returns — everything an
// agent needs to orient itself in one call.
type ProjectStatus struct {
	Project     ProjectMeta        `json:"project"`
	Caller      WhoAmI             `json:"caller"`
	Processes   []ProcessInfo      `json:"processes"`
	Scratchpads []scratchpad.Entry `json:"scratchpads"`
}

// ProjectMeta is the project root info echoed in many payloads.
type ProjectMeta struct {
	Path string `json:"path"`
	Key  string `json:"key"`
}

// ProcessOutput is the get_process_output payload. SPEC §7 enriches
// the old read_output result with screen geometry + version.
type ProcessOutput struct {
	Content       string `json:"content"`
	Mode          string `json:"mode"`
	NewOffset     int64  `json:"new_offset,omitempty"`
	ActiveScreen  string `json:"active_screen,omitempty"`
	Rows          int    `json:"rows,omitempty"`
	Cols          int    `json:"cols,omitempty"`
	Cursor        Cursor `json:"cursor"`
	IdleMS        int64  `json:"idle_ms,omitempty"`
	Status        string `json:"status,omitempty"`
	ScreenVersion int64  `json:"screen_version,omitempty"`
}

// RawOutput is the get_process_raw_output payload — ANSI preserved.
type RawOutput struct {
	Content   string `json:"content"`
	NewOffset int64  `json:"new_offset"`
	Status    string `json:"status,omitempty"`
}

// SearchResult is search_output's payload.
type SearchResult struct {
	Matches   []SearchMatch `json:"matches"`
	Truncated bool          `json:"truncated"`
}

type SearchMatch struct {
	LineNo int    `json:"line_no"`
	Text   string `json:"text"`
}

// PortSighting matches the per-child store in internal/app.
type PortSighting struct {
	Port        int    `json:"port"`
	URL         string `json:"url,omitempty"`
	FirstSeenAt string `json:"first_seen_at"`
}

// SpawnAgentArgs is the input shape for spawn_agent.
type SpawnAgentArgs struct {
	Agent             string `json:"agent"`
	AgentInstructions string `json:"agent_instructions"`
	Name              string `json:"name"`
}

// SpawnProcessArgs is the input shape for spawn_process.
type SpawnProcessArgs struct {
	Kind       string            `json:"kind"` // "terminal" | "command"
	Preset     string            `json:"preset"`
	Argv       []string          `json:"argv"`
	Name       string            `json:"name"`
	WorkingDir string            `json:"working_dir"`
	Env        map[string]string `json:"env"`
	Shell      bool              `json:"shell"`
}

// SendInputArgs is the input shape for send_input — covers text /
// paste / key with the optional wait+tail tail-after-send.
type SendInputArgs struct {
	ProcessID string `json:"process_id"`
	Kind      string `json:"kind"` // "text" | "paste" | "key"
	Text      string `json:"text"`
	Key       string `json:"key"`
	Submit    *bool  `json:"submit"`
	WaitMS    int    `json:"wait_ms"`
	TailMode  string `json:"tail_mode"` // "none" | "stream" | "grid"
}

// SendInputResult is the return shape of send_input.
type SendInputResult struct {
	OK   bool           `json:"ok"`
	Tail *ProcessOutput `json:"tail,omitempty"`
}

// WhoAmI is the whoami return shape.
type WhoAmI struct {
	ProcessID       string      `json:"process_id"`
	Name            string      `json:"name"`
	Role            CallerRole  `json:"role"`
	ParentProcessID string      `json:"parent_process_id,omitempty"`
	Project         ProjectMeta `json:"project"`
	AvailableTools  []string    `json:"available_tools"`
}

// HelpResponse is the help return shape.
type HelpResponse struct {
	Topic        string   `json:"topic"`
	Content      string   `json:"content"`
	RelatedTools []string `json:"related_tools,omitempty"`
}

func (s *Server) SetHost(h ToolHost) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.host = h
}

// dispatch routes a single JSON-RPC request. callerID is the ID of the
// process that owns this connection (resolved at greeting time).
// Returns nil for notifications (no id present), which tells the caller
// to skip writing a response. Otherwise returns a complete JSON-RPC
// reply ready to send.
func (s *Server) dispatch(callerID string, req []byte) []byte {
	var msg struct {
		JSONRPC string          `json:"jsonrpc"`
		ID      json.RawMessage `json:"id"`
		Method  string          `json:"method"`
		Params  json.RawMessage `json:"params"`
	}
	if err := json.Unmarshal(req, &msg); err != nil {
		return jsonRPCError(nil, codeParseError, "parse error: "+err.Error(), nil)
	}

	isNotification := len(msg.ID) == 0 || string(msg.ID) == "null"

	// MCP protocol-level methods (initialize, tools/list, tools/call,
	// ping, notifications) run before legacy direct-tool dispatch so
	// real MCP clients can hand-shake even when host isn't ready yet
	// (initialize doesn't touch the host).
	if result, handled, code, errMsg, data := s.handleProtocolMethod(callerID, msg.Method, msg.Params, isNotification); handled {
		if isNotification {
			return nil
		}
		if errMsg != "" {
			return jsonRPCError(msg.ID, code, errMsg, data)
		}
		return jsonRPCResult(msg.ID, result)
	}

	s.mu.Lock()
	host := s.host
	s.mu.Unlock()
	if host == nil {
		if isNotification {
			return nil
		}
		return jsonRPCError(msg.ID, codeInternal, "patterm: tool host not initialized", nil)
	}

	result, code, errMsg, data := callTool(host, callerID, msg.Method, msg.Params)
	if isNotification {
		return nil
	}
	if errMsg != "" {
		return jsonRPCError(msg.ID, code, errMsg, data)
	}
	return jsonRPCResult(msg.ID, result)
}

// toolError lets host implementations attach a structured `kind` to a
// JSON-RPC error so callers can branch on error type without parsing
// the message. The MCP layer maps recognized error kinds to their SPEC
// error codes; unknown errors fall through to codeInternal.
type toolError struct {
	Kind    string `json:"kind"`
	Message string `json:"message"`
}

func (e *toolError) Error() string { return e.Message }

// Errorf is a convenience for host implementations to return typed
// errors that map to JSON-RPC codes (needs_trust, role_forbidden, …).
func Errorf(kind, format string, a ...any) error {
	return &toolError{Kind: kind, Message: fmt.Sprintf(format, a...)}
}

// callTool is the central dispatch. Returns (result, code, errMsg,
// errData). On success errMsg is empty.
func callTool(h ToolHost, callerID, method string, params json.RawMessage) (any, int, string, any) {
	switch method {
	case "spawn_agent":
		var p SpawnAgentArgs
		if err := unmarshalParams(params, &p); err != nil {
			return nil, codeInvalidParams, err.Error(), nil
		}
		if p.Agent == "" {
			return nil, codeInvalidParams, "spawn_agent: agent required", nil
		}
		// Role gate: only orchestrators may call spawn_agent (SPEC §8).
		if h.CallerRole(callerID) == RoleSubAgent {
			return nil, codeRoleForbidden, "spawn_agent: sub-agents cannot spawn agents; use vendor-native subagent tooling or ask your parent", structuredKind("role_forbidden")
		}
		info, err := h.SpawnAgent(callerID, p)
		return mapToolResult(info, err)

	case "spawn_process":
		var p SpawnProcessArgs
		if err := unmarshalParams(params, &p); err != nil {
			return nil, codeInvalidParams, err.Error(), nil
		}
		info, err := h.SpawnProcess(callerID, p)
		return mapToolResult(info, err)

	case "start_process":
		var p struct {
			ProcessID string `json:"process_id"`
		}
		if err := unmarshalParams(params, &p); err != nil {
			return nil, codeInvalidParams, err.Error(), nil
		}
		info, err := h.StartProcess(callerID, p.ProcessID)
		return mapToolResult(info, err)

	case "restart_process":
		var p struct {
			ProcessID string `json:"process_id"`
			Signal    int    `json:"signal"`
		}
		if err := unmarshalParams(params, &p); err != nil {
			return nil, codeInvalidParams, err.Error(), nil
		}
		sig := syscall.Signal(p.Signal)
		info, err := h.RestartProcess(callerID, p.ProcessID, sig)
		return mapToolResult(info, err)

	case "stop_process":
		var p struct {
			ProcessID string `json:"process_id"`
			Signal    int    `json:"signal"`
		}
		if err := unmarshalParams(params, &p); err != nil {
			return nil, codeInvalidParams, err.Error(), nil
		}
		sig := syscall.Signal(p.Signal)
		info, err := h.StopProcess(callerID, p.ProcessID, sig)
		return mapToolResult(info, err)

	case "close_process":
		var p struct {
			ProcessID string `json:"process_id"`
		}
		if err := unmarshalParams(params, &p); err != nil {
			return nil, codeInvalidParams, err.Error(), nil
		}
		if err := h.CloseProcess(callerID, p.ProcessID); err != nil {
			return mapToolError(err)
		}
		return "ok", 0, "", nil

	case "rename_process":
		var p struct {
			ProcessID string `json:"process_id"`
			Name      string `json:"name"`
		}
		if err := unmarshalParams(params, &p); err != nil {
			return nil, codeInvalidParams, err.Error(), nil
		}
		if err := h.RenameProcess(callerID, p.ProcessID, p.Name); err != nil {
			return mapToolError(err)
		}
		return "ok", 0, "", nil

	case "select_process":
		var p struct {
			ProcessID string `json:"process_id"`
		}
		if err := unmarshalParams(params, &p); err != nil {
			return nil, codeInvalidParams, err.Error(), nil
		}
		if err := h.SelectProcess(callerID, p.ProcessID); err != nil {
			return mapToolError(err)
		}
		return "ok", 0, "", nil

	case "list_processes":
		var p struct {
			Kind string `json:"kind"`
		}
		_ = unmarshalParamsOptional(params, &p)
		return h.ListProcesses(callerID, p.Kind), 0, "", nil

	case "get_process_status":
		var p struct {
			ProcessID string `json:"process_id"`
		}
		if err := unmarshalParams(params, &p); err != nil {
			return nil, codeInvalidParams, err.Error(), nil
		}
		st, err := h.GetProcessStatus(callerID, p.ProcessID)
		if err != nil {
			return mapToolError(err)
		}
		return st, 0, "", nil

	case "get_project_status":
		ps, err := h.GetProjectStatus(callerID)
		if err != nil {
			return mapToolError(err)
		}
		return ps, 0, "", nil

	case "get_process_output":
		var p struct {
			ProcessID   string `json:"process_id"`
			Mode        string `json:"mode"`
			SinceOffset int64  `json:"since_offset"`
		}
		if err := unmarshalParams(params, &p); err != nil {
			return nil, codeInvalidParams, err.Error(), nil
		}
		if p.Mode == "" {
			p.Mode = "grid"
		}
		out, err := h.GetProcessOutput(callerID, p.ProcessID, p.Mode, p.SinceOffset)
		if err != nil {
			return mapToolError(err)
		}
		return out, 0, "", nil

	case "get_process_raw_output":
		var p struct {
			ProcessID   string `json:"process_id"`
			SinceOffset int64  `json:"since_offset"`
		}
		if err := unmarshalParams(params, &p); err != nil {
			return nil, codeInvalidParams, err.Error(), nil
		}
		out, err := h.GetProcessRawOutput(callerID, p.ProcessID, p.SinceOffset)
		if err != nil {
			return mapToolError(err)
		}
		return out, 0, "", nil

	case "search_output":
		var p struct {
			ProcessID string `json:"process_id"`
			Pattern   string `json:"pattern"`
			Kind      string `json:"kind"`
			Limit     int    `json:"limit"`
		}
		if err := unmarshalParams(params, &p); err != nil {
			return nil, codeInvalidParams, err.Error(), nil
		}
		if p.Limit <= 0 {
			p.Limit = 20
		}
		if p.Kind == "" {
			p.Kind = "rendered"
		}
		res, err := h.SearchOutput(callerID, p.ProcessID, p.Pattern, p.Kind, p.Limit)
		if err != nil {
			return mapToolError(err)
		}
		return res, 0, "", nil

	case "wait_for_pattern":
		var p struct {
			ProcessID      string  `json:"process_id"`
			Pattern        string  `json:"pattern"`
			TimeoutSeconds float64 `json:"timeout_seconds"`
			Scope          string  `json:"scope"`
		}
		if err := unmarshalParams(params, &p); err != nil {
			return nil, codeInvalidParams, err.Error(), nil
		}
		matched, snippet, err := h.WaitForPattern(callerID, p.ProcessID, p.Pattern, p.TimeoutSeconds, p.Scope)
		if err != nil {
			return mapToolError(err)
		}
		return map[string]any{"matched": matched, "snippet": snippet}, 0, "", nil

	case "get_process_ports":
		var p struct {
			ProcessID string `json:"process_id"`
		}
		if err := unmarshalParams(params, &p); err != nil {
			return nil, codeInvalidParams, err.Error(), nil
		}
		ports, err := h.GetProcessPorts(callerID, p.ProcessID)
		if err != nil {
			return mapToolError(err)
		}
		return map[string]any{"ports": ports}, 0, "", nil

	case "send_input":
		var p SendInputArgs
		if err := unmarshalParams(params, &p); err != nil {
			return nil, codeInvalidParams, err.Error(), nil
		}
		res, err := h.SendInput(callerID, p)
		if err != nil {
			return mapToolError(err)
		}
		return res, 0, "", nil

	case "send_message":
		var p struct {
			TargetProcessID string `json:"target_process_id"`
			Message         string `json:"message"`
		}
		if err := unmarshalParams(params, &p); err != nil {
			return nil, codeInvalidParams, err.Error(), nil
		}
		if err := h.SendMessage(callerID, p.TargetProcessID, p.Message); err != nil {
			return mapToolError(err)
		}
		return "ok", 0, "", nil

	case "request_human_attention":
		var p struct {
			ProcessID string `json:"process_id"`
			Reason    string `json:"reason"`
		}
		if err := unmarshalParams(params, &p); err != nil {
			return nil, codeInvalidParams, err.Error(), nil
		}
		if err := h.RequestHumanAttention(callerID, p.ProcessID, p.Reason); err != nil {
			return mapToolError(err)
		}
		return "ok", 0, "", nil

	case "timer_wait":
		var p struct {
			Seconds float64 `json:"seconds"`
			Label   string  `json:"label"`
		}
		if err := unmarshalParams(params, &p); err != nil {
			return nil, codeInvalidParams, err.Error(), nil
		}
		id, err := h.TimerWait(callerID, p.Seconds, p.Label)
		if err != nil {
			return mapToolError(err)
		}
		return map[string]string{"timer_id": id}, 0, "", nil

	case "scratchpad_list":
		entries, err := h.ScratchpadList()
		if err != nil {
			return nil, codeInternal, err.Error(), nil
		}
		return entries, 0, "", nil

	case "scratchpad_read":
		var p struct {
			Name string `json:"name"`
		}
		if err := unmarshalParams(params, &p); err != nil {
			return nil, codeInvalidParams, err.Error(), nil
		}
		content, rev, err := h.ScratchpadRead(p.Name)
		if err != nil {
			return nil, codeInternal, err.Error(), nil
		}
		return map[string]any{"content": content, "revision": rev}, 0, "", nil

	case "scratchpad_write":
		var p struct {
			Name             string `json:"name"`
			Content          string `json:"content"`
			ExpectedRevision string `json:"expected_revision"`
		}
		if err := unmarshalParams(params, &p); err != nil {
			return nil, codeInvalidParams, err.Error(), nil
		}
		rev, err := h.ScratchpadWrite(p.Name, p.Content, p.ExpectedRevision)
		if err != nil {
			// Optimistic-concurrency miss returns ok:false + current_revision
			// rather than a JSON-RPC error so callers can re-read + merge.
			var rme *scratchpad.RevisionMismatchError
			if errors.As(err, &rme) {
				return map[string]any{"ok": false, "current_revision": rme.CurrentRevision}, 0, "", nil
			}
			return nil, codeInternal, err.Error(), nil
		}
		return map[string]any{"ok": true, "revision": rev}, 0, "", nil

	case "scratchpad_append":
		var p struct {
			Name    string `json:"name"`
			Content string `json:"content"`
		}
		if err := unmarshalParams(params, &p); err != nil {
			return nil, codeInvalidParams, err.Error(), nil
		}
		if err := h.ScratchpadAppend(p.Name, p.Content); err != nil {
			return nil, codeInternal, err.Error(), nil
		}
		return map[string]any{"ok": true}, 0, "", nil

	case "whoami":
		return h.WhoAmI(callerID), 0, "", nil

	case "help":
		var p struct {
			Topic string `json:"topic"`
		}
		_ = unmarshalParamsOptional(params, &p)
		return h.Help(callerID, p.Topic), 0, "", nil
	}
	return nil, codeMethodNotFound, "method not found: " + method, nil
}

// mapToolResult is the (info, err) → JSON-RPC reply helper for the
// many handlers that return a ProcessInfo-ish struct.
func mapToolResult(v any, err error) (any, int, string, any) {
	if err != nil {
		return mapToolError(err)
	}
	return v, 0, "", nil
}

// mapToolError translates a host error to (code, message, data). If the
// error is a *toolError its Kind names the SPEC-defined category and we
// pick a stable JSON-RPC code for that category; otherwise we return a
// generic internal error.
func mapToolError(err error) (any, int, string, any) {
	var te *toolError
	if errors.As(err, &te) {
		code := codeInternal
		switch te.Kind {
		case ErrorKindInvalidArgs, ErrorKindInvalidKind:
			code = codeInvalidParams
		case ErrorKindNeedsTrust:
			code = codeNeedsTrust
		case ErrorKindRoleForbidden:
			code = codeRoleForbidden
		case ErrorKindNotRelated:
			code = codeNotRelated
		case ErrorKindNotFound:
			code = codeNotFound
		case ErrorKindWrongKind:
			code = codeWrongKind
		case ErrorKindUnknownAgent:
			code = codeUnknownAgent
		}
		return nil, code, te.Message, structuredKind(te.Kind)
	}
	return nil, codeInternal, err.Error(), nil
}

func structuredKind(kind string) map[string]string {
	return map[string]string{"kind": kind}
}

func unmarshalParams(raw json.RawMessage, out any) error {
	if len(raw) == 0 {
		return errors.New("missing params")
	}
	return json.Unmarshal(raw, out)
}

func unmarshalParamsOptional(raw json.RawMessage, out any) error {
	if len(raw) == 0 {
		return nil
	}
	return json.Unmarshal(raw, out)
}

func jsonRPCResult(id json.RawMessage, result any) []byte {
	resp := map[string]any{
		"jsonrpc": "2.0",
		"id":      id,
		"result":  result,
	}
	b, _ := json.Marshal(resp)
	return b
}

func jsonRPCError(id json.RawMessage, code int, message string, data any) []byte {
	errBody := map[string]any{
		"code":    code,
		"message": message,
	}
	if data != nil {
		errBody["data"] = data
	}
	resp := map[string]any{
		"jsonrpc": "2.0",
		"id":      id,
		"error":   errBody,
	}
	b, _ := json.Marshal(resp)
	return b
}