patterm/internal/app/mcp_inject.go

package app

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"

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

// patternMcpEntryName is the canonical name patterm uses when slotting
// itself into an external MCP config block (codex's mcp_servers,
// opencode's mcp, etc.). Stable on purpose: a single name means
// repeated spawns replace the previous entry instead of accumulating.
const patternMcpEntryName = "patterm"

// mcpConfigMerge prepares a vendored copy of the user's config file
// with patterm's MCP entry merged in, lays it out under a per-spawn
// home directory, and returns the env var assignment the child needs
// (e.g. "CODEX_HOME=/tmp/patterm-mcp-xxx").
//
// patterm never modifies the user's real config file in place. The
// merged copy lives under $XDG_RUNTIME_DIR/patterm/agents/<identity>/
// and is removed when the agent process exits.
func mcpConfigMerge(p *preset.Preset, inj *preset.MCPInjection, identity, bin, socket string) (envAssign, homeDir string, err error) {
	// Allow older preset files that pre-date the home_var / home_path /
	// format fields by falling back to known defaults for the well-known
	// agent config paths.
	homeVar, homePath, format := inj.HomeVar, inj.HomePath, strings.ToLower(inj.Format)
	if homeVar == "" || homePath == "" || format == "" {
		hv, hp, f := inferHomeFromPath(inj.Path)
		if homeVar == "" {
			homeVar = hv
		}
		if homePath == "" {
			homePath = hp
		}
		if format == "" {
			format = f
		}
	}
	if format == "" {
		switch strings.ToLower(filepath.Ext(inj.Path)) {
		case ".toml":
			format = "toml"
		case ".json":
			format = "json"
		}
	}
	if homeVar == "" || homePath == "" {
		return "", "", fmt.Errorf("preset %s: mcp_injection.config_file requires home_var and home_path (path %q not recognised; add the fields to the preset)", p.Name, inj.Path)
	}
	if inj.MergeKey == "" {
		return "", "", fmt.Errorf("preset %s: mcp_injection.config_file requires merge_key", p.Name)
	}
	if format == "" {
		return "", "", fmt.Errorf("preset %s: cannot infer mcp_injection.format from path %q", p.Name, inj.Path)
	}

	homeDir, err = mcpRuntimeDir(identity)
	if err != nil {
		return "", "", err
	}
	dest := filepath.Join(homeDir, homePath)
	if err := os.MkdirAll(filepath.Dir(dest), 0o700); err != nil {
		return "", "", err
	}

	src := expandUser(inj.Path)
	// Mirror the user's real agent-home directory (auth, sessions,
	// history, etc.) into the temp home via symlinks so codex / opencode
	// still see their credentials and prior state. Only the config file
	// itself is replaced with our merged copy.
	if err := mirrorAgentHome(filepath.Dir(src), filepath.Dir(dest), filepath.Base(dest)); err != nil {
		return "", "", err
	}

	srcBody, err := os.ReadFile(src)
	if err != nil && !errors.Is(err, os.ErrNotExist) {
		return "", "", fmt.Errorf("read %s: %w", src, err)
	}
	// srcBody stays nil if the user has no existing config — we'll
	// write a fresh minimal one with just the patterm entry.

	args := []string{"mcp-stdio", "--socket", socket, "--identity", identity}
	var merged []byte
	switch format {
	case "toml":
		merged, err = mergeTOMLMCP(srcBody, inj.MergeKey, bin, args)
	case "json":
		merged, err = mergeJSONMCP(srcBody, inj.MergeKey, bin, args)
	default:
		err = fmt.Errorf("preset %s: unsupported mcp_injection.format %q", p.Name, format)
	}
	if err != nil {
		return "", "", err
	}
	if err := os.WriteFile(dest, merged, 0o600); err != nil {
		return "", "", err
	}

	return homeVar + "=" + homeDir, homeDir, nil
}

