Cyrene_v2/internal/runtime/orchestration/spawntool.go

package orchestration

import (
	"context"
	"fmt"
	"io"
	"log/slog"
	"strings"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/loop"
)

// 子代理派发工具（design/05 §1 创建, §7 并行与安全）的常量。
const (
	// SpawnToolName 是 agent 调用以派发后台子代理的工具的注册名。
	SpawnToolName = "subagent_spawn"

	// DefaultSpawnWait 是前台派发在把回合交还模型之前等待其批次的时间。它沿用运行时
	// 既定的「前台运行，超过 2 秒自动移入后台」阈值
	// （DevDocs/design/04-运行时.md §8.5, internal/runtime/task）。子代理无论如何都
	// 继续运行；变的只是工具结果是否已带上它们的输出。
	DefaultSpawnWait = 2 * time.Second

	// maxSpawnBatch 限制一次派发最多创建多少个子代理。它约束单次模型调用能堆积的
	// 队列长度（design/05 §7 防失控）；并行执行上限本身仍由调度器的 MaxParallel 负责。
	maxSpawnBatch = 8

	// maxSpawnTextRunes 限制交回派发模型的单个子代理结论文本的字符数，使失控的子代理
	// 无法淹没前台会话的上下文。调度器在其快照上保留完整文本。
	maxSpawnTextRunes = 4000

	// minSpawnTimeoutSeconds 与 maxSpawnTimeoutSeconds 界定按调用覆盖
	// timeout_seconds 的合法区间（design/05 §7 裁决 02.4.4 的「按调用覆盖」）：
	// 下限挡住「立即超时」的畸形调用，上限挡住「永远不超时」的失控委托。
	minSpawnTimeoutSeconds = 1
	maxSpawnTimeoutSeconds = 24 * 60 * 60

	// SubagentResultSubject 是后台批次结算后广播其结果的主题，载荷是
	// *SubagentResultEvent（父会话 + 批次结果）。它是 design/05 §9「结果回填主
	// 会话并提示用户」的载体，消费者是 gateway 的中继：该主题被登记为 WS 下行帧
	// subagent_result（Docs/api-client.md §3.10），帧按载荷里的父会话归属，因此结果
	// 落回父会话的实时流而非悬空 hook；逐代理进行中状态仍走调度器进度事件。
	SubagentResultSubject = "subagent.result"
)

// SpawnTool 是让运行中的 agent 派发子代理的运行时入口（design/05 §1）。它是
// (*Scheduler).Spawn 唯一的生产调用方：模型像调用任何其他工具一样调用它，调度器施加
// 并行上限、按代理时限与中止机制，工具则交回合并结果，或者 —— 前台窗口内批次仍在跑时
// —— 一条派发确认。
//
// 它可安全并发使用，且不持有任何按调用的状态。
type SpawnTool struct {
	sched   *Scheduler
	wait    time.Duration
	timeout time.Duration
	log     *slog.Logger
}

// SpawnOption 定制由 NewSpawnTool 构建的 SpawnTool。
type SpawnOption func(*SpawnTool)

// WithSpawnWait 覆盖前台派发等待其批次的时间。非正值保持 DefaultSpawnWait。
func WithSpawnWait(d time.Duration) SpawnOption {
	return func(t *SpawnTool) {
		if d > 0 {
			t.wait = d
		}
	}
}

// WithSpawnLogger 替换该工具的 logger（标注来源 "orchestration/spawn"）。
func WithSpawnLogger(l *slog.Logger) SpawnOption {
	return func(t *SpawnTool) { t.log = spawnLog(l) }
}

// WithSpawnTimeout 设置派发层的默认子代理时限（design/05 §7 裁决 02.4.4
// 「agent_timeout 可配置」的接法）：正数时作为未按调用覆盖的批次默认，
// 覆盖调度器的 PerAgentTimeout；非正保持「不覆盖」，沿用调度器默认。
// 生产接线用 SUBAGENT_TIMEOUT_SEC（默认 10min）填充它。
func WithSpawnTimeout(d time.Duration) SpawnOption {
	return func(t *SpawnTool) {
		if d > 0 {
			t.timeout = d
		}
	}
}

// NewSpawnTool 在 sched 上构建派发工具。调度器为 nil 也被容忍：Execute 随后以
// ErrNoScheduler 快速失败，而不是 panic。
func NewSpawnTool(sched *Scheduler, opts ...SpawnOption) *SpawnTool {
	t := &SpawnTool{sched: sched, wait: DefaultSpawnWait, log: spawnLog(nil)}
	for _, opt := range opts {
		if opt != nil {
			opt(t)
		}
	}
	return t
}

