HomeAgent/internal/plugin/proc/shm_test.go

package proc

import (
	"fmt"
	"reflect"
	"strings"
	"sync"
	"testing"

	pubsdk "github.com/JianFeeeee/hologramagentsdk/sdk"
)

// 本文件验证共享内存 stage 并发的正确性——**整个迁移最关键的一环**（§4.4 风险 3.4）。
//
// 对照基线（今日 C ABI 副本模型）：
//   - 内置插件（共享 *StageContext + RWMutex）：0% 丢失
//   - 外部插件（快照-副本-写回）：35.8~36.8% 丢失（实验 12），现网量级百分之几脏数据
//
// 目标：共享内存 + 锁仲裁下，跨进程并发改写收敛到 0% 丢失。

// newTestSegment 造一块内存段模拟 mmap 区域（单测无需真 mmap：
// 编解码与 arena 逻辑与底层是 mmap 还是普通内存无关）。
func newTestSegment(t *testing.T, size int) *Segment {
	t.Helper()
	buf := make([]byte, size)
	seg, err := NewSegment(buf)
	if err != nil {
		t.Fatalf("NewSegment: %v", err)
	}
	return seg
}

func TestSegment_AttachValidatesMagicAndVersion(t *testing.T) {
	buf := make([]byte, 8192)
	if _, err := NewSegment(buf); err != nil {
		t.Fatalf("NewSegment: %v", err)
	}
	if _, err := AttachSegment(buf); err != nil {
		t.Fatalf("AttachSegment 应成功: %v", err)
	}

	// 魔数损坏
	bad := make([]byte, len(buf))
	copy(bad, buf)
	bad[0] ^= 0xFF
	if _, err := AttachSegment(bad); err == nil {
		t.Error("魔数不匹配应报错（避免版本不一致时静默错读）")
	}

	// 版本不匹配
	badVer := make([]byte, len(buf))
	copy(badVer, buf)
	badVer[4] = 99
	if _, err := AttachSegment(badVer); err == nil {
		t.Error("版本不匹配应报错")
	}
}

// 全部 16 个字段可跨进程往返——今日经 C ABI 只有 10 个字段可见（§8.3）。
func TestSegment_RoundTripAllFields(t *testing.T) {
	seg := newTestSegment(t, 16384)

	resp := "短路响应"
	src := &pubsdk.StageContext{
		RawMessage:       "原始输入",
		UserID:           "u1",
		GroupID:          "g1",
		LLMText:          "模型输出",
		ReasoningContent: "思考过程", // C ABI 下外部插件看不到
		FinalText:        "最终文本",
		Response:         &resp,
		Phase:            pubsdk.StageAfterToolcall,
		NoMemory:         true,
		ContextMsgs:      []map[string]interface{}{{"role": "user", "content": "hi"}}, // C ABI 看不到
		ToolCalls:        []pubsdk.ToolCall{{ID: "t1", Name: "weather_query", Plugin: "weather"}},
		ToolResults:      []pubsdk.ToolResult{{CallID: "t1", Name: "weather_query", Success: true, Result: "晴"}},
		Memory:           []pubsdk.MemItem{{Role: "user", Content: "记忆", Score: 0.9}}, // C ABI 看不到
		TokenUsage:       map[string]int{"prompt": 100, "completion": 50},             // C ABI 看不到
		Errors:           []string{"err1"},                                            // C ABI 看不到
		Extra: map[string]interface{}{
			ExtraKeyMediaType:     "image",
			ExtraKeyInputSource:   "qq",
			ExtraKeyOutputChannel: "qq",
		},
	}

	if err := seg.WriteAll(src); err != nil {
		t.Fatalf("WriteAll: %v", err)
	}

	var dst pubsdk.StageContext
	if err := seg.ReadInto(&dst); err != nil {
		t.Fatalf("ReadInto: %v", err)
	}

	if dst.RawMessage != src.RawMessage || dst.UserID != src.UserID || dst.GroupID != src.GroupID {
		t.Errorf("标量字段不一致: raw=%q uid=%q gid=%q", dst.RawMessage, dst.UserID, dst.GroupID)
	}
	if dst.ReasoningContent != "思考过程" {
		t.Errorf("ReasoningContent 应可见（C ABI 下不可见）: %q", dst.ReasoningContent)
	}
	if len(dst.ContextMsgs) != 1 {
		t.Errorf("ContextMsgs 应可见: %v", dst.ContextMsgs)
	}
	if len(dst.Memory) != 1 || dst.Memory[0].Score != 0.9 {
		t.Errorf("Memory 应可见: %v", dst.Memory)
	}
	if dst.TokenUsage["prompt"] != 100 {
		t.Errorf("TokenUsage 应可见: %v", dst.TokenUsage)
	}
	if len(dst.Errors) != 1 {
		t.Errorf("Errors 应可见: %v", dst.Errors)
	}
	if dst.Response == nil || *dst.Response != resp {
		t.Errorf("Response 应往返: %v", dst.Response)
	}
	if !dst.NoMemory {
		t.Error("NoMemory 标志应往返")
	}
	if len(dst.ToolResults) != 1 || dst.ToolResults[0].Result != "晴" {
		t.Errorf("ToolResults 应往返: %v", dst.ToolResults)
	}
	if dst.Extra[ExtraKeyMediaType] != "image" {
		t.Errorf("Extra 提升字段应往返: %v", dst.Extra)
	}
}

