Agent-Squad/adapters/native/register.go

// Package native 把 agent-native 接成小队的 Executor 与 ModelClient。
package native

import (
	"context"
	"fmt"
	"strings"
	"time"

	nativexec "agent-native/adapters/native"
	"agent-native/exec"
	"agent-native/runtime"
	"agent-native/runtime/checkpoint"
	"agent-native/runtime/hooks"
	"agent-native/runtime/llm"
	"agent-native/runtime/mcp"
	"agent-native/runtime/subagent"
	"agent-native/runtime/tools"
	"agent-squad/squad"
)

// Config 是组合根交给 agent-native 后端的装配。零值只做模型执行。
type Config struct {
	Client            *llm.Client
	DefaultModel      string
	MaxTurns          int
	Timeout           time.Duration
	Search            tools.SearchProvider
	Exec              *tools.ExecConfig
	Fetch             *tools.FetchConfig
	MCP               *mcp.Manager
	SkillsRoots       []string
	Subagent          subagent.Factory
	Compact           *nativexec.CompactConfig
	HostTools         []tools.Tool
	OnEvent           func(ev squad.ProcessEvent)
	Stream            llm.StreamFunc
	MaxToolCalls      int
	BeforeTool        func(call llm.ToolCall) error
	AfterTool         func(call llm.ToolCall, res llm.Message)
	Stop              func(last llm.Message, info runtime.StopInfo) runtime.StopDecision
	Reminders         func(turn int) []string
	RepeatGuard       bool
	Hooks             *hooks.Bus
	Checkpoints       *checkpoint.Log
	CheckpointAnchor  func() string
	Trace             func(name string, attrs map[string]string) func()
	CodeOrchestration bool
	Ask               func(ctx context.Context, question string) (string, error)
	// PermissionAsk 把内核钩子的 ask 决策转交宿主裁决（nil = 内核保持
	// ask 按拒绝处理的旧行为）。经 nativexec.GateOptions 透传。
	PermissionAsk func(ctx context.Context, call llm.ToolCall) runtime.AskVerdict
}

// Register 把 "agent-native" 工厂注册进解析器。
func Register(r *squad.BackendResolver, client *llm.Client, defaultModel string, maxTurns int, timeout time.Duration) {
	RegisterFull(r, Config{Client: client, DefaultModel: defaultModel, MaxTurns: maxTurns, Timeout: timeout})
}

// RegisterFull 注册带工具装配的 "agent-native" 工厂。
func RegisterFull(r *squad.BackendResolver, cfg Config) {
	r.Register("agent-native", func(role squad.RoleSpec) (squad.Executor, error) {
		opts, err := buildOptions(cfg, role)
		if err != nil {
			return nil, err
		}
		next := &bridge{next: nativexec.New(opts)}
		if cfg.OnEvent == nil {
			return next, nil
		}
		return &eventTap{next: next, on: cfg.OnEvent}, nil
	})
}

func buildOptions(cfg Config, role squad.RoleSpec) (nativexec.Options, error) {
	model := role.Model
	if model == "" {
		model = cfg.DefaultModel
	}
	if model == "" {
		return nativexec.Options{}, fmt.Errorf("agent-native backend requires a model (role=%s)", role.Role)
	}
	opts := nativexec.Options{
		ModelOptions: nativexec.ModelOptions{
			Model: model, MaxTurns: cfg.MaxTurns, Timeout: cfg.Timeout, Stream: cfg.Stream,
		},
		ToolOptions: nativexec.ToolOptions{
			Search: cfg.Search, Fetch: cfg.Fetch, MCP: cfg.MCP,
			SkillsRoots: cfg.SkillsRoots, Subagent: cfg.Subagent, HostTools: cfg.HostTools,
			Checkpoints: cfg.Checkpoints, CheckpointAnchor: cfg.CheckpointAnchor,
			CodeOrchestration: cfg.CodeOrchestration,
		},
		ObserveOptions: nativexec.ObserveOptions{Trace: cfg.Trace},
		SessionOptions: nativexec.SessionOptions{Compact: cfg.Compact, MaxToolCalls: cfg.MaxToolCalls},
		GateOptions: nativexec.GateOptions{
			BeforeTool: cfg.BeforeTool, AfterTool: cfg.AfterTool, Stop: cfg.Stop,
			Reminders: cfg.Reminders, RepeatGuard: cfg.RepeatGuard, Hooks: cfg.Hooks,
			Ask: cfg.Ask, PermissionAsk: cfg.PermissionAsk,
		},
	}
	if cfg.Client != nil {
		opts.APIBase, opts.APIKey = cfg.Client.BaseURL, cfg.Client.APIKey
	}
	if cfg.Exec != nil && (role.BashPermission == squad.BashPermissionAllow || role.BashPermission == squad.BashPermissionAsk) {
		opts.Exec = cfg.Exec
	}
	return opts, nil
}

