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"
)

// Constants for the sub-agent dispatch tool (design/05 §1 创建, §7 并行与安全).
const (
	// SpawnToolName is the registered name of the tool the agent calls to
	// dispatch background sub-agents.
	SpawnToolName = "subagent_spawn"

	// DefaultSpawnWait is how long a foreground dispatch waits for its batch
	// before handing the turn back to the model. It mirrors the runtime's
	// established "前台运行，超过 2 秒自动移入后台" threshold
	// (DevDocs/design/04-运行时.md §8.5, internal/runtime/task). The sub-agents keep
	// running either way; only whether the tool result already carries their
	// output changes.
	DefaultSpawnWait = 2 * time.Second

	// maxSpawnBatch caps how many sub-agents one dispatch may create. It bounds the
	// queue a single model call can build (design/05 §7 防失控); the parallel
	// execution limit itself stays with the scheduler's MaxParallel.
	maxSpawnBatch = 8

	// maxSpawnTextRunes caps how many runes of one sub-agent's result text are
	// handed back to the spawning model, so a runaway sub-agent cannot flood the
	// front session's context. The scheduler keeps the full text on its snapshot.
	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 is the runtime entry point that lets the running agent dispatch
// sub-agents (design/05 §1). It is the only production caller of
// (*Scheduler).Spawn: the model invokes it like any other tool, the scheduler
// applies the parallel limit, per-agent timeout and abort machinery, and the tool
// hands back either the merged result or — when the batch is still running after
// the foreground wait — a dispatch acknowledgement.
//
// It is safe for concurrent use and holds no per-call state.
type SpawnTool struct {
	sched   *Scheduler
	wait    time.Duration
	timeout time.Duration
	log     *slog.Logger
}

// SpawnOption customizes a SpawnTool built by NewSpawnTool.
type SpawnOption func(*SpawnTool)

// WithSpawnWait overrides how long a foreground dispatch waits for its batch. A
// non-positive value keeps DefaultSpawnWait.
func WithSpawnWait(d time.Duration) SpawnOption {
	return func(t *SpawnTool) {
		if d > 0 {
			t.wait = d
		}
	}
}

// WithSpawnLogger replaces the tool's logger (tagged source "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 builds the dispatch tool over sched. A nil scheduler is tolerated:
// Execute then fails fast with ErrNoScheduler instead of panicking.
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 returns the tool identifier.
func (t *SpawnTool) Name() string { return SpawnToolName }

// Description explains the tool to the model: what a sub-agent is, that it runs in
// an isolated context in parallel, and that a slow batch is handed to the
// background rather than blocking the conversation (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 returns the input JSON Schema: either a single goal shorthand or a
// tasks array for a parallel batch, plus an optional persona (the role the
// sub-agent runs as, design/22 §8-2), an optional background flag, an optional
// per-dispatch tool whitelist (design/05 §7) and an optional per-call timeout
// override.
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 dispatches the requested sub-agents, then either returns their merged
// results (the batch settled inside the foreground window) or a dispatch
// acknowledgement (it moved to the background, reported later on the bus).
//
// The sub-agents are spawned on a detached context: their lifetime is their own
// (bounded by the scheduler's per-agent timeout and Abort/CancelAll), never the
// dispatching turn's, so a slow batch is not killed the moment the tool returns
// (design/04 §2 后台不阻塞前台). Context values are kept so the sub-agent span —
// and the flow lane it roots — still chains under the step that dispatched it.
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.FlowOwnerFromContext(ctx), t.timeout)
	if err != nil {
		return nil, fmt.Errorf("%s: %w", SpawnToolName, err)
	}
	// Reject an unknown target role before anything starts: a batch must be
	// all-or-nothing, otherwise the sub-agents already dispatched would keep running
	// with their result never reported to the model (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 waits up to t.wait for the dispatched batch to settle. A settled
