Cyrene_v2/internal/runtime/orchestration/spawnpersona_test.go

package orchestration

import (
	"context"
	"io"
	"log/slog"
	"strings"
	"sync"
	"testing"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/restrict"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/db"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/agent"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/loop"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/memory"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
)

// This file locks design/22 §8-2 at the dispatch layer: `subagent_spawn` may name a
// target persona, and the sub-agent then runs AS that role — that role's instructions
// become the leading system message and that role's memory scope is what it reads and
// writes ("她以该角色身份临时替班"). Everything is exercised in process over a REAL
// *loop.Loop with a per-role agent pool and a REAL memory store, so the assertions are
// about what the sub-agent's run actually saw, not about wiring.
//
// The other half of the file pins the backward-compatibility contract: omitting the
// persona must reproduce the pre-change behavior exactly (the sub-agent inherits the
// spawning turn's role), an unknown persona must be rejected instead of silently
// degrading, and naming a persona must never widen the identity/permission lane.

const (
	// parentRole is the role the dispatching session is bound to — deliberately NOT
	// the role a persona dispatch asks for, so a leak shows up as the wrong prompt.
	parentRole = "xilian"
	// targetRole is the role a persona dispatch asks the sub-agent to stand in for.
	targetRole = "yuki"
)

// parentPrompt and targetPrompt are the two roles' instructions; a test tells which
// pool entry the loop resolved from the leading system message alone.
const (
	parentPrompt = "你是昔涟。"
	targetPrompt = "你是由纪。"
)

// roleAdapter records every request the loop builds and answers with a fixed final
// text, so a test can read the exact prompt and memory block a sub-agent's run got.
type roleAdapter struct {
	mu   sync.Mutex
	reqs []model.Request
}

func (a *roleAdapter) Name() string { return "mock" }

func (a *roleAdapter) Chat(_ context.Context, req model.Request, _ model.Options) (model.Response, error) {
	a.mu.Lock()
	a.reqs = append(a.reqs, req)
	a.mu.Unlock()
	return model.Response{Model: "mock-1", Content: "好的", StopReason: "end_turn"}, nil
}

func (a *roleAdapter) ChatStream(ctx context.Context, req model.Request, _ func(model.StreamEvent)) (model.Response, error) {
	return a.Chat(ctx, req, model.Options{})
}

// request returns the n-th recorded request, failing the test when the run produced
// fewer. Requests are recorded in call order, i.e. one per sub-agent run.
func (a *roleAdapter) request(t *testing.T, n int) model.Request {
	t.Helper()
	a.mu.Lock()
	defer a.mu.Unlock()
	if len(a.reqs) <= n {
		t.Fatalf("adapter recorded %d requests, want at least %d", len(a.reqs), n+1)
	}
	return a.reqs[n]
}

// leadingSystem returns the first message of the request, which is where the agent's
// instructions land (agent/request.go 首条 system 消息).
func (a *roleAdapter) leadingSystem(t *testing.T, n int) string {
	t.Helper()
	msgs := a.request(t, n).Messages
	if len(msgs) == 0 {
		t.Fatalf("request %d carried no messages", n)
	}
	return messageText(msgs[0])
}

// carries reports whether any message of the n-th request contains want.
func (a *roleAdapter) carries(t *testing.T, n int, want string) bool {
	t.Helper()
	for _, m := range a.request(t, n).Messages {
		if strings.Contains(messageText(m), want) {
			return true
		}
	}
	return false
}

// personaRig is the in-process e2e rig: a real loop with a two-role agent pool over
// a capturing adapter, the real scheduler on top (the production shape), the dispatch
// tool above that, and a real memory store so role scoping is asserted on persisted
// items. The turn hook records what role the loop resolved for each settled run.
type personaRig struct {
	tool     *SpawnTool
	sched    *Scheduler
	ad       *roleAdapter
	memStore *memory.Store

	mu    sync.Mutex
	turns []loop.TurnEnd
}

// newPersonaRig builds the rig. personaExists is the role resolver under test: nil
// mirrors a composition root that wired no pool.
func newPersonaRig(t *testing.T, personaExists func(string) bool) *personaRig {
	t.Helper()
	rig := &personaRig{ad: &roleAdapter{}, memStore: newPersonaMemoryStore(t)}
	pool := make(map[string]*agent.Agent, 2)
	for id, instructions := range map[string]string{parentRole: parentPrompt, targetRole: targetPrompt} {
		ag, err := agent.New(instructions, nil, nil, agent.WithModel(rig.ad, nil))
		if err != nil {
			t.Fatalf("agent.New(%s): %v", id, err)
		}
		pool[id] = ag
	}

	lp := loop.New(loop.Options{
		Agent:     pool[parentRole],
		Agents:    pool,
		PersonaID: parentRole,
		Memory:    rig.memStore,
		Retriever: memory.NewRetriever(rig.memStore),
		// Built without a persona, exactly as cmd/cyrene builds it: stamping the
		// turn's role onto each item is the loop's job.
		Extractor: memory.NewManualExtractor("", "owner"),
		OnTurnEnd: func(end loop.TurnEnd) string {
			rig.mu.Lock()
			rig.turns = append(rig.turns, end)
			rig.mu.Unlock()
			return ""
		},
	})
	rig.sched = newScheduler(t, Options{MaxParallel: 2, Loop: lp, PersonaExists: personaExists})
	rig.tool = NewSpawnTool(rig.sched, WithSpawnWait(5*time.Second))
	return rig
}

