Cyrene_v2/internal/runtime/loop/attribution_test.go

package loop

import (
	"context"
	"strings"
	"testing"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/agent"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/memory"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
)

// This file covers design/22 §7-5 from the model's side: a conversation whose
// assistant messages were spoken by different roles must reach the provider with
// that attribution visible, without a single character being added to a message
// body and without disturbing the assembled prompt's established order.

// speakerMsg builds one conversation message labelled with the role that spoke
// it, exactly the way the gateway's transcript seeding labels an assistant turn
// (model.Message.SenderName carries the speaking role id).
func speakerMsg(role model.Role, personaID, text string) model.Message {
	m := model.TextMessage(role, text)
	m.SenderName = personaID
	return m
}

// attributedLoop builds a two-role loop: instructions say which role answered, and
// the name resolver mirrors the persona registry's id→display-name mapping.
func attributedLoop(t *testing.T, ad model.Adapter) *Loop {
	t.Helper()
	xilian := personaAgent(t, "你是昔涟。", ad)
	yuki := personaAgent(t, "你是由纪。", ad)
	names := map[string]string{"xilian": "昔涟", "yuki": "由纪"}
	return New(Options{
		Agent:       xilian,
		Agents:      map[string]*agent.Agent{"xilian": xilian, "yuki": yuki},
		PersonaID:   "xilian",
		PersonaName: func(id string) string { return names[id] },
	})
}

// seedMemory writes one persona-scoped memory item for personaID so the round's
// retrieval has something to inject.
func seedMemory(t *testing.T, store *memory.Store, personaID, content string) {
	t.Helper()
	if _, err := store.Add(context.Background(), memory.MemoryItem{
		PersonaID: personaID, Content: content,
		Scope: memory.ScopePersona, Kind: memory.KindRelationship,
	}); err != nil {
		t.Fatalf("seed memory: %v", err)
	}
}

// systemTexts returns the text of every system message of req, in order.
func systemTexts(req model.Request) []string {
	out := make([]string, 0, len(req.Messages))
	for _, m := range req.Messages {
		if m.Role == model.RoleSystem {
			out = append(out, messageText(m))
		}
	}
	return out
}

// indexOfSystemContaining returns the index of the first message whose text
// contains want, or -1.
func indexOfSystemContaining(req model.Request, want string) int {
	for i, m := range req.Messages {
		if strings.Contains(messageText(m), want) {
			return i
		}
	}
	return -1
}

// TestAttribution_SpeakersReachTheModelRequest is the acceptance case of
// design/22 §7-5: the session now speaks as Y while the context still holds
// utterances by X. The request the adapter receives must (a) open with Y's own
// instructions, (b) declare which role is answering, (c) mark where the speaker
// changed so "第 N 条是谁说的" is decidable, and (d) leave every message body
// untouched — the attribution travels only in system messages.
func TestAttribution_SpeakersReachTheModelRequest(t *testing.T) {
	t.Parallel()
	ad := &mockAdapter{chat: []mgResp{{resp: model.Response{Content: "在呢"}}}}
	l := attributedLoop(t, ad)

	msgs := []model.Message{
		speakerMsg(model.RoleAssistant, "xilian", "我记住了"),
		speakerMsg(model.RoleAssistant, "xilian", "还有一件事"),
		model.TextMessage(model.RoleUser, "继续"),
		speakerMsg(model.RoleAssistant, "yuki", "我在呢"),
		model.TextMessage(model.RoleUser, "你好"),
	}
	if _, err := l.Run(WithSessionEnv(context.Background(), roleEnv("yuki")), "sess-y", msgs); err != nil {
		t.Fatalf("Run: %v", err)
	}
	req := lastChatRequest(t, ad)

	// (a) the role that answers now still owns the first system message.
	if got := messageText(req.Messages[0]); got != "你是由纪。" {
		t.Fatalf("first system = %q, want the bound role's instructions", got)
	}

	sys := systemTexts(req)
	// instructions + the role-context segment + one boundary note per speaker
	// stretch (X's two assistant turns are one stretch).
	if len(sys) != 4 {
		t.Fatalf("system messages = %d (%q), want 4 (instructions + role context + two speaker boundaries)", len(sys), sys)
	}
	if sys[0] != "你是由纪。" {
		t.Fatalf("system[0] = %q, want the instructions untouched at the head", sys[0])
	}
	if !strings.HasPrefix(sys[1], roleContextMarker) {
		t.Fatalf("system[1] = %q, want the role context segment", sys[1])
	}
	for _, want := range []string{"「由纪」（yuki）", "不同角色相继发言"} {
		if !strings.Contains(sys[1], want) {
			t.Fatalf("role context %q missing %q", sys[1], want)
		}
	}
	if strings.Contains(sys[1], speakerSwitchMarker) {
		t.Fatalf("role context %q must not be mistaken for a boundary note", sys[1])
	}
	if !strings.HasPrefix(sys[2], speakerSwitchMarker) || !strings.Contains(sys[2], "「昔涟」（xilian）") || !strings.Contains(sys[2], "不是你说的") {
		t.Fatalf("system[2] = %q, want X's boundary note marking its utterances as not yours", sys[2])
	}
	if !strings.HasPrefix(sys[3], speakerSwitchMarker) || !strings.Contains(sys[3], "来自你自己") || !strings.Contains(sys[3], "「由纪」（yuki）") {
		t.Fatalf("system[3] = %q, want Y's own boundary note naming the answering role", sys[3])
	}

	// (c) the boundary sits directly before the first utterance it covers, and the
	// second X utterance inherits the same speaker without a second note.
	firstX := indexOfSystemContaining(req, "我记住了")
	if firstX < 0 || req.Messages[firstX-1].Role != model.RoleSystem {
		t.Fatalf("X's first utterance at %d is not preceded by a boundary note (%+v)", firstX, req.Messages)
	}
	if idx := indexOfSystemContaining(req, "「昔涟」（xilian）"); idx > firstX {
		t.Fatalf("X boundary at %d must precede its utterance at %d", idx, firstX)
	}
	// The answering role's own history is re-opened by its own boundary note.
	self := indexOfSystemContaining(req, "我在呢")
	if self < 0 || !strings.Contains(messageText(req.Messages[self-1]), "来自你自己") {
		t.Fatalf("Y's utterance is not preceded by its self boundary: %+v", req.Messages)
	}

	// (d) no body was rewritten: user text stays verbatim and carries no metadata,
	// and the assistant's own words are unchanged.
	var users []string
	for _, m := range req.Messages {
		if m.Role != model.RoleUser {
			continue
		}
		users = append(users, messageText(m))
	}
	if len(users) != 2 || users[0] != "继续" || users[1] != "你好" {
		t.Fatalf("user bodies = %q, want the verbatim [继续 你好] (no system metadata in the user lane)", users)
	}
	for _, u := range users {
		if strings.Contains(u, speakerSwitchMarker) || strings.Contains(u, roleContextMarker) {
			t.Fatalf("user body %q carries system attribution metadata", u)
		}
	}
	for _, body := range []string{"我记住了", "还有一件事", "我在呢"} {
		if idx := indexOfSystemContaining(req, body); idx >= 0 && strings.Contains(messageText(req.Messages[idx]), "「") {
			t.Fatalf("assistant body %q was rewritten with attribution", messageText(req.Messages[idx]))
		}
	}
}

