Cyrene_v2/internal/runtime/loop/memory_async_test.go

package loop

import (
	"context"
	"testing"
	"time"

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

// The LLM extractor is the async-capable implementation behind design/03 §1: the
// loop type-asserts this capability to detach extraction to a background
// goroutine, so the contract is locked here rather than discovered at runtime.
var (
	_ memory.Extractor = (*memory.LLMExtractor)(nil)
	_ asyncExtractor   = (*memory.LLMExtractor)(nil)
	_ memory.Extractor = (*memory.ManualExtractor)(nil)
)

// gatedExtractor is a memory.Extractor that also implements asyncExtractor. Its
// extraction blocks until release is closed (or ctx is canceled), so a test can
// prove the loop's turn finishes while extraction is still in flight, and that
// the extraction survives the caller's own cancellation.
type gatedExtractor struct {
	release  chan struct{}
	started  chan struct{}
	items    []memory.MemoryItem
	dispatch int
}

func newGatedExtractor() *gatedExtractor {
	return &gatedExtractor{
		release: make(chan struct{}),
		started: make(chan struct{}, 1),
		items: []memory.MemoryItem{{
			Scope:       memory.ScopeUser,
			Kind:        memory.KindRelationship,
			Content:     "用户习惯：爱喝冰美式",
			Confidence:  0.9,
			DecayWeight: 1,
		}},
	}
}

func (g *gatedExtractor) Extract(ctx context.Context, _ []model.Message) ([]memory.MemoryItem, error) {
	select {
	case g.started <- struct{}{}:
	default:
	}
	select {
	case <-g.release:
		return g.items, nil
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

func (g *gatedExtractor) ExtractAsync(ctx context.Context, conv []model.Message) (<-chan memory.AsyncResult, error) {
	g.dispatch++
	out := make(chan memory.AsyncResult, 1)
	go func() {
		defer close(out)
		items, err := g.Extract(ctx, conv)
		out <- memory.AsyncResult{Items: items, Err: err}
	}()
	return out, nil
}

// TestExtractMemoryAsync_LLMExtractorThroughLoop wires the real LLM extractor
// (not a fake) into the loop: after the turn settles, the extractor's own model
// call runs in the background and its parsed items land in the store, so the
// production path (design/03 §1 LLM 异步自动提取) is exercised end to end.
func TestExtractMemoryAsync_LLMExtractorThroughLoop(t *testing.T) {
	t.Parallel()
	store := newTestMemoryStore(t)
	ad := &mockAdapter{chat: []mgResp{
		{resp: model.Response{Content: "记住了"}},
		{resp: model.Response{Content: `[{"content":"用户爱喝冰美式","scope":"user","kind":"relationship","confidence":0.92}]`}},
	}}
	a := newTestAgent(t, ad, nil)
	l := New(Options{
		Agent: a, Memory: store, PersonaID: "persona-a", MaxRounds: 8,
		Extractor: memory.NewLLMExtractor(ad, "persona-a", "owner"),
	})

	if _, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "我喜欢冰美式")}); err != nil {
		t.Fatalf("Run: %v", err)
	}
	waitForMemoryCount(t, store, "persona-a", 1)

	items, err := store.ListByPersona(context.Background(), "persona-a")
	if err != nil {
		t.Fatalf("ListByPersona: %v", err)
	}
	if items[0].Content != "用户爱喝冰美式" || items[0].Scope != memory.ScopeUser {
		t.Fatalf("stored item = %+v, want the extracted user-scope memory", items[0])
	}
	if items[0].PersonaID != "persona-a" {
		t.Fatalf("item persona = %q, want persona-a", items[0].PersonaID)
	}

	extractCalls := 0
	for _, req := range ad.reqs {
		if req.Purpose == "memory_extract" {
			extractCalls++
		}
	}
	if extractCalls != 1 {
		t.Fatalf("memory_extract model calls = %d, want 1 (the LLM extractor drove the background call)", extractCalls)
	}
}

// waitForMemoryCount polls the persona's stored memories until want is reached,
// failing the test when the extraction never lands.
func waitForMemoryCount(t *testing.T, store *memory.Store, personaID string, want int) {
	t.Helper()
	deadline := time.Now().Add(2 * time.Second)
	for {
		items, err := store.ListByPersona(context.Background(), personaID)
		if err != nil {
			t.Fatalf("ListByPersona: %v", err)
		}
		if len(items) == want {
			return
		}
		if time.Now().After(deadline) {
			t.Fatalf("stored memories = %d, want %d (extraction never landed)", len(items), want)
		}
		time.Sleep(5 * time.Millisecond)
	}
}