// Name 返回工具标识符。
func (t *SpawnTool) Name() string { return SpawnToolName }

// Description 向模型解释该工具：什么是子代理、它在隔离上下文中并行运行、慢批次会被
// 交给后台而不阻塞对话（design/05 §1/§2）。
func (t *SpawnTool) Description() string {
	return "Dispatch one or more sub-agents to work in parallel on independent parts of a task. " +
		"Each sub-agent runs in its own isolated context and returns a summary plus its result. " +
		"A batch that does not finish within 2 seconds keeps running in the background and reports its " +
		"results later; prefer several sub-agents over one serial chain for independent work."
}

// Parameters 返回输入的 JSON Schema：单个 goal 简写，或用于并行批次的 tasks 数组，
// 外加可选的 persona（子代理运行时所充当的角色，design/22 §8-2）、可选的 background
// 标志、可选的按派发工具白名单（design/05 §7）与可选的按调用超时覆盖。
func (t *SpawnTool) Parameters() (map[string]any, error) {
	return map[string]any{
		"type": "object",
		"properties": map[string]any{
			"goal":            spawnStringProp("Objective the sub-agent works toward alone. Required unless tasks is given."),
			"persona":         spawnPersonaProp(),
			"role":            spawnStringProp("Role/persona id to bind the sub-agent to (optional). Superseded by persona, which is the role the sub-agent actually runs as."),
			"mode":            spawnStringProp("Behavior mode for the role, e.g. engineer (optional)."),
			"attach_info":     spawnStringProp("Extra context or data to hand the sub-agent (optional)."),
			"tasks":           spawnTasksProp(),
			"tool_whitelist":  spawnWhitelistProp("Applies to every sub-agent of this dispatch when a task does not set its own."),
			"timeout_seconds": spawnTimeoutProp(t.timeout),
			"background": map[string]any{
				"type":        "boolean",
				"description": "Dispatch and return immediately without waiting for the foreground window (optional, default false).",
			},
		},
		"anyOf": []any{
			map[string]any{"required": []any{"goal"}},
			map[string]any{"required": []any{"tasks"}},
		},
	}, nil
}

// Execute 派发所请求的子代理，随后要么返回它们的合并结果（批次在前台窗口内结算），
// 要么返回一条派发确认（批次已转入后台，稍后经总线报告）。
//
// 子代理在一个脱离的上下文上派发：它们的生命周期属于自己（受调度器的按代理时限与
// Abort/CancelAll 约束），从不是派发回合的，因此慢批次不会在工具返回的那一刻被杀
// （design/04 §2 后台不阻塞前台）。上下文值被保留，使子代理 span —— 以及它作为根的
// 泳道 —— 仍串在派发它的那一步之下。
func (t *SpawnTool) Execute(ctx context.Context, args any) (any, error) {
	m, ok := args.(map[string]any)
	if !ok {
		return nil, fmt.Errorf("%s: args must be an object", SpawnToolName)
	}
	if t.sched == nil {
		return nil, fmt.Errorf("%s: %w", SpawnToolName, ErrNoScheduler)
	}
	specs, err := spawnSpecs(m, loop.FlowSessionFromContext(ctx), t.timeout)
	if err != nil {
		return nil, fmt.Errorf("%s: %w", SpawnToolName, err)
	}
	// flow 归属桶（`{会话}[#{路}]`）另取一次：ParentSession 必须是**真会话 id**
	// （后台批次的 `subagent_result` 帧按它归属到父会话），而图归属要精确到**哪一路**，
	// 否则真并行时两路子代理会混进同一张图（design/15 §2.5.10-1）。
	owner := loop.FlowOwnerFromContext(ctx)
	for i := range specs {
		specs[i].FlowOwner = owner
	}
	// 在任何东西启动之前就拒绝未知的目标角色：批次必须是全有或全无，否则已经派发出的
	// 子代理会继续运行，而它们的结果永远不上报给模型（design/22 §8-2: 未知角色明确报错，
	// 不静默回落到默认角色).
	if perr := t.sched.CheckPersonas(specs); perr != nil {
		return nil, fmt.Errorf("%s: %w", SpawnToolName, perr)
	}

	spawnCtx := context.WithoutCancel(ctx)
	ids := make([]string, 0, len(specs))
	for _, spec := range specs {
		sub, serr := t.sched.Spawn(spawnCtx, spec)
		if serr != nil {
			return nil, fmt.Errorf("%s: %w", SpawnToolName, serr)
		}
		ids = append(ids, sub.ID)
	}
	t.log.Info("sub-agents dispatched", "count", len(ids), "role", specs[0].RoleID,
		"persona", specs[0].PersonaID, "mode", specs[0].Mode, "session", specs[0].ParentSession)

	if boolArg(m, "background") {
		go t.await(ids, specs[0].ParentSession)
		return runningSpawnResult(ids, specs), nil
	}
	return t.foreground(ids, specs), nil
}