// TestSegment_RoundTripEveryField 是**反射驱动的全字段往返**判据。
//
// ## 它取代了什么
//
// 上面那条 `TestSegment_RoundTripAllFields` 是逐字段手写断言，历史价值很大
// （C ABI 时代它能抓"哪些字段外部插件看不到"），但有两个结构性缺陷：
//
//  1. **漏字段无声**：它只断言了 Response/NoMemory/Memory/TokenUsage/Errors/
//     ToolResults/Extra 的 media_type。实测漏了 Phase、LLMText、FinalText、
//     ToolCalls、ExtraKeyMediaBlocks、ExtraBag —— 而那些字段确实在往返，
//     只是没人验。**手写断言永远不会因为"新增字段"而变红。**
//  2. 它验证的是"我记得的那几个"，不是"全部"。
//
// 本判据改用反射：遍历 `pubsdk.StageContext` 的每个导出字段，
// 构造非零值 → WriteAll → ReadInto → 逐字段比。**新增 SDK 字段而不改
// shmcodec 时，这里会红并指名是哪个字段。**
//
// ## 为什么不用 == 整体比较
//
// StageContext 内含 `sync.RWMutex`（不可比较）与 map/slice（不可比较），
// 只能逐字段用 reflect.DeepEqual。Extra 需要单独归一（具名键提升 + bag）。
func TestSegment_RoundTripEveryField(t *testing.T) {
	seg := newTestSegment(t, 65536)
	src := fullStageContext(t)

	if err := seg.WriteAll(src); err != nil {
		t.Fatalf("WriteAll: %v", err)
	}
	var dst pubsdk.StageContext
	if err := seg.ReadInto(&dst); err != nil {
		t.Fatalf("ReadInto: %v", err)
	}

	sv := reflect.ValueOf(src).Elem()
	dv := reflect.ValueOf(&dst).Elem()
	var lost []string
	for i := 0; i < sv.NumField(); i++ {
		f := sv.Type().Field(i)
		if !f.IsExported() {
			continue
		}
		if name := f.Name; name == "Extra" {
			continue // 单独验（见下）
		}
		if !reflect.DeepEqual(sv.Field(i).Interface(), dv.Field(i).Interface()) {
			lost = append(lost, fmt.Sprintf("%s: want=%v got=%v",
				f.Name, sv.Field(i).Interface(), dv.Field(i).Interface()))
		}
	}
	if len(lost) > 0 {
		t.Fatalf("StageContext 有字段跨进程往返后**丢值或变形**：\n  %s\n"+
			"（修法：在 shmcodec.go 的 captureLocal/writeLocal/readLocal 三处接线）",
			strings.Join(lost, "\n  "))
	}

	// Extra 的四类具名键 + bag 必须都保住。
	for _, k := range []string{
		ExtraKeyMediaBlocks, ExtraKeyMediaType, ExtraKeyInputSource, ExtraKeyOutputChannel,
	} {
		if !reflect.DeepEqual(src.Extra[k], dst.Extra[k]) {
			t.Errorf("Extra[%s] 未往返：want=%v got=%v", k, src.Extra[k], dst.Extra[k])
		}
	}
	if !reflect.DeepEqual(src.Extra["qq_auth"], dst.Extra["qq_auth"]) {
		t.Errorf("Extra 的自建键（bag）未往返：want=%v got=%v",
			src.Extra["qq_auth"], dst.Extra["qq_auth"])
	}
}