// mcpConfigEnv reads the user's existing config file, merges patterm's
// MCP entry into it, and returns an env-var assignment (e.g.
// `OPENCODE_CONFIG_CONTENT={...}`) the child can read directly. No
// file is written and XDG_CONFIG_HOME is not touched — the agent's
// auth/state/skill dirs continue to resolve from the user's real
// $HOME exactly as they do without patterm.
func mcpConfigEnv(p *preset.Preset, inj *preset.MCPInjection, identity, bin, socket string) (string, error) {
	if inj.Var == "" {
		return "", fmt.Errorf("preset %s: mcp_injection.config_env requires var", p.Name)
	}
	if inj.MergeKey == "" {
		return "", fmt.Errorf("preset %s: mcp_injection.config_env requires merge_key", p.Name)
	}
	format := strings.ToLower(inj.Format)
	if format == "" {
		switch strings.ToLower(filepath.Ext(inj.Path)) {
		case ".toml":
			format = "toml"
		case ".json":
			format = "json"
		}
	}
	if format == "" {
		return "", fmt.Errorf("preset %s: cannot infer mcp_injection.format from path %q", p.Name, inj.Path)
	}

	var srcBody []byte
	if inj.Path != "" {
		body, err := os.ReadFile(expandUser(inj.Path))
		if err != nil && !errors.Is(err, os.ErrNotExist) {
			return "", fmt.Errorf("read %s: %w", inj.Path, err)
		}
		srcBody = body
	}

	args := []string{"mcp-stdio", "--socket", socket, "--identity", identity}
	var merged []byte
	var err error
	switch format {
	case "toml":
		merged, err = mergeTOMLMCP(srcBody, inj.MergeKey, bin, args)
	case "json":
		merged, err = mergeJSONMCP(srcBody, inj.MergeKey, bin, args)
	default:
		err = fmt.Errorf("preset %s: unsupported mcp_injection.format %q", p.Name, format)
	}
	if err != nil {
		return "", err
	}
	return inj.Var + "=" + string(merged), nil
}

// mcpCLIOverrideArgs builds the `-c key=value` argv tail for the
// `cli_override` injection kind. The agent merges these into its
// in-memory config at startup, so there's no filesystem footprint at
// all — codex picks up patterm's MCP server without us touching
// ~/.codex/config.toml or hijacking CODEX_HOME (which would mask
// auth.json and saved sessions).
func mcpCLIOverrideArgs(p *preset.Preset, inj *preset.MCPInjection, identity, bin, socket string) ([]string, error) {
	flag := inj.Flag
	if flag == "" {
		flag = "-c"
	}
	prefix := inj.KeyPrefix
	if prefix == "" {
		return nil, fmt.Errorf("preset %s: mcp_injection.cli_override requires key_prefix", p.Name)
	}
	args := []string{"mcp-stdio", "--socket", socket, "--identity", identity}

	// We hard-code TOML scalar encoding because every consumer in the
	// wild (codex today; future cli_override targets are expected to
	// be the same) parses overrides as TOML expressions. Quoting the
	// command preserves spaces in paths; quoting each args element
	// keeps the array shape intact.
	cmdVal := tomlString(bin)
	var argsVal strings.Builder
	argsVal.WriteString("[")
	for i, a := range args {
		if i > 0 {
			argsVal.WriteString(", ")
		}
		argsVal.WriteString(tomlString(a))
	}
	argsVal.WriteString("]")

	return []string{
		flag, prefix + ".command=" + cmdVal,
		flag, prefix + ".args=" + argsVal.String(),
	}, nil
}

// tomlString renders a Go string as a TOML basic string literal. TOML
// uses the same escape conventions as JSON for backslash and quote,
// which keeps this implementation small.
func tomlString(s string) string {
	b, _ := json.Marshal(s)
	return string(b)
}

// inferHomeFromPath maps the well-known agent config paths to the env
// var + relative path patterm should use when merging. Lets older
// preset files (without home_var/home_path/format) keep working.
func inferHomeFromPath(path string) (homeVar, homePath, format string) {
	switch {
	case strings.HasSuffix(path, "/.codex/config.toml"):
		return "CODEX_HOME", "config.toml", "toml"
	case strings.HasSuffix(path, "/opencode/opencode.json"):
		return "XDG_CONFIG_HOME", "opencode/opencode.json", "json"
	}
	return "", "", ""
}

// mirrorAgentHome populates mirroredDir with symlinks pointing at each
// entry of srcDir, except for skipBase (which the caller is replacing
// with a freshly-written file). This lets agents that root every piece
// of their per-user state at one dir — codex via CODEX_HOME, opencode
// via XDG_CONFIG_HOME/opencode — keep reading their real auth.json,
// sessions, history, etc. even when patterm overrides the home root.
func mirrorAgentHome(srcDir, mirroredDir, skipBase string) error {
	if err := os.MkdirAll(mirroredDir, 0o700); err != nil {
		return err
	}
	entries, err := os.ReadDir(srcDir)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			return nil
		}
		return err
	}
	for _, e := range entries {
		if e.Name() == skipBase {
			continue
		}
		src := filepath.Join(srcDir, e.Name())
		dst := filepath.Join(mirroredDir, e.Name())
		// Replace any stale symlink/file at dst — the runtime dir is
		// per-identity so this should be a no-op on first spawn, but
		// being defensive keeps re-spawn semantics sane if the dir is
		// reused.
		_ = os.Remove(dst)
		if err := os.Symlink(src, dst); err != nil {
			return fmt.Errorf("symlink %s -> %s: %w", src, dst, err)
		}
	}
	return nil
}