// foreground 至多等待 t.wait 让已派发的批次结算。结算的批次被合并进工具结果；窗口
// 关闭时仍在运行的批次改到后台宣告，并立即被确认。
func (t *SpawnTool) foreground(ids []string, specs []AgentSpec) any {
	waitCtx, cancel := context.WithTimeout(context.Background(), t.wait)
	defer cancel()
	res, err := t.sched.Wait(waitCtx, ids)
	if err != nil {
		t.log.Info("sub-agent batch moved to background", "count", len(ids), "waited", t.wait.String())
		go t.await(ids, specs[0].ParentSession)
		return runningSpawnResult(ids, specs)
	}
	return doneSpawnResult(res)
}

// await 阻塞到批次结算，并把合并结果作为 design/05 §2「结束后写执行总结 → 返回主会话」
// 的交接广播出去。它运行在脱离的 goroutine 上；总线为 nil 或 id 未知时安静结束。
// parentSession 是该批次所属的父会话，随载荷一起广播，使中继能把回填结果
// 归属到父会话的实时流，而不是发一张无归属的帧。
func (t *SpawnTool) await(ids []string, parentSession string) {
	res, err := t.sched.Wait(context.Background(), ids)
	if err != nil {
		t.log.Warn("sub-agent batch wait failed", "count", len(ids), "error", err)
		if res == nil {
			return
		}
	}
	b := t.sched.Bus()
	if b == nil {
		return
	}
	payload := &SubagentResultEvent{SessionID: parentSession, SpawnResult: doneSpawnResult(res)}
	if perr := b.Publish(context.Background(), SubagentResultSubject, payload); perr != nil {
		t.log.Warn("publish sub-agent results", "error", perr)
	}
}

// SpawnResult 是派发工具返回给派发模型的载荷：批次在前台窗口内结算时是它的合并结果，
// 转入后台时是一条派发确认。
type SpawnResult struct {
	// Status 是 "done"（已结算）或 "running"（已转入后台）。
	Status string `json:"status"`
	// Summary 是对该批次的单行聚合。
	Summary string `json:"summary"`
	// Agents 持有每个已派发子代理的一个条目，按派发顺序排列。
	Agents []SpawnAgentResult `json:"agents"`
}

// SpawnAgentResult 是 SpawnResult 中单个子代理的条目。
type SpawnAgentResult struct {
	// ID 是子代理 id（也是 WS subagent_event 帧所携带的 id）。
	ID string `json:"id"`
	// Goal 是交给该子代理的目标。
	Goal string `json:"goal"`
	// Role 是绑定的角色/persona（若有）：指名了派发目标角色时用该角色，否则是旧的
	// 角色标签。
	Role string `json:"role,omitempty"`
	// Mode 是绑定的行为模式（若有）。
	Mode string `json:"mode,omitempty"`
	// Status 是返回时刻的生命周期状态。
	Status string `json:"status"`
	// Summary 是子代理结算时写入的执行摘要。
	Summary string `json:"summary,omitempty"`
	// Text 是子代理的结果文本（上限 maxSpawnTextRunes）。
	Text string `json:"text,omitempty"`
	// Steps 是子代理执行的循环步数。
	Steps int `json:"steps,omitempty"`
}

// SubagentResultEvent 是后台批次结算后发布的回填载荷（design/05 §9；契约见
// Docs/api-client.md §3.10）：SessionID 为父会话，其余字段即派发工具返回给模型的结果。
type SubagentResultEvent struct {
	SessionID string `json:"sessionId,omitempty"`
	*SpawnResult
}