type bridge struct{ next exec.Backend }

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

func (b *bridge) Execute(ctx context.Context, req squad.ExecRequest) (*squad.ExecResult, error) {
	res, err := b.next.Execute(ctx, toNative(req))
	if res == nil {
		return nil, err
	}
	return fromNative(res), err
}

func toNative(req squad.ExecRequest) exec.Request {
	out := exec.Request{
		WorkDir: req.WorkDir, Prompt: req.Prompt, Instructions: req.Instructions,
		Targets: req.Targets, Timeout: req.Timeout, ToolProfile: req.ToolProfile,
		Model: req.Model, ToolAllowlist: req.ToolAllowlist, Skills: req.Skills,
		SessionScope: req.SessionScope, ExecutionID: req.ExecutionID, Env: req.Env,
	}
	if req.Budget != nil {
		out.Budget = &exec.Budget{
			MaxPromptTokens: req.Budget.MaxPromptTokens, MaxCompletionTokens: req.Budget.MaxCompletionTokens,
			MaxToolCalls: req.Budget.MaxToolCalls, Deadline: req.Budget.Deadline,
		}
	}
	if req.OnEvent != nil {
		prev := req.OnEvent
		out.OnEvent = func(ev exec.ProcessEvent) {
			prev(squad.ProcessEvent{
				Type: ev.Type, Text: ev.Text, Tool: ev.Tool, Detail: ev.Detail,
				PromptTokens: ev.PromptTokens, CompletionTokens: ev.CompletionTokens,
				LastPromptTokens: ev.LastPromptTokens, LastCompletionTokens: ev.LastCompletionTokens,
				CacheReadTokens: ev.CacheReadTokens, CacheWriteTokens: ev.CacheWriteTokens,
			})
		}
	}
	return out
}

func fromNative(res *exec.Result) *squad.ExecResult {
	return &squad.ExecResult{
		Success: res.Success, Patches: res.Patches, Logs: res.Logs, ErrorReason: res.ErrorReason,
		Retryable: res.Retryable, PromptTokens: res.PromptTokens, CompletionTokens: res.CompletionTokens,
	}
}

type eventTap struct {
	next squad.Executor
	on   func(squad.ProcessEvent)
}

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

func (b *eventTap) Execute(ctx context.Context, req squad.ExecRequest) (*squad.ExecResult, error) {
	prev := req.OnEvent
	req.OnEvent = func(ev squad.ProcessEvent) {
		b.on(ev)
		if prev != nil {
			prev(ev)
		}
	}
	return b.next.Execute(ctx, req)
}

// LLM 把 llm.Client 适配成 ModelClient。
type LLM struct{ Inner *llm.Client }

func (c LLM) Complete(ctx context.Context, turn squad.ModelTurn) (string, error) {
	if c.Inner == nil {
		return "", fmt.Errorf("llm planner: no model client")
	}
	msgs := []llm.Message{
		{Role: llm.RoleSystem, Content: turn.System},
		{Role: llm.RoleUser, Content: turn.User},
	}
	opts := llm.StreamOptions{MaxTokens: turn.MaxTokens, APIBase: turn.APIBase, APIKey: turn.APIKey}
	var sb strings.Builder
	for ev := range c.Inner.Stream(ctx, turn.Model, msgs, nil, opts) {
		if ev.Kind == llm.KindError {
			return "", fmt.Errorf("%s", ev.ErrorMessage)
		}
		if ev.Kind == llm.KindTextDelta {
			sb.WriteString(ev.Text)
		}
		if ev.Kind == llm.KindMessageEnd {
			if ev.Message.IsError {
				return "", fmt.Errorf("assistant error: %s", ev.Message.Content)
			}
			if sb.Len() == 0 {
				sb.WriteString(ev.Message.Content)
			}
			break
		}
	}
	return sb.String(), nil
}