// batch is merged into the tool result; a batch still running when the window
// closes is announced in the background and acknowledged immediately.
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 blocks until the batch settles and broadcasts the merged outcome as the
// design/05 §2 "结束后写执行总结 → 返回主会话" hand-off. It runs on a detached
// goroutine; a nil bus or an unknown id ends it quietly.
// 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 is the payload the dispatch tool returns to the spawning model: the
// batch's merged outcome when it settled inside the foreground window, or a
// dispatch acknowledgement when it moved to the background.
type SpawnResult struct {
	// Status is "done" (settled) or "running" (dispatched to the background).
	Status string `json:"status"`
	// Summary is a one-line aggregate over the batch.
	Summary string `json:"summary"`
	// Agents holds one entry per dispatched sub-agent, in dispatch order.
	Agents []SpawnAgentResult `json:"agents"`
}

// SpawnAgentResult is one sub-agent's entry in a SpawnResult.
type SpawnAgentResult struct {
	// ID is the sub-agent id (also the id the WS subagent_event frames carry).
	ID string `json:"id"`
	// Goal is the objective the sub-agent was given.
	Goal string `json:"goal"`
	// Role is the bound role/persona, if any: the dispatched target persona when one
	// was named, otherwise the legacy role label.
	Role string `json:"role,omitempty"`
	// Mode is the bound behavior mode, if any.
	Mode string `json:"mode,omitempty"`
	// Status is the lifecycle status at return time.
	Status string `json:"status"`
	// Summary is the execution summary written when the sub-agent settled.
	Summary string `json:"summary,omitempty"`
	// Text is the sub-agent's result text (capped at maxSpawnTextRunes).
	Text string `json:"text,omitempty"`
	// Steps is how many loop steps the sub-agent ran.
	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 maps a settled batch onto the tool result.
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 acknowledges a batch that is still running in the background.
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 decodes the tool arguments into one AgentSpec per sub-agent. A
// non-empty tasks array is a parallel batch and takes precedence over the single
// goal shorthand. dftTimeout is the dispatch-layer default (WithSpawnTimeout);
// a non-positive value leaves the scheduler's per-agent timeout in charge.
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 reads one AgentSpec out of a decoded arguments object (the top-level
// object or one tasks entry). A blank goal wraps 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"))
	// A named persona is the role the sub-agent actually runs as, so it also fills
	// the binding metadata: leaving role and persona disagreeing would make the tool
	// result, the progress events and the ws subagent_event frame name a role the
	// sub-agent is not answering as. Omitting persona keeps the pre-existing
	// role-only path byte for byte (design/22 §8-2 backward compatibility).
	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 reads a tasks array of objects. A missing or empty array yields nil; a
// non-array or an entry that is not an object is rejected so a malformed batch
// never silently degrades into a single spawn.
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 builds a string property schema.
func spawnStringProp(desc string) map[string]any {
	return map[string]any{"type": "string", "description": desc}
}

// spawnTasksProp builds the tasks array property schema for a parallel batch.
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 builds the tool_whitelist property 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 builds the timeout_seconds property 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 builds the persona property schema. It is the design/22 §8-2 派发
// switch: naming a role makes the sub-agent stand in for it (that role's
// instructions and memory scope) instead of working as the spawning session's role.
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 reads a string argument, tolerating a missing or non-string value.
func stringArg(m map[string]any, key string) string {
	s, _ := m[key].(string)
	return s
}

// boolArg reads a boolean argument, tolerating a missing or non-boolean value.
func boolArg(m map[string]any, key string) bool {
	b, _ := m[key].(bool)
	return b
}

// spawnLog returns l tagged with source "orchestration/spawn", replacing a nil
// logger with a discard one so the tool can always log.
func spawnLog(l *slog.Logger) *slog.Logger {
	if l == nil {
		l = slog.New(slog.NewTextHandler(io.Discard, nil))
	}
	return logger.WithSource(l, "orchestration/spawn")
}