// fullStageContext 构造一个**所有导出字段都非零**的 StageContext。
//
// ★ 它自身带防腐蚀：反射断言每个导出字段都非零。新增 SDK 字段而没在这里
// 赋值时，本函数会直接失败 —— 否则那条新字段永远不会被上面那条判据覆盖，
// 判据看起来全绿却罩不住新字段（"假绿"）。
func fullStageContext(t *testing.T) *pubsdk.StageContext {
	t.Helper()
	resp := "resp"
	sc := &pubsdk.StageContext{
		RawMessage:       "raw-输入",
		UserID:           "u1",
		GroupID:          "g1",
		LLMText:          "llm-输出",
		ReasoningContent: "think",
		FinalText:        "final",
		Response:         &resp,
		Phase:            pubsdk.StageAfterToolcall,
		NoMemory:         true,
		ContextMsgs:      []map[string]interface{}{{"role": "user", "content": "hi"}},
		ToolCalls:        []pubsdk.ToolCall{{ID: "c1", Name: "t1", Plugin: "p"}},
		ToolResults:      []pubsdk.ToolResult{{CallID: "c1", Name: "t1", Success: true, Result: "o"}},
		Memory:           []pubsdk.MemItem{{Role: "user", Content: "m", Score: 0.5}},
		TokenUsage:       map[string]int{"prompt": 1},
		Errors:           []string{"e1"},
		Extra: map[string]interface{}{
			ExtraKeyMediaBlocks:   []interface{}{map[string]interface{}{"type": "text"}},
			ExtraKeyMediaType:     "image",
			ExtraKeyInputSource:   "qq",
			ExtraKeyOutputChannel: "qq",
			"qq_auth":             "self-key",
		},
	}

	// 防腐蚀：每个导出字段都必须非零。
	v := reflect.ValueOf(sc).Elem()
	for i := 0; i < v.NumField(); i++ {
		f := v.Type().Field(i)
		if !f.IsExported() {
			continue
		}
		if v.Field(i).IsZero() {
			t.Fatalf("StageContext.%s 在 fullStageContext 里没赋值（仍是零值）—— "+
				"全字段往返判据覆盖不到它。请给它一个非零值。", f.Name)
		}
	}
	return sc
}