// TestExtractMemoryAsync_DoesNotBlockTurn proves the turn returns while the LLM
// extraction is still running, and that the extraction completes (and persists)
// after the caller's context has already been canceled — the detached context
// keeps a request-scoped cancellation from aborting the background extraction.
func TestExtractMemoryAsync_DoesNotBlockTurn(t *testing.T) {
	t.Parallel()
	store := newTestMemoryStore(t)
	ext := newGatedExtractor()
	ad := &mockAdapter{chat: []mgResp{{resp: model.Response{Content: "记住了"}}}}
	a := newTestAgent(t, ad, nil)
	l := New(Options{Agent: a, Memory: store, Extractor: ext, PersonaID: "persona-a", MaxRounds: 8})

	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() {
		_, err := l.Run(ctx, "sess", []model.Message{model.TextMessage(model.RoleUser, "记住 我爱喝冰美式")})
		done <- err
	}()

	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("Run: %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("Run blocked on memory extraction: the async path must not gate the reply")
	}

	if ext.dispatch != 1 {
		t.Fatalf("async dispatch count = %d, want 1 (one extraction per settled turn)", ext.dispatch)
	}
	items, err := store.ListByPersona(context.Background(), "persona-a")
	if err != nil {
		t.Fatalf("ListByPersona: %v", err)
	}
	if len(items) != 0 {
		t.Fatalf("store already holds %d items after Run, want 0 (extraction still in flight)", len(items))
	}

	// The turn is over, so its context is canceled; the background extraction
	// must still finish and persist.
	cancel()
	select {
	case <-ext.started:
	case <-time.After(2 * time.Second):
		t.Fatal("extraction goroutine never started")
	}
	close(ext.release)
	waitForMemoryCount(t, store, "persona-a", 1)
}

// TestExtractMemoryAsync_RestrictedEnvNeverDispatches proves the group-chat /
// stranger isolation still gates the async path: nothing is dispatched, so no
// owner-private memory can be written from a read-only conversation.
func TestExtractMemoryAsync_RestrictedEnvNeverDispatches(t *testing.T) {
	t.Parallel()
	store := newTestMemoryStore(t)
	ext := newGatedExtractor()
	l := New(Options{Memory: store, Extractor: ext, PersonaID: "persona-a"})

	ctx := WithSessionEnv(context.Background(), SessionEnv{ActorEnv: restrict.ActorGroupMember, MemWritable: false})
	l.extractMemoryAsync(ctx, []model.Message{model.TextMessage(model.RoleUser, "记住 我爱喝冰美式")}, 1)

	if ext.dispatch != 0 {
		t.Fatalf("async dispatch count = %d, want 0 (restricted context is read-only)", ext.dispatch)
	}
	items, err := store.ListByPersona(context.Background(), "persona-a")
	if err != nil {
		t.Fatalf("ListByPersona: %v", err)
	}
	if len(items) != 0 {
		t.Fatalf("restricted async extract wrote %d items, want 0 (read-only isolation)", len(items))
	}
}

// TestExtractMemoryAsync_ManualExtractorStaysSync proves the deterministic path
// is unchanged: an extractor without the async capability is still run inline, so
// the explicit "记住 X" memory is stored by the time the call returns.
func TestExtractMemoryAsync_ManualExtractorStaysSync(t *testing.T) {
	t.Parallel()
	store := newTestMemoryStore(t)
	l := New(Options{Memory: store, Extractor: memory.NewManualExtractor("persona-a", "owner"), PersonaID: "persona-a"})

	ctx := WithSessionEnv(context.Background(), OwnerEnv())
	l.extractMemoryAsync(ctx, []model.Message{model.TextMessage(model.RoleUser, "记住 我爱喝冰美式")}, 1)

	items, err := store.ListByPersona(context.Background(), "persona-a")
	if err != nil {
		t.Fatalf("ListByPersona: %v", err)
	}
	if len(items) != 1 {
		t.Fatalf("stored memories = %d, want 1 (deterministic extractor runs inline)", len(items))
	}
}