Agent-Squad/squad/backend.go

package squad

import (
	"fmt"
	"sync"

	"error-code-system/errcode"
)

// BackendFactory 按角色构造后端：角色级 Model/ToolProfile 在工厂闭包层绑定
// （agent-native ExecRequest 无 Model 字段，绑定不侵入执行契约）。
type BackendFactory func(role RoleSpec) (Executor, error)

// BackendResolver 角色 → 智能体后端解析器。
// 内置 "stub"（agent-native 离线桩）；真实智能体经 RegisterAgentNative 注册
// agent-native 自研内核，产品也可注册 "pi-agent" 等外部智能体工厂。
// 未知名字 fail-closed，绝不静默落 stub。角色声明 MaxConcurrentTasks 时
// 以并发限流包装实现对象级串行（每个角色一份缓存实例，限流跨步骤共享）。
type BackendResolver struct {
	mu         sync.RWMutex
	factories  map[string]BackendFactory
	underlying map[string]Executor // 无工厂直挂的后端实例
	wrapped    map[string]Executor // 角色（含模型绑定）→ 已包装实例缓存
}

// NewBackendResolver 创建解析器并注册内置后端。
func NewBackendResolver() *BackendResolver {
	r := &BackendResolver{
		factories:  make(map[string]BackendFactory),
		underlying: make(map[string]Executor),
		wrapped:    make(map[string]Executor),
	}
	r.RegisterInstance("stub", &StubExecutor{})
	return r
}

// Register 注册后端工厂（名字如 "agent-native"、"pi-agent"）。
func (r *BackendResolver) Register(name string, f BackendFactory) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.factories[name] = f
}

// RegisterInstance 直接注册后端实例（忽略角色级绑定）。
func (r *BackendResolver) RegisterInstance(name string, b Executor) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.underlying[name] = b
}

// ListBackends 返回可用后端名（排序稳定）。
func (r *BackendResolver) ListBackends() []string {
	r.mu.RLock()
	defer r.mu.RUnlock()
	seen := map[string]bool{}
	for n := range r.factories {
		seen[n] = true
	}
	for n := range r.underlying {
		seen[n] = true
	}
	out := make([]string, 0, len(seen))
	for n := range seen {
		out = append(out, n)
	}
	// 简单插入排序，输出稳定。
	for i := 1; i < len(out); i++ {
		for j := i; j > 0 && out[j] < out[j-1]; j-- {
			out[j], out[j-1] = out[j-1], out[j]
		}
	}
	return out
}

// Resolve 解析角色的执行后端：Role.Backend → Options.DefaultBackend → "stub"。
// MaxConcurrentTasks > 1 时以 agent-native 并发限流包装（per-role 缓存，跨步骤共享；
// 创建段持写锁防并行首解析竞态——否则并发步骤会各拿一份限流器）。
func (r *BackendResolver) Resolve(role RoleSpec, defaultBackend string) (Executor, string, error) {
	name := role.Backend
	if name == "" {
		name = defaultBackend
	}
	if name == "" {
		name = "stub"
	}
	cacheKey := name + "\x00" + role.Role + "\x00" + role.Model

	// 后端名快照在锁外获取（Resolve 持锁期间再调 ListBackends 会自死锁）。
	names := r.ListBackends()
	r.mu.RLock()
	cached, ok := r.wrapped[cacheKey]
	factory, hasFactory := r.factories[name]
	under, hasUnder := r.underlying[name]
	r.mu.RUnlock()
	if ok {
		return cached, name, nil
	}

	r.mu.Lock()
	defer r.mu.Unlock()
	// 双检：等写锁期间可能已被并发解析填充。
	if b, ok := r.wrapped[cacheKey]; ok {
		return b, name, nil
	}

	var b Executor
	if hasUnder {
		b = under
	} else if hasFactory {
		var err error
		b, err = factory(role)
		if err != nil {
			return nil, name, err
		}
	} else {
		return nil, name, errcode.New(errcode.SquadUnknownBackend, fmt.Sprintf("unknown backend %q (registered: %v); refusing to silently fall back", name, names))
	}

	// MaxConcurrentTasks>=1 即包装限流（1 = 对象级串行；0 = 不限）。
	if max := role.MaxConcurrentTasks; max >= 1 {
		b = WrapWithConcurrencyLimit(b, max)
	}
	r.wrapped[cacheKey] = b
	return b, name, nil
}