// doneSpawnResult 把已结算的批次映射为工具结果。
func doneSpawnResult(res *MergedResult) *SpawnResult {
	out := &SpawnResult{Status: "done", Agents: make([]SpawnAgentResult, 0)}
	if res == nil {
		return out
	}
	out.Summary = res.Summary
	for _, a := range res.Agents {
		out.Agents = append(out.Agents, SpawnAgentResult{
			ID:      a.ID,
			Goal:    a.Goal,
			Role:    a.RoleID,
			Mode:    a.Mode,
			Status:  string(a.Status),
			Summary: a.Summary,
			Text:    compactText(a.Text, maxSpawnTextRunes),
			Steps:   a.Steps,
		})
	}
	return out
}

// runningSpawnResult 确认一个仍在后台运行的批次。
func runningSpawnResult(ids []string, specs []AgentSpec) *SpawnResult {
	out := &SpawnResult{
		Status:  "running",
		Summary: fmt.Sprintf("已派发 %d 个子代理，正在后台运行；完成后结果回填", len(ids)),
		Agents:  make([]SpawnAgentResult, 0, len(ids)),
	}
	for i, id := range ids {
		spec := specs[i]
		out.Agents = append(out.Agents, SpawnAgentResult{
			ID:     id,
			Goal:   spec.Goal,
			Role:   spec.RoleID,
			Mode:   spec.Mode,
			Status: string(StatusRunning),
		})
	}
	return out
}

// spawnSpecs 把工具参数解码成每个子代理一个 AgentSpec。非空的 tasks 数组是并行批次，
// 优先于单个 goal 简写。dftTimeout 是派发层默认值（WithSpawnTimeout）；非正值则让
// 调度器的按代理时限说了算。
func spawnSpecs(m map[string]any, parentSession string, dftTimeout time.Duration) ([]AgentSpec, error) {
	tasks, err := taskList(m["tasks"])
	if err != nil {
		return nil, err
	}
	if len(tasks) > maxSpawnBatch {
		return nil, fmt.Errorf("too many tasks: %d (max %d)", len(tasks), maxSpawnBatch)
	}
	if len(tasks) == 0 {
		spec, serr := specFromMap(m, parentSession)
		if serr != nil {
			return nil, serr
		}
		// 派发层默认时限只补位：按调用未给出 timeout_seconds 时才生效。
		if spec.Timeout == 0 {
			spec.Timeout = dftTimeout
		}
		return []AgentSpec{spec}, nil
	}
	// 批次级的白名单是未自带名单任务的缺省值（缺省全开：未列出即继承全量）。
	batchWhitelist := stringListArg(m, "tool_whitelist")
	batchTimeout := timeoutArg(m, "timeout_seconds", dftTimeout)
	specs := make([]AgentSpec, 0, len(tasks))
	for i, task := range tasks {
		spec, serr := specFromMap(task, parentSession)
		if serr != nil {
			return nil, fmt.Errorf("tasks[%d]: %w", i, serr)
		}
		if len(spec.ToolWhitelist) == 0 {
			spec.ToolWhitelist = cloneStrings(batchWhitelist)
		}
		if spec.Timeout == 0 {
			spec.Timeout = batchTimeout
		}
		specs = append(specs, spec)
	}
	return specs, nil
}

// specFromMap 从已解码的参数对象（顶层对象或 tasks 中的一个条目）读取一个 AgentSpec。
// goal 为空时包装 ErrMissingGoal。
func specFromMap(m map[string]any, parentSession string) (AgentSpec, error) {
	goal := strings.TrimSpace(stringArg(m, "goal"))
	if goal == "" {
		return AgentSpec{}, ErrMissingGoal
	}
	personaID := strings.TrimSpace(stringArg(m, "persona"))
	roleID := strings.TrimSpace(stringArg(m, "role"))
	// 被指名的 persona 就是子代理实际充当的角色，因此它也填充绑定元数据：让 role 与
	// persona 互相矛盾会导致工具结果、进度事件与 ws subagent_event 帧点名一个子代理
	// 并未以之作答的角色。省略 persona 则逐字节保持既有的纯 role 路径
	// （design/22 §8-2 向后兼容）。
	if personaID != "" {
		roleID = personaID
	}
	return AgentSpec{
		Goal:          goal,
		RoleID:        roleID,
		PersonaID:     personaID,
		Mode:          strings.TrimSpace(stringArg(m, "mode")),
		AttachInfo:    strings.TrimSpace(stringArg(m, "attach_info")),
		ParentSession: parentSession,
		ToolWhitelist: stringListArg(m, "tool_whitelist"),
		Timeout:       timeoutArg(m, "timeout_seconds", 0),
	}, nil
}

