PaddedCell/ego-mgr/src/main.go

package main

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"syscall"

	"github.com/spf13/cobra"
)

const (
	expectedAgentVerify = "IF YOU ARE AN AGENT/MODEL, YOU SHOULD NEVER TOUCH THIS ENV VARIABLE"
	egoFileName         = "ego.json"
)

// Exit codes per spec
const (
	ExitSuccess          = 0
	ExitUsageError      = 1
	ExitColumnNotFound  = 2
	ExitColumnExists    = 3
	ExitPermission      = 4
	ExitLockFailed      = 5
	ExitJSONError       = 6
	ExitNotFound       = 7
)

// EgoData is the on-disk JSON structure
type EgoData struct {
	Columns       []string                       `json:"columns"`
	PublicColumns []string                        `json:"public-columns"`
	PublicScope   map[string]string               `json:"public-scope"`
	AgentScope    map[string]map[string]string    `json:"agent-scope"`
}

func resolveOpenclawPath() string {
	if p := os.Getenv("OPENCLAW_PATH"); p != "" {
		return p
	}
	home, _ := os.UserHomeDir()
	return filepath.Join(home, ".openclaw")
}

func egoFilePath() string {
	return filepath.Join(resolveOpenclawPath(), egoFileName)
}

func currentAgentID() string {
	return os.Getenv("AGENT_ID")
}

func requirePcguard() {
	if os.Getenv("AGENT_VERIFY") != expectedAgentVerify {
		fmt.Fprintln(os.Stderr, "Error: must be invoked via pcexec (AGENT_VERIFY mismatch)")
		os.Exit(ExitPermission)
	}
	if os.Getenv("AGENT_ID") == "" {
		fmt.Fprintln(os.Stderr, "Error: AGENT_ID not set — must be invoked via pcexec")
		os.Exit(ExitPermission)
	}
	if os.Getenv("AGENT_WORKSPACE") == "" {
		fmt.Fprintln(os.Stderr, "Error: AGENT_WORKSPACE not set — must be invoked via pcexec")
		os.Exit(ExitPermission)
	}
}

// readEgoData reads and parses the ego.json file
func readEgoData() (*EgoData, error) {
	fp := egoFilePath()
	raw, err := os.ReadFile(fp)
	if err != nil {
		if os.IsNotExist(err) {
			// Return empty structure
			return &EgoData{
				Columns:       []string{},
				PublicColumns: []string{},
				PublicScope:   map[string]string{},
				AgentScope:    map[string]map[string]string{},
			}, nil
		}
		return nil, fmt.Errorf("failed to read %s: %w", fp, err)
	}
	var data EgoData
	if err := json.Unmarshal(raw, &data); err != nil {
		return nil, fmt.Errorf("failed to parse %s: %w", fp, err)
	}
	// Ensure maps are initialized
	if data.PublicScope == nil {
		data.PublicScope = map[string]string{}
	}
	if data.AgentScope == nil {
		data.AgentScope = map[string]map[string]string{}
	}
	if data.Columns == nil {
		data.Columns = []string{}
	}
	if data.PublicColumns == nil {
		data.PublicColumns = []string{}
	}
	return &data, nil
}

// writeEgoData writes ego data to ego.json with file locking
func writeEgoData(data *EgoData) error {
	fp := egoFilePath()

	// Acquire file lock
	lockPath := fp + ".lock"
	lockFile, err := os.OpenFile(lockPath, os.O_CREATE|os.O_RDWR, 0600)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error: cannot create lock file: %v\n", err)
		os.Exit(ExitLockFailed)
	}
	defer func() {
		lockFile.Close()
		os.Remove(lockPath)
	}()

	if err := syscall.Flock(int(lockFile.Fd()), syscall.LOCK_EX|syscall.LOCK_NB); err != nil {
		fmt.Fprintln(os.Stderr, "Error: failed to acquire file lock (another process is writing)")
		os.Exit(ExitLockFailed)
	}
	defer syscall.Flock(int(lockFile.Fd()), syscall.LOCK_UN)

	raw, err := json.MarshalIndent(data, "", "  ")
	if err != nil {
		return fmt.Errorf("failed to marshal JSON: %w", err)
	}

	if err := os.WriteFile(fp, append(raw, '\n'), 0644); err != nil {
		return fmt.Errorf("failed to write %s: %w", fp, err)
	}
	return nil
}