// nil Response 与空字符串 Response 必须可区分（短路语义依赖此）。
func TestSegment_ResponseNilVsEmpty(t *testing.T) {
	seg := newTestSegment(t, 8192)

	if err := seg.WriteAll(&pubsdk.StageContext{RawMessage: "x"}); err != nil {
		t.Fatalf("WriteAll: %v", err)
	}
	var d1 pubsdk.StageContext
	if err := seg.ReadInto(&d1); err != nil {
		t.Fatalf("ReadInto: %v", err)
	}
	if d1.Response != nil {
		t.Errorf("未设置的 Response 应为 nil，实际 %q", *d1.Response)
	}

	empty := ""
	seg2 := newTestSegment(t, 8192)
	if err := seg2.WriteAll(&pubsdk.StageContext{Response: &empty}); err != nil {
		t.Fatalf("WriteAll: %v", err)
	}
	var d2 pubsdk.StageContext
	if err := seg2.ReadInto(&d2); err != nil {
		t.Fatalf("ReadInto: %v", err)
	}
	if d2.Response == nil {
		t.Error("显式设为空串的 Response 不应读成 nil（短路语义会丢）")
	} else if *d2.Response != "" {
		t.Errorf("Response 应为空串，实际 %q", *d2.Response)
	}
}

// 只读插件的 WriteDirty 必须零写入——**这是消除 lost update 的核心断言**。
func TestSegment_WriteDirty_ReadOnlyPluginWritesNothing(t *testing.T) {
	seg := newTestSegment(t, 16384)
	base := &pubsdk.StageContext{
		RawMessage:  "查天气",
		ToolResults: []pubsdk.ToolResult{{CallID: "c1", Name: "weather_query", Result: "已清洗"}},
	}
	if err := seg.WriteAll(base); err != nil {
		t.Fatalf("WriteAll: %v", err)
	}

	// 插件侧：读入 → 只读 → 写回
	var local pubsdk.StageContext
	if err := seg.ReadInto(&local); err != nil {
		t.Fatalf("ReadInto: %v", err)
	}
	snap := TakeSnapshot(&local)
	_ = local.ToolResults[0].Result // 只读，不改

	n, err := seg.WriteDirty(&local, snap)
	if err != nil {
		t.Fatalf("WriteDirty: %v", err)
	}
	if n != 0 {
		t.Fatalf("只读插件应零写回，实际写回 %d 个字段（会覆盖他人改写）", n)
	}
}

// 原地改切片元素必须被识别为脏 —— C ABI 侧修 11.3 时踩过的坑。
func TestSegment_WriteDirty_InPlaceSliceMutationDetected(t *testing.T) {
	seg := newTestSegment(t, 16384)
	if err := seg.WriteAll(&pubsdk.StageContext{
		ToolResults: []pubsdk.ToolResult{{CallID: "c1", Result: "带\x1b[31mANSI\x1b[0m"}},
	}); err != nil {
		t.Fatalf("WriteAll: %v", err)
	}

	var local pubsdk.StageContext
	if err := seg.ReadInto(&local); err != nil {
		t.Fatalf("ReadInto: %v", err)
	}
	snap := TakeSnapshot(&local)
	local.ToolResults[0].Result = "带ANSI" // 原地改元素（sanitizer 的实际行为）

	n, err := seg.WriteDirty(&local, snap)
	if err != nil {
		t.Fatalf("WriteDirty: %v", err)
	}
	if n != 1 {
		t.Fatalf("原地改切片元素应被识别为 1 个脏字段，实际 %d", n)
	}

	var after pubsdk.StageContext
	if err := seg.ReadInto(&after); err != nil {
		t.Fatalf("ReadInto: %v", err)
	}
	if after.ToolResults[0].Result != "带ANSI" {
		t.Errorf("清洗结果未写回: %v", after.ToolResults[0].Result)
	}
}