// turn returns the n-th settled turn the loop reported, failing the test when the
// rig observed fewer.
func (r *personaRig) turn(t *testing.T, n int) loop.TurnEnd {
	t.Helper()
	r.mu.Lock()
	defer r.mu.Unlock()
	if len(r.turns) <= n {
		t.Fatalf("loop reported %d settled turns, want at least %d", len(r.turns), n+1)
	}
	return r.turns[n]
}

// parentTurnCtx is the context the dispatching turn runs under: the parent session's
// own role on the owner-full identity lane, as the gateway resolves it. A persona
// dispatch must be able to change the first without touching the second.
func parentTurnCtx() context.Context {
	env := loop.OwnerEnv()
	env.PersonaID = parentRole
	return loop.WithSessionEnv(context.Background(), env)
}

// newPersonaMemoryStore opens an in-memory store for the rig.
func newPersonaMemoryStore(t *testing.T) *memory.Store {
	t.Helper()
	ctx := context.Background()
	d, err := db.Open(ctx, ":memory:", slog.New(slog.NewTextHandler(io.Discard, nil)))
	if err != nil {
		t.Fatalf("db.Open: %v", err)
	}
	t.Cleanup(func() { _ = d.Close() })
	store, err := memory.NewStore(ctx, d, slog.New(slog.NewTextHandler(io.Discard, nil)))
	if err != nil {
		t.Fatalf("memory.NewStore: %v", err)
	}
	return store
}

// seedRoleMemory files one persona-scope memory under roleID.
func seedRoleMemory(t *testing.T, store *memory.Store, roleID, content string) {
	t.Helper()
	_, err := store.Add(context.Background(), memory.MemoryItem{
		PersonaID: roleID,
		RoleID:    roleID,
		Scope:     memory.ScopePersona,
		Kind:      memory.KindRelationship,
		Content:   content,
	})
	if err != nil {
		t.Fatalf("seed %s memory: %v", roleID, err)
	}
}

// storedPersonaOf returns the persona a memory with the given content was filed
// under, or "" when it is absent.
func storedPersonaOf(t *testing.T, store *memory.Store, content string) string {
	t.Helper()
	items, err := store.List(context.Background())
	if err != nil {
		t.Fatalf("memory list: %v", err)
	}
	for _, it := range items {
		if strings.Contains(it.Content, content) {
			return it.PersonaID
		}
	}
	return ""
}

