Agent-Squad/squad/port.go

package squad

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"
)

// Executor 是小队核心对一次智能体执行的端口。实现放在适配器模块。
type Executor interface {
	Name() string
	Execute(ctx context.Context, req ExecRequest) (*ExecResult, error)
}

// ModelClient 是规划器对一次模型补全的端口。实现放在适配器模块。
type ModelClient interface {
	Complete(ctx context.Context, turn ModelTurn) (string, error)
}

// ModelTurn 是一次规划补全的输入。适配器负责把它变成供应商协议。
type ModelTurn struct {
	Model     string
	System    string
	User      string
	MaxTokens int
	APIBase   string
	APIKey    string
}

const (
	ToolProfileWorkspace = "workspace"
	ToolProfileReasoning = "reasoning"
)

// ExecRequest 是一次步骤执行的输入。字段与 native exec.Request 的 JSON 对齐，类型不共享。
type ExecRequest struct {
	WorkDir       string             `json:"work_dir"`
	Prompt        string             `json:"prompt"`
	Instructions  string             `json:"instructions"`
	Targets       []string           `json:"targets"`
	Timeout       time.Duration      `json:"timeout"`
	ToolProfile   string             `json:"tool_profile,omitempty"`
	OnEvent       func(ProcessEvent) `json:"-"`
	Model         string             `json:"model,omitempty"`
	ToolAllowlist []string           `json:"tool_allowlist,omitempty"`
	Skills        []string           `json:"skills,omitempty"`
	Budget        *ExecBudget        `json:"budget,omitempty"`
	SessionScope  string             `json:"session_scope,omitempty"`
	ExecutionID   string             `json:"execution_id,omitempty"`
	Env           map[string]string  `json:"env,omitempty"`
}

// ExecBudget 是一次步骤执行的共享预算。
type ExecBudget struct {
	MaxPromptTokens     int       `json:"max_prompt_tokens,omitempty"`
	MaxCompletionTokens int       `json:"max_completion_tokens,omitempty"`
	MaxToolCalls        int       `json:"max_tool_calls,omitempty"`
	Deadline            time.Time `json:"deadline,omitempty"`
}

// ProcessEvent 是过程帧。
type ProcessEvent struct {
	Type             string
	Text             string
	Tool             string
	Detail           string
	PromptTokens     int
	CompletionTokens int
}

// ExecResult 是一次步骤执行的终态。
type ExecResult struct {
	Success          bool              `json:"success"`
	Patches          map[string][]byte `json:"patches"`
	Logs             string            `json:"logs"`
	ErrorReason      string            `json:"error_reason"`
	Retryable        bool              `json:"retryable,omitempty"`
	PromptTokens     int               `json:"prompt_tokens,omitempty"`
	CompletionTokens int               `json:"completion_tokens,omitempty"`
}

// StubExecutor 是离线桩，测试与未接模型时使用。
type StubExecutor struct{}

func (s *StubExecutor) Name() string { return "stub" }

func (s *StubExecutor) Execute(ctx context.Context, req ExecRequest) (*ExecResult, error) {
	patches := make(map[string][]byte)
	for _, target := range req.Targets {
		outPath := filepath.Join(req.WorkDir, target)
		_ = os.MkdirAll(filepath.Dir(outPath), 0o755)
		var content []byte
		if strings.HasSuffix(target, ".json") {
			if existing, err := os.ReadFile(outPath); err == nil {
				var dummy any
				if json.Unmarshal(existing, &dummy) == nil {
					content = existing
				}
			}
			if len(content) == 0 {
				content = []byte(`{"status":"mock_ok"}`)
			}
		} else {
			content = []byte(fmt.Sprintf("# 模拟生成结果\n\n- 目标文件: %s\n- 提示词摘要: %s\n", target, req.Prompt))
		}
		_ = os.WriteFile(outPath, content, 0o644)
		patches[target] = content
	}
	return &ExecResult{Success: true, Patches: patches, Logs: "Stub agent execution completed."}, nil
}

// LimitingExecutor 用有界信号量限制底层执行并发。
type LimitingExecutor struct {
	inner Executor
	sem   chan struct{}
}

// WrapWithConcurrencyLimit 包装执行器；max<=0 时默认 8。
func WrapWithConcurrencyLimit(inner Executor, max int) Executor {
	if inner == nil {
		return nil
	}
	if max <= 0 {
		max = 8
	}
	return &LimitingExecutor{inner: inner, sem: make(chan struct{}, max)}
}

func (l *LimitingExecutor) Name() string {
	if l.inner == nil {
		return "limiting"
	}
	return l.inner.Name()
}

func (l *LimitingExecutor) Execute(ctx context.Context, req ExecRequest) (*ExecResult, error) {
	select {
	case l.sem <- struct{}{}:
		defer func() { <-l.sem }()
	case <-ctx.Done():
		return &ExecResult{
			Success:     false,
			ErrorReason: fmt.Sprintf("agent concurrency wait canceled: %v", ctx.Err()),
		}, ctx.Err()
	}
	return l.inner.Execute(ctx, req)
}