// 复刻现网场景（实验 13）：sanitizer 改写 + weather 只读并发，清洗结果不得被覆盖。
// 这是 C ABI 副本模型下量级百分之几脏数据的直接来源。
func TestSegment_ProductionScenario_SanitizerNotOverwrittenByWeather(t *testing.T) {
	seg := newTestSegment(t, 16384)
	dirty := "天气：晴 \x1b[31m28°C\x1b[0m"
	clean := "天气：晴 28°C"

	if err := seg.WriteAll(&pubsdk.StageContext{
		RawMessage:  "查天气",
		Phase:       pubsdk.StageAfterToolcall,
		ToolResults: []pubsdk.ToolResult{{CallID: "c1", Name: "weather_query", Result: dirty}},
	}); err != nil {
		t.Fatalf("WriteAll: %v", err)
	}

	lock := newStageLock()

	// sanitizer：拿锁 → 读 → 清洗 → 写脏字段 → 放锁
	runSanitizer := func() error {
		if err := lock.Acquire("sanitizer"); err != nil {
			return err
		}
		defer lock.Release("sanitizer")
		var local pubsdk.StageContext
		if err := seg.ReadInto(&local); err != nil {
			return err
		}
		snap := TakeSnapshot(&local)
		if len(local.ToolResults) > 0 {
			s, _ := local.ToolResults[0].Result.(string)
			local.ToolResults[0].Result = strings.NewReplacer("\x1b[31m", "", "\x1b[0m", "").Replace(s)
		}
		_, err := seg.WriteDirty(&local, snap)
		return err
	}

	// weather：拿锁 → 读 → 只读 → 零写回 → 放锁
	runWeather := func() error {
		if err := lock.Acquire("weather"); err != nil {
			return err
		}
		defer lock.Release("weather")
		var local pubsdk.StageContext
		if err := seg.ReadInto(&local); err != nil {
			return err
		}
		snap := TakeSnapshot(&local)
		if len(local.ToolResults) > 0 {
			_ = local.ToolResults[0].Result // 只读
		}
		n, err := seg.WriteDirty(&local, snap)
		if err != nil {
			return err
		}
		if n != 0 {
			return fmt.Errorf("weather 只读却写回 %d 个字段", n)
		}
		return nil
	}

	// 并发扇出（保留原始设计），weather 后完成是最坏情形
	var wg sync.WaitGroup
	errs := make(chan error, 2)
	wg.Add(2)
	go func() { defer wg.Done(); errs <- runSanitizer() }()
	go func() { defer wg.Done(); errs <- runWeather() }()
	wg.Wait()
	close(errs)
	for err := range errs {
		if err != nil {
			t.Fatalf("插件执行失败: %v", err)
		}
	}

	var final pubsdk.StageContext
	if err := seg.ReadInto(&final); err != nil {
		t.Fatalf("ReadInto: %v", err)
	}
	got, _ := final.ToolResults[0].Result.(string)
	if got != clean {
		t.Fatalf("清洗结果被覆盖：期望 %q，实际 %q", clean, got)
	}
}