func mcpRuntimeDir(identity string) (string, error) {
	if runtime := os.Getenv("XDG_RUNTIME_DIR"); runtime != "" {
		dir := filepath.Join(runtime, "patterm", "agents", identity)
		if err := os.MkdirAll(dir, 0o700); err != nil {
			return "", err
		}
		return dir, nil
	}
	dir := filepath.Join(os.TempDir(), "patterm-agents-"+identity)
	if err := os.MkdirAll(dir, 0o700); err != nil {
		return "", err
	}
	return dir, nil
}

func expandUser(p string) string {
	if strings.HasPrefix(p, "~/") {
		home, err := os.UserHomeDir()
		if err == nil {
			return filepath.Join(home, p[2:])
		}
	}
	return p
}

// mergeJSONMCP parses src as JSON, slots patterm's MCP entry under the
// merge key, and reserializes. If src is empty/whitespace, we start
// from an empty object. opencode uses a `command` array shape with
// `type: "local"`; codex JSON variants (uncommon) reuse the codex
// command/args shape. We emit the opencode shape because it's the
// only JSON consumer in the default preset set.
func mergeJSONMCP(src []byte, mergeKey, bin string, args []string) ([]byte, error) {
	var root map[string]any
	trimmed := strings.TrimSpace(string(src))
	if trimmed == "" {
		root = map[string]any{}
	} else {
		if err := json.Unmarshal([]byte(trimmed), &root); err != nil {
			return nil, fmt.Errorf("parse json config: %w", err)
		}
	}

	mcp, _ := root[mergeKey].(map[string]any)
	if mcp == nil {
		mcp = map[string]any{}
	}

	entry := map[string]any{
		"type":    "local",
		"command": append([]string{bin}, args...),
		"enabled": true,
	}
	mcp[patternMcpEntryName] = entry
	root[mergeKey] = mcp

	out, err := json.MarshalIndent(root, "", "  ")
	if err != nil {
		return nil, err
	}
	return append(out, '\n'), nil
}

// mergeTOMLMCP merges a `[<mergeKey>.patterm]` block into a TOML
// document. We deliberately avoid pulling in a full TOML parser:
// codex's config.toml is human-edited but the patterm entry is
// well-bounded, so a string-level "strip the old patterm section,
// append a fresh one" suffices for the merge use case.
func mergeTOMLMCP(src []byte, mergeKey, bin string, args []string) ([]byte, error) {
	stripped := stripTOMLSection(string(src), mergeKey+"."+patternMcpEntryName)

	if stripped != "" && !strings.HasSuffix(stripped, "\n") {
		stripped += "\n"
	}
	if stripped != "" {
		stripped += "\n"
	}

	var b strings.Builder
	b.WriteString(stripped)
	b.WriteString("# managed by patterm — re-written on each spawn\n")
	fmt.Fprintf(&b, "[%s.%s]\n", mergeKey, patternMcpEntryName)
	fmt.Fprintf(&b, "command = %q\n", bin)
	b.WriteString("args = [")
	for i, a := range args {
		if i > 0 {
			b.WriteString(", ")
		}
		fmt.Fprintf(&b, "%q", a)
	}
	b.WriteString("]\n")
	return []byte(b.String()), nil
}

// stripTOMLSection returns src with the `[header]` table (and the
// lines until the next top-level `[...]` header or EOF) removed.
// Lines that begin with `header.` as a subsection of the target are
// also dropped so we don't leave stale per-key dotted assignments.
func stripTOMLSection(src, header string) string {
	if src == "" {
		return ""
	}
	wantTable := "[" + header + "]"
	wantSubPrefix := "[" + header + "."
	lines := strings.Split(src, "\n")
	out := make([]string, 0, len(lines))
	inTarget := false
	for _, line := range lines {
		trimmed := strings.TrimSpace(line)
		if strings.HasPrefix(trimmed, "[") && strings.HasSuffix(trimmed, "]") {
			if trimmed == wantTable || strings.HasPrefix(trimmed, wantSubPrefix) {
				inTarget = true
				continue
			}
			inTarget = false
		}
		if inTarget {
			continue
		}
		out = append(out, line)
	}
	joined := strings.Join(out, "\n")
	return strings.TrimRight(joined, "\n")
}