// taskList 读取由对象组成的 tasks 数组。缺失或空数组产出 nil；非数组或条目不是对象则
// 被拒绝，使畸形批次绝不静默退化成单次派发。
func taskList(v any) ([]map[string]any, error) {
	if v == nil {
		return nil, nil
	}
	raw, ok := v.([]any)
	if !ok {
		return nil, fmt.Errorf("tasks must be an array")
	}
	out := make([]map[string]any, 0, len(raw))
	for i, item := range raw {
		m, ok := item.(map[string]any)
		if !ok {
			return nil, fmt.Errorf("tasks[%d] must be an object", i)
		}
		out = append(out, m)
	}
	return out, nil
}

// spawnStringProp 构建一个字符串属性 schema。
func spawnStringProp(desc string) map[string]any {
	return map[string]any{"type": "string", "description": desc}
}

// spawnTasksProp 构建并行批次的 tasks 数组属性 schema。
func spawnTasksProp() map[string]any {
	return map[string]any{
		"type":        "array",
		"description": "Sub-agents to dispatch in parallel. When given, it overrides goal/persona/role/mode/attach_info.",
		"maxItems":    maxSpawnBatch,
		"items": map[string]any{
			"type": "object",
			"properties": map[string]any{
				"goal":            spawnStringProp("Objective this sub-agent works toward alone."),
				"persona":         spawnPersonaProp(),
				"role":            spawnStringProp("Role/persona id to bind the sub-agent to (optional). Superseded by persona."),
				"mode":            spawnStringProp("Behavior mode for the role (optional)."),
				"attach_info":     spawnStringProp("Extra context or data for this sub-agent (optional)."),
				"tool_whitelist":  spawnWhitelistProp("Only this sub-agent; empty means every registered tool."),
				"timeout_seconds": spawnTimeoutProp(0),
			},
			"required": []any{"goal"},
		},
	}
}

// spawnWhitelistProp 构建 tool_whitelist 属性 schema（design/05 §7
// 裁决 02.4.1）：缺省全开，显式列出则只放行所列。
func spawnWhitelistProp(desc string) map[string]any {
	return map[string]any{
		"type": "array",
		"items": map[string]any{
			"type": "string",
		},
		"description": "Tool names the sub-agent may invoke. Omitted or empty means every registered tool " +
			"(the default, the main agent's full set); listing names restricts the sub-agent to exactly those " +
			"tools. " + desc,
	}
}

// spawnTimeoutProp 构建 timeout_seconds 属性 schema（design/05 §7
// 裁决 02.4.4 的按调用覆盖）。dft 非正时描述为「沿用派发默认」。
func spawnTimeoutProp(dft time.Duration) map[string]any {
	desc := "Per-call timeout override in seconds for the sub-agent run (optional; omitted keeps the " +
		"configured sub-agent default)."
	if dft > 0 {
		desc = fmt.Sprintf("Per-call timeout override in seconds for the sub-agent run (optional, default %s).", dft)
	}
	return map[string]any{
		"type":        "integer",
		"minimum":     minSpawnTimeoutSeconds,
		"maximum":     maxSpawnTimeoutSeconds,
		"description": desc,
	}
}

// spawnPersonaProp 构建 persona 属性 schema。它是 design/22 §8-2 的派发开关：指名一个
// 角色会让子代理替班该角色（用该角色的指令与记忆作用域），而不是以派发会话的角色工作。
func spawnPersonaProp() map[string]any {
	return spawnStringProp("Persona/role id the sub-agent runs AS: it answers with that role's " +
		"instructions and reads/writes that role's memory scope (optional, default: the current " +
		"role). A role id this runtime does not know is rejected. Supersedes role.")
}

// stringArg 读取一个字符串参数，容忍缺失或非字符串值。
func stringArg(m map[string]any, key string) string {
	s, _ := m[key].(string)
	return s
}

// boolArg 读取一个布尔参数，容忍缺失或非布尔值。
func boolArg(m map[string]any, key string) bool {
	b, _ := m[key].(bool)
	return b
}

// spawnLog 返回标注了来源 "orchestration/spawn" 的 l，并把 nil logger 换成丢弃式
// logger，使该工具总能记录日志。
func spawnLog(l *slog.Logger) *slog.Logger {
	if l == nil {
		l = slog.New(slog.NewTextHandler(io.Discard, nil))
	}
	return logger.WithSource(l, "orchestration/spawn")
}