Agent-Squad/squad/backend.go

package squad

import (
	"context"
	"fmt"
	"strings"
	"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 时
// 以并发限流包装实现对象级串行（每个角色一份缓存实例，限流跨步骤共享）。
//
// 缓存键含 ConfigIdentity 与工厂世代，故 BashPermission/工具配置收紧或工厂
// 替换在下一次 Resolve 即生效。并发 gate 按「后端名+角色」独立于执行器版本
// 持有：配置换代不会另起限流器而绕过角色上限。
//
// 边界：已在飞的任务不会被回滚或取消；收紧权限/降低并发只影响后续 Resolve
// 与后续 Acquire。需要即时撤权的宿主应在每次工具调用前再校验。
type BackendResolver struct {
	mu         sync.RWMutex
	factories  map[string]BackendFactory
	underlying map[string]Executor // 无工厂直挂的后端实例
	wrapped    map[string]Executor // 执行器配置身份 → 已包装实例缓存
	factoryGen map[string]uint64   // Register/RegisterInstance 世代，替换即递增
	gates      map[string]*roleConcurrencyGate
}

// NewBackendResolver 创建解析器并注册内置后端。
func NewBackendResolver() *BackendResolver {
	r := &BackendResolver{
		factories:  make(map[string]BackendFactory),
		underlying: make(map[string]Executor),
		wrapped:    make(map[string]Executor),
		factoryGen: make(map[string]uint64),
		gates:      make(map[string]*roleConcurrencyGate),
	}
	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
	r.invalidateLocked(name)
}

// RegisterInstance 直接注册后端实例（忽略角色级绑定）。
// 替换同名实例会使该后端下已缓存的包装执行器失效；"stub" 语义保持不变。
func (r *BackendResolver) RegisterInstance(name string, b Executor) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.underlying[name] = b
	r.invalidateLocked(name)
}

func (r *BackendResolver) invalidateLocked(name string) {
	r.factoryGen[name]++
	prefix := name + "\x00"
	for k := range r.wrapped {
		if strings.HasPrefix(k, prefix) {
			delete(r.wrapped, k)
		}
	}
	// gates 按角色身份保留：换工厂/配置不得丢掉进行中的并发计数。
}

// 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 时以共享角色 gate 包装（跨执行器版本共用；
// 创建段持写锁防并行首解析竞态——否则并发步骤会各拿一份执行器）。
func (r *BackendResolver) Resolve(role RoleSpec, defaultBackend string) (Executor, string, error) {
	name := role.Backend
	if name == "" {
		name = defaultBackend
	}
	if name == "" {
		name = "stub"
	}

	// 后端名快照在锁外获取（Resolve 持锁期间再调 ListBackends 会自死锁）。
	names := r.ListBackends()

	r.mu.RLock()
	gen := r.factoryGen[name]
	cacheKey := resolveCacheKey(name, role, gen)
	cached, ok := r.wrapped[cacheKey]
	r.mu.RUnlock()
	if ok {
		if max := role.MaxConcurrentTasks; max >= 1 {
			r.mu.Lock()
			if g := r.gates[limiterKey(name, role.Role)]; g != nil {
				g.setMax(max)
			}
			r.mu.Unlock()
		}
		return cached, name, nil
	}

	r.mu.Lock()
	defer r.mu.Unlock()
	gen = r.factoryGen[name]
	cacheKey = resolveCacheKey(name, role, gen)
	// 双检：等写锁期间可能已被并发解析填充。
	if b, ok := r.wrapped[cacheKey]; ok {
		if max := role.MaxConcurrentTasks; max >= 1 {
			if g := r.gates[limiterKey(name, role.Role)]; g != nil {
				g.setMax(max)
			}
		}
		return b, name, nil
	}

	var b Executor
	if under, hasUnder := r.underlying[name]; hasUnder {
		b = under
	} else if factory, hasFactory := r.factories[name]; 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 = 不限）。
	// gate 按角色身份共享，与 ConfigIdentity / 工厂世代无关。
	if max := role.MaxConcurrentTasks; max >= 1 {
		g := r.gateLocked(name, role.Role, max)
		b = &gatedExecutor{inner: b, gate: g}
	}
	r.wrapped[cacheKey] = b
	return b, name, nil
}

func resolveCacheKey(name string, role RoleSpec, gen uint64) string {
	return name + "\x00" + role.Role + "\x00" + ConfigIdentity(role) + "\x00" +
		fmt.Sprintf("%d", gen) + "\x00" + limiterStamp(role.MaxConcurrentTasks)
}

func limiterKey(backend, role string) string {
	return backend + "\x00" + role
}

func (r *BackendResolver) gateLocked(backend, role string, max int) *roleConcurrencyGate {
	key := limiterKey(backend, role)
	g, ok := r.gates[key]
	if !ok {
		g = newRoleConcurrencyGate(max)
		r.gates[key] = g
		return g
	}
	g.setMax(max)
	return g
}

// roleConcurrencyGate 是按角色共享的并发上限。
// 执行器配置换代仍复用同一 gate，故旧版本在飞任务计入同一上限。
type roleConcurrencyGate struct {
	mu     sync.Mutex
	max    int
	active int
	notify chan struct{} // 关闭并替换以唤醒等待者
}

func newRoleConcurrencyGate(max int) *roleConcurrencyGate {
	if max <= 0 {
		max = 1
	}
	return &roleConcurrencyGate{max: max, notify: make(chan struct{})}
}

func (g *roleConcurrencyGate) setMax(max int) {
	if max <= 0 {
		max = 1
	}
	g.mu.Lock()
	defer g.mu.Unlock()
	g.max = max
	g.signalLocked()
}

func (g *roleConcurrencyGate) signalLocked() {
	close(g.notify)
	g.notify = make(chan struct{})
}

func (g *roleConcurrencyGate) acquire(ctx context.Context) error {
	for {
		g.mu.Lock()
		if g.active < g.max {
			g.active++
			g.mu.Unlock()
			return nil
		}
		notify := g.notify
		g.mu.Unlock()
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-notify:
		}
	}
}

func (g *roleConcurrencyGate) release() {
	g.mu.Lock()
	defer g.mu.Unlock()
	if g.active > 0 {
		g.active--
	}
	g.signalLocked()
}

// gatedExecutor 把共享角色 gate 套在具体执行器上。
type gatedExecutor struct {
	inner Executor
	gate  *roleConcurrencyGate
}

func (e *gatedExecutor) Name() string {
	if e.inner == nil {
		return "gated"
	}
	return e.inner.Name()
}

func (e *gatedExecutor) Execute(ctx context.Context, req ExecRequest) (*ExecResult, error) {
	if err := e.gate.acquire(ctx); err != nil {
		return &ExecResult{
			Success:     false,
			ErrorReason: fmt.Sprintf("agent concurrency wait canceled: %v", err),
		}, err
	}
	defer e.gate.release()
	return e.inner.Execute(ctx, req)
}