// TestSpawnToolRunsAsTargetPersona is the design/22 §8-2 acceptance, end to end in
// process: a dispatch naming a persona runs with THAT role's instructions and THAT
// role's memory, while the dispatching session's own role stays untouched. The same
// rig then proves the omitted-persona path is unchanged, which is the pre-change
// behavior: the sub-agent inherits the spawning turn's role.
func TestSpawnToolRunsAsTargetPersona(t *testing.T) {
	t.Parallel()
	rig := newPersonaRig(t, func(id string) bool { return id == parentRole || id == targetRole })
	// One memory per role, so the injected block names which scope the run read from.
	seedRoleMemory(t, rig.memStore, targetRole, "由纪的暗号是香菜")
	seedRoleMemory(t, rig.memStore, parentRole, "昔涟的暗号是乌冬")

	raw, err := rig.tool.Execute(parentTurnCtx(), map[string]any{
		"goal": "替我确认一下口味", "persona": targetRole,
	})
	if err != nil {
		t.Fatalf("Execute: %v", err)
	}
	got := decodeSpawnResult(t, raw)
	if len(got.Agents) != 1 || got.Agents[0].Status != string(StatusDone) {
		t.Fatalf("agents = %+v, want one settled sub-agent", got.Agents)
	}
	// The reported role is the target persona, so the tool result, the progress events
	// and the ws subagent_event frame all name the role it really answered as.
	if got.Agents[0].Role != targetRole {
		t.Fatalf("reported role = %q, want the target persona %q", got.Agents[0].Role, targetRole)
	}

	// 1) The prompt the model actually received is the TARGET role's instructions —
	// not the dispatching role's, and not a shared default.
	if prompt := rig.ad.leadingSystem(t, 0); prompt != targetPrompt {
		t.Fatalf("sub-agent system prompt = %q, want the target persona's %q", prompt, targetPrompt)
	}
	// 2) The role the loop resolved for the run (which scopes memory, emotion and the
	// thinking knobs) is the target persona, although the dispatching env said xilian.
	if turn := rig.turn(t, 0); turn.PersonaID != targetRole {
		t.Fatalf("sub-agent turn persona = %q, want %q", turn.PersonaID, targetRole)
	}
	// 3) Memory reads follow the target role: only yuki's memory was injected.
	if !rig.ad.carries(t, 0, "由纪的暗号是香菜") {
		t.Fatalf("sub-agent did not recall the target persona's memory: %+v", rig.ad.request(t, 0).Messages)
	}
	if rig.ad.carries(t, 0, "昔涟的暗号是乌冬") {
		t.Fatalf("sub-agent read the dispatching persona's memory: %+v", rig.ad.request(t, 0).Messages)
	}
	// 4) The identity lane travelled unchanged: naming a persona is never a permission
	// (project red line: 身份识别靠元数据 + 策略引擎).
	env := rig.turn(t, 0).Env
	if env.ActorEnv != restrict.ActorOwner || !env.MemWritable {
		t.Fatalf("sub-agent env = %+v, want the dispatching owner lane unchanged", env)
	}

	// Backward compatibility: a dispatch without persona still runs as the spawning
	// role (xilian) and reads xilian's memory — byte for byte the pre-change behavior.
	raw, err = rig.tool.Execute(parentTurnCtx(), map[string]any{
		"goal": "老路子", "role": "legacy-label",
	})
	if err != nil {
		t.Fatalf("Execute without persona: %v", err)
	}
	got = decodeSpawnResult(t, raw)
	if prompt := rig.ad.leadingSystem(t, 1); prompt != parentPrompt {
		t.Fatalf("persona-less sub-agent prompt = %q, want the spawning role's %q", prompt, parentPrompt)
	}
	if turn := rig.turn(t, 1); turn.PersonaID != parentRole {
		t.Fatalf("persona-less sub-agent persona = %q, want the inherited %q", turn.PersonaID, parentRole)
	}
	if rig.ad.carries(t, 1, "由纪的暗号是香菜") {
		t.Fatalf("persona-less sub-agent read another role's memory: %+v", rig.ad.request(t, 1).Messages)
	}
	if !rig.ad.carries(t, 1, "昔涟的暗号是乌冬") {
		t.Fatalf("persona-less sub-agent lost the inherited role's memory: %+v", rig.ad.request(t, 1).Messages)
	}
	// The legacy role label still round-trips untouched when no persona is named.
	if got.Agents[0].Role != "legacy-label" {
		t.Fatalf("legacy role = %q, want it untouched when persona is omitted", got.Agents[0].Role)
	}
}

// TestSpawnedSubAgentMemoryWriteScopeIsTargetPersona pins the write half of design/22
// §8-2: a memory the sub-agent's turn extracts is filed under the target persona, not
// under the role of the session that dispatched it. It also pins that a persona never
// widens the identity lane: in a restricted (read-only) env the same dispatch still
// writes nothing, even though it names a role.
func TestSpawnedSubAgentMemoryWriteScopeIsTargetPersona(t *testing.T) {
	t.Parallel()
	rig := newPersonaRig(t, func(string) bool { return true })

	sub, err := rig.sched.Spawn(parentTurnCtx(), AgentSpec{
		Goal:      "记住主人的口味",
		RoleID:    targetRole,
		PersonaID: targetRole,
		Context:   []model.Message{model.TextMessage(model.RoleUser, "记住 我讨厌香菜")},
	})
	if err != nil {
		t.Fatalf("Spawn: %v", err)
	}
	if _, err := rig.sched.Wait(context.Background(), []string{sub.ID}); err != nil {
		t.Fatalf("Wait: %v", err)
	}
	if got := storedPersonaOf(t, rig.memStore, "我讨厌香菜"); got != targetRole {
		t.Fatalf("extracted memory filed under %q, want the target persona %q", got, targetRole)
	}

	// Restricted env: the dispatching turn was a group chat (read-only). Naming a
	// persona must not create a write path where the lane forbids one.
	restricted := loop.WithSessionEnv(context.Background(), loop.SessionEnv{
		ActorEnv:    restrict.ActorGroupMember,
		MemWritable: false,
		PersonaID:   parentRole,
	})
	sub, err = rig.sched.Spawn(restricted, AgentSpec{
		Goal:      "群聊里的她被派去记住东西",
		PersonaID: targetRole,
		Context:   []model.Message{model.TextMessage(model.RoleUser, "记住 我喜欢吃辣")},
	})
	if err != nil {
		t.Fatalf("Spawn (restricted): %v", err)
	}
	if _, err := rig.sched.Wait(context.Background(), []string{sub.ID}); err != nil {
		t.Fatalf("Wait (restricted): %v", err)
	}
	if got := storedPersonaOf(t, rig.memStore, "我喜欢吃辣"); got != "" {
		t.Fatalf("a restricted sub-agent wrote memory under %q, want no write at all", got)
	}
	env := rig.turn(t, 1).Env
	if env.MemWritable || env.ActorEnv != restrict.ActorGroupMember {
		t.Fatalf("restricted sub-agent env = %+v, want the read-only group lane preserved", env)
	}
}