// ensureAgent ensures the current agent has an entry in agent-scope
func ensureAgent(data *EgoData) {
	agentID := currentAgentID()
	if _, ok := data.AgentScope[agentID]; !ok {
		data.AgentScope[agentID] = map[string]string{}
	}
}

// isPublicColumn checks if a column name is in public-columns
func isPublicColumn(data *EgoData, name string) bool {
	for _, c := range data.PublicColumns {
		if c == name {
			return true
		}
	}
	return false
}

// isAgentColumn checks if a column name is in columns (agent scope)
func isAgentColumn(data *EgoData, name string) bool {
	for _, c := range data.Columns {
		if c == name {
			return true
		}
	}
	return false
}

// columnExists checks if a column name exists in either scope
func columnExists(data *EgoData, name string) bool {
	return isPublicColumn(data, name) || isAgentColumn(data, name)
}

func main() {
	rootCmd := &cobra.Command{
		Use:   "ego-mgr",
		Short: "Agent identity/profile manager for OpenClaw",
		Long: `ego-mgr manages agent personal information (name, email, timezone, etc.).

Fields can be Agent Scope (per-agent) or Public Scope (shared by all agents).

Examples:
  ego-mgr add column name
  ego-mgr add public-column timezone --default UTC
  ego-mgr set name "小智"
  ego-mgr get name
  ego-mgr show
  ego-mgr list columns
  ego-mgr delete name`,
	}

	rootCmd.AddCommand(addCmd(), deleteCmd(), setCmd(), getCmd(), showCmd(), listCmd(), lookupCmd())

	if err := rootCmd.Execute(); err != nil {
		os.Exit(ExitUsageError)
	}
}

func addCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "add",
		Short: "Add a new column",
	}
	cmd.AddCommand(addColumnCmd(), addPublicColumnCmd())
	return cmd
}

func addColumnCmd() *cobra.Command {
	var defaultVal string
	cmd := &cobra.Command{
		Use:   "column <column-name>",
		Short: "Add an agent-scope column",
		Args:  cobra.ExactArgs(1),
		Run: func(cmd *cobra.Command, args []string) {
			requirePcguard()
			colName := args[0]

			data, err := readEgoData()
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				os.Exit(ExitJSONError)
			}

			if columnExists(data, colName) {
				fmt.Fprintf(os.Stderr, "Error: column '%s' already exists\n", colName)
				os.Exit(ExitColumnExists)
			}

			data.Columns = append(data.Columns, colName)

			// Set default value for all existing agents
			if defaultVal != "" {
				for agentID := range data.AgentScope {
					data.AgentScope[agentID][colName] = defaultVal
				}
			}

			if err := writeEgoData(data); err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				os.Exit(ExitJSONError)
			}
		},
	}
	cmd.Flags().StringVar(&defaultVal, "default", "", "Default value for the column")
	return cmd
}

func addPublicColumnCmd() *cobra.Command {
	var defaultVal string
	cmd := &cobra.Command{
		Use:   "public-column <column-name>",
		Short: "Add a public-scope column",
		Args:  cobra.ExactArgs(1),
		Run: func(cmd *cobra.Command, args []string) {
			requirePcguard()
			colName := args[0]

			data, err := readEgoData()
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				os.Exit(ExitJSONError)
			}

			if columnExists(data, colName) {
				fmt.Fprintf(os.Stderr, "Error: column '%s' already exists\n", colName)
				os.Exit(ExitColumnExists)
			}

			data.PublicColumns = append(data.PublicColumns, colName)
			if defaultVal != "" {
				data.PublicScope[colName] = defaultVal
			} else {
				data.PublicScope[colName] = ""
			}

			if err := writeEgoData(data); err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				os.Exit(ExitJSONError)
			}
		},
	}
	cmd.Flags().StringVar(&defaultVal, "default", "", "Default value for the column")
	return cmd
}

func deleteCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "delete <column-name>",
		Short: "Delete a column and all its values",
		Args:  cobra.ExactArgs(1),
		Run: func(cmd *cobra.Command, args []string) {
			requirePcguard()
			colName := args[0]

			data, err := readEgoData()
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				os.Exit(ExitJSONError)
			}

			if !columnExists(data, colName) {
				fmt.Fprintf(os.Stderr, "Error: column '%s' does not exist\n", colName)
				os.Exit(ExitColumnNotFound)
			}

			// Remove from public-columns if present
			if isPublicColumn(data, colName) {
				newCols := []string{}
				for _, c := range data.PublicColumns {
					if c != colName {
						newCols = append(newCols, c)
					}
				}
				data.PublicColumns = newCols
				delete(data.PublicScope, colName)
			}

			// Remove from agent columns if present
			if isAgentColumn(data, colName) {
				newCols := []string{}
				for _, c := range data.Columns {
					if c != colName {
						newCols = append(newCols, c)
					}
				}
				data.Columns = newCols
				// Remove from all agent scopes
				for agentID := range data.AgentScope {
					delete(data.AgentScope[agentID], colName)
				}
			}

			if err := writeEgoData(data); err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				os.Exit(ExitJSONError)
			}
		},
	}
	return cmd
}

func setCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "set <column-name> <value>",
		Short: "Set a field value",
		Args:  cobra.ExactArgs(2),
		Run: func(cmd *cobra.Command, args []string) {
			requirePcguard()
			colName := args[0]
			value := args[1]

			data, err := readEgoData()
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				os.Exit(ExitJSONError)
			}

			if !columnExists(data, colName) {
				fmt.Fprintf(os.Stderr, "Error: column '%s' does not exist (use 'ego-mgr add column' or 'ego-mgr add public-column' first)\n", colName)
				os.Exit(ExitColumnNotFound)
			}

			// Auto-register agent
			ensureAgent(data)

			if isPublicColumn(data, colName) {
				data.PublicScope[colName] = value
			} else {
				agentID := currentAgentID()
				data.AgentScope[agentID][colName] = value
			}

			if err := writeEgoData(data); err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				os.Exit(ExitJSONError)
			}
		},
	}
	return cmd
}

func getCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "get <column-name>",
		Short: "Get a field value",
		Args:  cobra.ExactArgs(1),
		Run: func(cmd *cobra.Command, args []string) {
			requirePcguard()
			colName := args[0]

			data, err := readEgoData()
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				os.Exit(ExitJSONError)
			}

			if !columnExists(data, colName) {
				fmt.Fprintf(os.Stderr, "Error: column '%s' does not exist\n", colName)
				os.Exit(ExitColumnNotFound)
			}

			// Auto-register agent
			ensureAgent(data)

			if isPublicColumn(data, colName) {
				fmt.Print(data.PublicScope[colName])
			} else {
				agentID := currentAgentID()
				fmt.Print(data.AgentScope[agentID][colName])
			}
		},
	}
	return cmd
}

func showCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "show",
		Short: "Show all fields and values",
		Run: func(cmd *cobra.Command, args []string) {
			requirePcguard()

			data, err := readEgoData()
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				os.Exit(ExitJSONError)
			}

			// Auto-register agent
			ensureAgent(data)
			agentID := currentAgentID()

			// Print public scope first
			for _, col := range data.PublicColumns {
				val := data.PublicScope[col]
				fmt.Printf("%s: %s\n", col, val)
			}

			// Then agent scope
			for _, col := range data.Columns {
				val := ""
				if agentData, ok := data.AgentScope[agentID]; ok {
					val = agentData[col]
				}
				fmt.Printf("%s: %s\n", col, val)
			}
		},
	}
	return cmd
}

func listCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "list",
		Short: "List information",
	}
	cmd.AddCommand(listColumnsCmd())
	return cmd
}

func listColumnsCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "columns",
		Short: "List all column names",
		Run: func(cmd *cobra.Command, args []string) {
			requirePcguard()

			data, err := readEgoData()
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				os.Exit(ExitJSONError)
			}

			// Print public columns first
			for _, col := range data.PublicColumns {
				fmt.Println(col)
			}

			// Then agent columns
			for _, col := range data.Columns {
				fmt.Println(col)
			}
		},
	}
	return cmd
}

func lookupCmd() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "lookup <username>",
		Short: "Look up an agent ID by default-username",
		Args:  cobra.ExactArgs(1),
		Run: func(cmd *cobra.Command, args []string) {
			username := args[0]

			data, err := readEgoData()
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error: %v\n", err)
				os.Exit(ExitJSONError)
			}

			// Verify default-username column exists
			if !isAgentColumn(data, "default-username") {
				fmt.Fprintf(os.Stderr, "Error: column 'default-username' does not exist\n")
				os.Exit(ExitColumnNotFound)
			}

			for agentID, agentData := range data.AgentScope {
				if agentData["default-username"] == username {
					fmt.Print(agentID)
					return
				}
			}

			fmt.Fprintf(os.Stderr, "Error: no agent found with default-username '%s'\n", username)
			os.Exit(ExitNotFound)
		},
	}
	return cmd
}