YunShu/registry.go

package main

import (
	"fmt"
	"os"
	"path/filepath"
	"strings"
)

type AgentRegistry struct {
	mains map[string]*AgentDef
	subs  map[string]*AgentDef
}

func ScanAgents() *AgentRegistry {
	r := &AgentRegistry{
		mains: make(map[string]*AgentDef),
		subs:  make(map[string]*AgentDef),
	}

	dirs := []string{
		"agents",
		filepath.Join(ConfigDir(), "agents"),
	}

	seen := make(map[string]bool)

	for _, dir := range dirs {
		entries, err := os.ReadDir(dir)
		if err != nil {
			continue
		}
		for _, e := range entries {
			if e.IsDir() || !strings.HasSuffix(e.Name(), ".md") || seen[e.Name()] {
				continue
			}
			seen[e.Name()] = true

			agentPath := filepath.Join(dir, e.Name())
			def, err := LoadAgent(agentPath)
			if err != nil {
				warnLog("跳过 agent", "file", e.Name(), "err", err)
				continue
			}

			switch def.Type {
			case "main":
				r.mains[def.Name] = def
			case "sub":
				r.subs[def.Name] = def
			}
		}
	}

	return r
}

func (r *AgentRegistry) GetMain(name string) *AgentDef {
	return r.mains[name]
}

func (r *AgentRegistry) GetSub(name string) *AgentDef {
	return r.subs[name]
}

func (r *AgentRegistry) ListMains() []*AgentDef {
	list := make([]*AgentDef, 0, len(r.mains))
	for _, def := range r.mains {
		list = append(list, def)
	}
	return list
}

func (r *AgentRegistry) ListSubs() []*AgentDef {
	list := make([]*AgentDef, 0, len(r.subs))
	for _, def := range r.subs {
		list = append(list, def)
	}
	return list
}

func (r *AgentRegistry) String() string {
	var b strings.Builder
	b.WriteString(fmt.Sprintf("mains: %d, subs: %d\n", len(r.mains), len(r.subs)))
	for _, def := range r.mains {
		b.WriteString(fmt.Sprintf("  [main] %s: %s\n", def.Name, def.Description))
	}
	for _, def := range r.subs {
		b.WriteString(fmt.Sprintf("  [sub]  %s: %s\n", def.Name, def.Description))
	}
	return b.String()
}