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

	// SubagentResultSubject is the bus subject a finished background batch is
	// announced on, carrying the *MergedResult. It is the design/05 §9 结果回填 hook
	// for whichever component owns the session message queue (design/15); the
	// WS-visible per-agent status already travels on the scheduler's own progress
	// events, so this subject has no consumer until that queue lands.
	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
	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) }
}

// 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 background flag.
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."),
			"role":        spawnStringProp("Role/persona id to bind the sub-agent to (optional)."),
			"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(),
			"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))
	if err != nil {
		return nil, fmt.Errorf("%s: %w", SpawnToolName, err)
	}

	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,
		"mode", specs[0].Mode, "session", specs[0].ParentSession)

	if boolArg(m, "background") {
		go t.await(ids)
		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)
		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.
func (t *SpawnTool) await(ids []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
	}
	if perr := b.Publish(context.Background(), SubagentResultSubject, res); 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.
	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"`
}

// 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.
func spawnSpecs(m map[string]any, parentSession string) ([]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
		}
		return []AgentSpec{spec}, nil
	}
	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)
		}
		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
	}
	return AgentSpec{
		Goal:          goal,
		RoleID:        strings.TrimSpace(stringArg(m, "role")),
		Mode:          strings.TrimSpace(stringArg(m, "mode")),
		AttachInfo:    strings.TrimSpace(stringArg(m, "attach_info")),
		ParentSession: parentSession,
	}, 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/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."),
				"role":        spawnStringProp("Role/persona id to bind the sub-agent to (optional)."),
				"mode":        spawnStringProp("Behavior mode for the role (optional)."),
				"attach_info": spawnStringProp("Extra context or data for this sub-agent (optional)."),
			},
			"required": []any{"goal"},
		},
	}
}

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