Agent-Squad/squad/registry.go

package squad

import (
	"fmt"
	"sort"
	"strings"
	"sync"
)

// RoleRegistry 角色注册中心：空表起步，无任何预设角色。
// 线程安全；角色可来自配置文件、HTTP API、CLI 或规划器运行期派生。
type RoleRegistry struct {
	mu    sync.RWMutex
	roles map[string]RoleSpec
	// OnChange 可选钩子：角色注册/派生后回调（store 层用于落盘）。
	OnChange func(spec RoleSpec)
}

// NewRoleRegistry 创建空注册表。
func NewRoleRegistry() *RoleRegistry {
	return &RoleRegistry{roles: make(map[string]RoleSpec)}
}

// Register 注册或更新角色；Role 为空或 BashPermission 非法时返回错误。
func (r *RoleRegistry) Register(spec RoleSpec) error {
	if strings.TrimSpace(spec.Role) == "" {
		return fmt.Errorf("role key cannot be empty")
	}
	switch spec.BashPermission {
	case "", "deny", "ask", "allow":
	default:
		return fmt.Errorf("bash_permission must be allow/ask/deny, got %q", spec.BashPermission)
	}
	if spec.BashPermission == "" {
		spec.BashPermission = "deny"
	}
	if spec.Source == "" {
		spec.Source = RoleSourceRegistered
	}
	if spec.Name == "" {
		spec.Name = spec.Role
	}
	r.mu.Lock()
	r.roles[spec.Role] = spec
	r.mu.Unlock()
	if r.OnChange != nil {
		r.OnChange(spec)
	}
	return nil
}

// Get 按 Key 获取角色。
func (r *RoleRegistry) Get(key string) (RoleSpec, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()
	spec, ok := r.roles[key]
	return spec, ok
}

// List 列出全部角色（按 Role 排序，输出稳定）。
func (r *RoleRegistry) List() []RoleSpec {
	r.mu.RLock()
	defer r.mu.RUnlock()
	out := make([]RoleSpec, 0, len(r.roles))
	for _, spec := range r.roles {
		out = append(out, spec)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Role < out[j].Role })
	return out
}

// Remove 删除角色；不存在时返回 false。
func (r *RoleRegistry) Remove(key string) bool {
	r.mu.Lock()
	_, ok := r.roles[key]
	if ok {
		delete(r.roles, key)
	}
	r.mu.Unlock()
	return ok
}

// Recommend 按查询文本推荐最贴切角色：Keywords 命中计 10 分/词，
// Responsibilities 命中计 5 分；无命中返回 false（不猜测兜底角色）。
func (r *RoleRegistry) Recommend(query string) (RoleSpec, bool) {
	q := strings.ToLower(query)
	var best RoleSpec
	bestScore := 0
	for _, spec := range r.List() {
		score := 0
		for _, kw := range spec.Keywords {
			if kw != "" && strings.Contains(q, strings.ToLower(kw)) {
				score += 10
			}
		}
		for _, resp := range spec.Responsibilities {
			if resp != "" && strings.Contains(q, strings.ToLower(resp)) {
				score += 5
			}
		}
		if score > bestScore {
			best, bestScore = spec, score
		}
	}
	return best, bestScore > 0
}

// DeriveRequest 派生角色请求：以 BaseRole 为基线叠加领域特化。
type DeriveRequest struct {
	RoleKey          string   `json:"role_key"`
	Name             string   `json:"name,omitempty"`
	BaseRole         string   `json:"base_role,omitempty"`
	Backend          string   `json:"backend,omitempty"`
	Model            string   `json:"model,omitempty"`
	ToolProfile      string   `json:"tool_profile,omitempty"`
	Tools            []string `json:"tools,omitempty"`
	Skills           []string `json:"skills,omitempty"`
	BashPermission   string   `json:"bash_permission,omitempty"`
	DomainTopic      string   `json:"domain_topic,omitempty"`
	Responsibilities []string `json:"responsibilities,omitempty"`
	ForbiddenActions []string `json:"forbidden_actions,omitempty"`
	Instructions     string   `json:"instructions,omitempty"`
}

// Derive 派生特设角色：继承基线角色的 Backend/Model/Tools/指令，
// 叠加领域职责与禁忌；BaseRole 缺省取注册表中第一个角色（无角色则内联最小基线）。
func (r *RoleRegistry) Derive(req DeriveRequest) (RoleSpec, error) {
	if strings.TrimSpace(req.RoleKey) == "" {
		return RoleSpec{}, fmt.Errorf("role_key cannot be empty")
	}
	base, hasBase := r.Get(req.BaseRole)
	if !hasBase {
		list := r.List()
		if len(list) > 0 {
			base = list[0]
		}
	}
	spec := RoleSpec{
		Role:             req.RoleKey,
		Name:             req.Name,
		Source:           RoleSourceDerived,
		BaseRole:         base.Role,
		Backend:          firstNonEmpty(req.Backend, base.Backend),
		Model:            firstNonEmpty(req.Model, base.Model),
		ToolProfile:      firstNonEmpty(req.ToolProfile, base.ToolProfile),
		Tools:            firstNonEmptySlice(req.Tools, base.Tools),
		Skills:           firstNonEmptySlice(req.Skills, base.Skills),
		BashPermission:   firstNonEmpty(req.BashPermission, base.BashPermission),
		Responsibilities: append([]string{}, req.Responsibilities...),
		ForbiddenActions: append([]string{}, req.ForbiddenActions...),
		Instructions:     firstNonEmpty(req.Instructions, base.Instructions),
		InstructionsFile: base.InstructionsFile,
	}
	if spec.Name == "" {
		spec.Name = fmt.Sprintf("%s专员", req.RoleKey)
	}
	if len(spec.Responsibilities) == 0 && req.DomainTopic != "" {
		spec.Responsibilities = append(spec.Responsibilities, fmt.Sprintf("专责推演与论证【%s】", req.DomainTopic))
	}
	spec.ForbiddenActions = append(spec.ForbiddenActions,
		"严禁违规越权执行未经规划的破坏性命令", "严禁直接代用户做终审拍板")
	if err := r.Register(spec); err != nil {
		return RoleSpec{}, err
	}
	return spec, nil
}

// firstNonEmptySlice 首个非空切片，否则回退后者。
func firstNonEmptySlice(a, b []string) []string {
	if len(a) > 0 {
		return append([]string{}, a...)
	}
	return append([]string{}, b...)
}

func firstNonEmpty(vals ...string) string {
	for _, v := range vals {
		if v != "" {
			return v
		}
	}
	return ""
}