// TestAssembleContext_AttributionKeepsBlockOrder pins the placement rule: the role
// declaration is added AFTER the established leading blocks (time, working memory,
// memory) and immediately before the conversation, so the memory/time block order
// and the instructions-first rule are untouched.
func TestAssembleContext_AttributionKeepsBlockOrder(t *testing.T) {
	t.Parallel()
	store := newTestMemoryStore(t)
	seedMemory(t, store, "yuki", "她喜欢冰美式")
	ad := &mockAdapter{chat: []mgResp{{resp: model.Response{Content: "嗯"}}}}
	xilian := personaAgent(t, "你是昔涟。", ad)
	yuki := personaAgent(t, "你是由纪。", ad)
	l := New(Options{
		Agent:       xilian,
		Agents:      map[string]*agent.Agent{"xilian": xilian, "yuki": yuki},
		PersonaID:   "xilian",
		PersonaName: func(id string) string { return map[string]string{"xilian": "昔涟", "yuki": "由纪"}[id] },
		Retriever:   memory.NewRetriever(store),
		Timezone:    "Asia/Shanghai",
	})

	msgs := []model.Message{
		speakerMsg(model.RoleAssistant, "xilian", "我记住了"),
		model.TextMessage(model.RoleUser, "继续"),
	}
	if _, err := l.Run(WithSessionEnv(context.Background(), roleEnv("yuki")), "sess-order", msgs); err != nil {
		t.Fatalf("Run: %v", err)
	}
	req := lastChatRequest(t, ad)

	timeIdx := indexOfSystemContaining(req, "当前时间")
	memIdx := indexOfSystemContaining(req, "冰美式")
	declIdx := indexOfSystemContaining(req, roleContextMarker)
	noteIdx := indexOfSystemContaining(req, speakerSwitchMarker)
	lastUser := -1
	for i, m := range req.Messages {
		if m.Role == model.RoleUser {
			lastUser = i
		}
	}
	if timeIdx < 0 || memIdx < 0 || declIdx < 0 || noteIdx < 0 || lastUser < 0 {
		t.Fatalf("missing block: time=%d mem=%d decl=%d note=%d user=%d (%+v)", timeIdx, memIdx, declIdx, noteIdx, lastUser, req.Messages)
	}
	if !(memIdx < declIdx && declIdx < noteIdx && noteIdx < lastUser) {
		t.Fatalf("block order = time:%d mem:%d decl:%d note:%d user:%d, want instructions < time < memory < declaration < notes < user", timeIdx, memIdx, declIdx, noteIdx, lastUser)
	}
	if got := messageText(req.Messages[0]); got != "你是由纪。" {
		t.Fatalf("first system = %q, want the instructions untouched at the head", got)
	}
}

// TestAttribution_UnlabeledConversationIsUntouched locks the backward-compatible
// half: a caller that labels nothing (every single-role path, IM, backthink, the
// autonomous passes) gets exactly the prompt it had before, with no role segment
// and no boundary notes.
func TestAttribution_UnlabeledConversationIsUntouched(t *testing.T) {
	t.Parallel()
	ad := &mockAdapter{chat: []mgResp{{resp: model.Response{Content: "done"}}}}
	l := attributedLoop(t, ad)

	msgs := []model.Message{
		model.TextMessage(model.RoleAssistant, "我记住了"),
		model.TextMessage(model.RoleUser, "继续"),
	}
	if _, err := l.Run(WithSessionEnv(context.Background(), roleEnv("yuki")), "sess-plain", msgs); err != nil {
		t.Fatalf("Run: %v", err)
	}
	req := lastChatRequest(t, ad)
	if len(req.Messages) != len(msgs)+1 {
		t.Fatalf("messages = %d, want %d (instructions + conversation, nothing added)", len(req.Messages), len(msgs)+1)
	}
	for _, s := range systemTexts(req) {
		if strings.Contains(s, roleContextMarker) || strings.Contains(s, speakerSwitchMarker) {
			t.Fatalf("unlabeled conversation got attribution %q", s)
		}
	}
}