Agent-Squad/squad/backend.go

package squad

import (
	"context"
	"fmt"
	"sync"
	"time"

	"agent-native/adapters/native"
	"agent-native/exec"
	"agent-native/runtime/llm"
	"agent-native/runtime/mcp"
	"agent-native/runtime/subagent"
	"agent-native/runtime/tools"

	"error-code-system/errcode"
)

// BackendFactory 按角色构造后端：角色级 Model/ToolProfile 在工厂闭包层绑定
// （agent-native exec.Request 无 Model 字段，绑定不侵入执行契约）。
type BackendFactory func(role RoleSpec) (exec.Backend, 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]exec.Backend // 无工厂直挂的后端实例
	wrapped    map[string]exec.Backend // 角色（含模型绑定）→ 已包装实例缓存
}

// NewBackendResolver 创建解析器并注册内置后端。
func NewBackendResolver() *BackendResolver {
	r := &BackendResolver{
		factories:  make(map[string]BackendFactory),
		underlying: make(map[string]exec.Backend),
		wrapped:    make(map[string]exec.Backend),
	}
	r.RegisterInstance("stub", &exec.StubBackend{})
	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 exec.Backend) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.underlying[name] = b
}

// AgentNativeConfig 声明 agent-native 后端的全量能力装配（组合根提供）。
// 零值 = 仅基础模型执行：搜索/取回/执行/MCP/技能/委派工具一律不注册（fail-closed）。
type AgentNativeConfig struct {
	Client       *llm.Client
	DefaultModel string
	MaxTurns     int
	Timeout      time.Duration
	// 以下均为 agent-native Options 的直通注入，详见其注释。
	Search      tools.SearchProvider
	Exec        *tools.ExecConfig
	Fetch       *tools.FetchConfig
	MCP         *mcp.Manager
	SkillsRoots []string
	Subagent    subagent.Factory
	Compact     *native.CompactConfig
	// HostTools 注册进内核工作区。与内置工具同名时该步骤失败，不覆盖。
	HostTools []tools.Tool
	// OnEvent 收到内核的逐步过程事件（文本、工具、用量）。与 Request.OnEvent 同时生效。
	OnEvent func(ev exec.ProcessEvent)
	// Stream 测试注入。为空时按 Client 访问模型。
	Stream llm.StreamFunc
	// MaxToolCalls 单步工具调用上限。0 表示不限制。
	MaxToolCalls int
}

// buildAgentNativeOptions 按角色把配置换算为 native.Options。
// BashPermission 三态在此执行：deny = 不装配执行工具（无论组合根是否提供 Exec）；
// ask/allow = 装配（ask 的人工门由引擎 permission 挂起点先行把关）。
func buildAgentNativeOptions(cfg AgentNativeConfig, role RoleSpec) (native.Options, error) {
	model := role.Model
	if model == "" {
		model = cfg.DefaultModel
	}
	if model == "" {
		return native.Options{}, fmt.Errorf("agent-native backend requires a model (role=%s)", role.Role)
	}
	opts := native.Options{
		Model:        model,
		MaxTurns:     cfg.MaxTurns,
		Timeout:      cfg.Timeout,
		Search:       cfg.Search,
		Fetch:        cfg.Fetch,
		MCP:          cfg.MCP,
		SkillsRoots:  cfg.SkillsRoots,
		Subagent:     cfg.Subagent,
		Compact:      cfg.Compact,
		HostTools:    cfg.HostTools,
		Stream:       cfg.Stream,
		MaxToolCalls: cfg.MaxToolCalls,
	}
	if cfg.Client != nil {
		opts.APIBase, opts.APIKey = cfg.Client.BaseURL, cfg.Client.APIKey
	}
	// 显式 allow/ask 才装配（注册层默认 deny，但派生/程序构造的 RoleSpec
	// 未必经注册——工厂层兜底，空白串不放行）。
	if cfg.Exec != nil && (role.BashPermission == BashPermissionAllow || role.BashPermission == BashPermissionAsk) {
		opts.Exec = cfg.Exec
	}
	return opts, nil
}

// processEventBackend 把配置级过程回调接到每次 Execute，同时保留请求自带的 OnEvent。
type processEventBackend struct {
	next exec.Backend
	on   func(ev exec.ProcessEvent)
}

func (b *processEventBackend) Name() string { return b.next.Name() }

func (b *processEventBackend) Execute(ctx context.Context, req exec.Request) (*exec.Result, error) {
	if b.on != nil {
		prev := req.OnEvent
		req.OnEvent = func(ev exec.ProcessEvent) {
			b.on(ev)
			if prev != nil {
				prev(ev)
			}
		}
	}
	return b.next.Execute(ctx, req)
}

// RegisterAgentNative 注册 "agent-native" 后端工厂（基础签名，向后兼容）：
// 角色级 Model 覆盖 defaultModel，无附加能力装配。
func (r *BackendResolver) RegisterAgentNative(client *llm.Client, defaultModel string, maxTurns int, timeout time.Duration) {
	r.RegisterAgentNativeFull(AgentNativeConfig{
		Client: client, DefaultModel: defaultModel, MaxTurns: maxTurns, Timeout: timeout,
	})
}

// RegisterAgentNativeFull 注册带全量能力装配的 "agent-native" 工厂：
// 搜索/取回/执行/MCP/技能/委派经配置注入；BashPermission=deny 的角色
// 不装配执行工具。智能体执行全部委托 agent-native，本库不实现 Agent。
func (r *BackendResolver) RegisterAgentNativeFull(cfg AgentNativeConfig) {
	r.Register("agent-native", func(role RoleSpec) (exec.Backend, error) {
		opts, err := buildAgentNativeOptions(cfg, role)
		if err != nil {
			return nil, err
		}
		next := native.New(opts)
		if cfg.OnEvent == nil {
			return next, nil
		}
		return &processEventBackend{next: next, on: cfg.OnEvent}, nil
	})
}

// 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) (exec.Backend, 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 exec.Backend
	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 = exec.WrapWithConcurrencyLimit(b, max)
	}
	r.wrapped[cacheKey] = b
	return b, name, nil
}