// 多插件高并发累加同一字段：总写入次数必须等于最终长度（零丢失零撕裂）。
// 对应实验 8（5 进程 × 300 轮），这里在单进程内用 goroutine 模拟并发扇出，
// 验证共享段 + 锁仲裁 + 脏字段写回三者组合的正确性。
func TestSegment_ConcurrentAppend_NoLostUpdate(t *testing.T) {
	// arena 需容纳 append-only 的中间垃圾：每轮写入长度递增，
	// 5 插件 × 40 轮 → 最长 200 字符，累计约 200*201/2 = 20100 字节，留足余量。
	seg := newTestSegment(t, 128*1024)
	if err := seg.WriteAll(&pubsdk.StageContext{FinalText: ""}); err != nil {
		t.Fatalf("WriteAll: %v", err)
	}

	lock := newStageLock()
	tags := []string{"A", "B", "C", "D", "E"}
	const iters = 40

	var wg sync.WaitGroup
	errCh := make(chan error, len(tags)*iters)

	for _, tag := range tags {
		wg.Add(1)
		go func(tag string) {
			defer wg.Done()
			for i := 0; i < iters; i++ {
				if err := lock.Acquire(tag); err != nil {
					errCh <- err
					return
				}
				var local pubsdk.StageContext
				if err := seg.ReadInto(&local); err != nil {
					lock.Release(tag)
					errCh <- err
					return
				}
				snap := TakeSnapshot(&local)
				local.FinalText += tag // 读-改-写
				if _, err := seg.WriteDirty(&local, snap); err != nil {
					lock.Release(tag)
					errCh <- err
					return
				}
				if err := lock.Release(tag); err != nil {
					errCh <- err
					return
				}
			}
		}(tag)
	}
	wg.Wait()
	close(errCh)
	for err := range errCh {
		if err != nil {
			t.Fatalf("并发写入失败: %v", err)
		}
	}

	var final pubsdk.StageContext
	if err := seg.ReadInto(&final); err != nil {
		t.Fatalf("ReadInto: %v", err)
	}

	// 关键断言：各标记出现次数之和 == 最终长度 ⇒ 无丢失、无撕裂
	total := 0
	counts := map[string]int{}
	for _, tag := range tags {
		c := strings.Count(final.FinalText, tag)
		counts[tag] = c
		total += c
	}
	if total != len(final.FinalText) {
		t.Fatalf("出现撕裂：各标记计数之和 %d != 最终长度 %d（counts=%v）",
			total, len(final.FinalText), counts)
	}
	want := len(tags) * iters
	if total != want {
		t.Fatalf("出现 lost update：期望 %d 次写入全部保留，实际 %d（counts=%v）",
			want, total, counts)
	}
	for tag, c := range counts {
		if c != iters {
			t.Errorf("插件 %s 的写入丢失：期望 %d 次，实际 %d 次", tag, iters, c)
		}
	}
}

// arena 用尽必须显式报错，不得静默截断（§4.4 风险登记）。
func TestSegment_ArenaExhaustionReturnsError(t *testing.T) {
	seg := newTestSegment(t, headerSize+ctxSize+256) // 极小 arena
	big := strings.Repeat("x", 1024)
	err := seg.WriteAll(&pubsdk.StageContext{FinalText: big})
	if err == nil {
		t.Fatal("arena 不足应报错，而非静默截断")
	}
	if !strings.Contains(err.Error(), "arena 空间不足") {
		t.Errorf("错误信息应说明 arena 不足，实际: %v", err)
	}
}

// 压实回收 append-only 垃圾，且不破坏现存字段。
func TestSegment_CompactReclaimsGarbage(t *testing.T) {
	seg := newTestSegment(t, 32*1024)
	if err := seg.WriteAll(&pubsdk.StageContext{FinalText: "初始"}); err != nil {
		t.Fatalf("WriteAll: %v", err)
	}

	// 反复改写同一字段，制造 append-only 垃圾
	for i := 0; i < 50; i++ {
		var local pubsdk.StageContext
		if err := seg.ReadInto(&local); err != nil {
			t.Fatalf("ReadInto: %v", err)
		}
		snap := TakeSnapshot(&local)
		local.FinalText = fmt.Sprintf("第%d次改写内容", i)
		if _, err := seg.WriteDirty(&local, snap); err != nil {
			t.Fatalf("WriteDirty: %v", err)
		}
	}

	usedBefore := seg.ArenaUsed()
	var beforeCtx pubsdk.StageContext
	if err := seg.ReadInto(&beforeCtx); err != nil {
		t.Fatalf("ReadInto: %v", err)
	}

	reclaimed := seg.Compact()
	if reclaimed == 0 {
		t.Error("应回收到垃圾空间")
	}
	if seg.ArenaUsed() >= usedBefore {
		t.Errorf("压实后已用空间应下降：%d → %d", usedBefore, seg.ArenaUsed())
	}

	var afterCtx pubsdk.StageContext
	if err := seg.ReadInto(&afterCtx); err != nil {
		t.Fatalf("压实后 ReadInto: %v", err)
	}
	if afterCtx.FinalText != beforeCtx.FinalText {
		t.Errorf("压实破坏了字段内容：%q → %q", beforeCtx.FinalText, afterCtx.FinalText)
	}
}