Cyrene_v2/internal/runtime/loop/memory.go

package loop

import (
	"context"
	"strings"
	"time"

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

// memoryExtractTimeout bounds one background extraction. It only governs the
// detached extraction context, so a hung model call cannot leak a goroutine for
// the process lifetime.
const memoryExtractTimeout = 60 * time.Second

// asyncExtractor is the extractor capability the loop detaches to a background
// goroutine: memory.LLMExtractor implements it through ExtractAsync (design/03 §1
// LLM 异步自动提取), returning a channel that delivers exactly one result.
// Deterministic extractors (memory.ManualExtractor) do not implement it and stay
// on the synchronous path.
type asyncExtractor interface {
	ExtractAsync(ctx context.Context, conversation []model.Message) (<-chan memory.AsyncResult, error)
}

// injectMemory retrieves the persona's relevant long-term memories and prepends
// them as a leading system message, so the model can ground its answer in recalled
// context (design §6 / §2). It is re-run every round via assembleContext so the
// retrieval reflects the latest conversation rather than a first-turn snapshot. It
// returns the input messages unchanged when no retriever/persona is configured,
// retrieval fails, or nothing relevant is found.
//
// The session env on ctx gates the retrieval: the owner-full env keeps the full
// recall (unchanged behavior); the restricted env (group chat / stranger,
// design/02 §3, design/11 §2) is read-only isolation — only non-owner-private
// project memory is injected so no owner-sensitive memory leaks to a group or
// stranger.
func (l *Loop) injectMemory(ctx context.Context, messages []model.Message) []model.Message {
	block := l.memoryBlock(ctx, messages)
	if block == "" {
		return messages
	}
	out := make([]model.Message, 0, len(messages)+1)
	out = append(out, model.TextMessage(model.RoleSystem, block))
	out = append(out, messages...)
	return out
}

// memoryBlock renders the round's memory block, or "" when there is nothing to
// inject (no retriever/persona, a retrieval error, or no relevant memory). It is
// split out from injectMemory because the assembly places other leading blocks
// between the memory block and the conversation (the role context of design/22
// §7-5), so it needs the block itself rather than a pre-prefixed slice.
func (l *Loop) memoryBlock(ctx context.Context, messages []model.Message) string {
	personaID := l.personaFor(ctx)
	if l.retriever == nil || personaID == "" {
		return ""
	}
	env := sessionEnvFrom(ctx)
	block, err := l.retrieveMemoryBlock(ctx, env, personaID, queryFromMessages(messages))
	if err != nil {
		l.log.Warn("memory inject", "error", err, "env", env.ActorEnv)
		return ""
	}
	return block
}

// retrieveMemoryBlock returns the memory injection block for env and the turn's
// role. The owner-full env keeps the pre-existing Inject behavior. The restricted
// env (group chat / stranger) injects only project-scope memory (人人可读, not
// owner-private) and marks the block as group-isolated so callers can tell an
// isolated read from the owner's full read.
func (l *Loop) retrieveMemoryBlock(ctx context.Context, env SessionEnv, personaID, query string) (string, error) {
	if !env.Restricted() {
		return l.retriever.Inject(ctx, personaID, query, 0)
	}
	results, err := l.retriever.Retrieve(ctx, personaID, query, 0)
	if err != nil {
		return "", err
	}
	shared := make([]memory.Result, 0, len(results))
	for _, r := range results {
		if r.Item.Scope == memory.ScopeProject {
			shared = append(shared, r)
		}
	}
	if len(shared) == 0 {
		return "", nil
	}
	return formatMemoryBlock(shared, "群聊隔离"), nil
}

// formatMemoryBlock builds the leading system block from retrieved results. The
// non-empty marker prefixes the block so an isolated group read is distinguishable
// from the owner's full read.
func formatMemoryBlock(results []memory.Result, marker string) string {
	var b strings.Builder
	if marker != "" {
		b.WriteString(marker)
		b.WriteString("：")
	}
	b.WriteString("相关记忆（按相关度 × 衰减排序）：")
	for _, res := range results {
		b.WriteByte('\n')
		b.WriteString("- [")
		b.WriteString(string(res.Item.Kind))
		b.WriteString("] ")
		b.WriteString(res.Item.Content)
	}
	return b.String()
}

// assembleContext builds the model-visible message list for one round (design
// §1.3 每轮 assemble, the loop's equivalent of OpenClaw's transformContext): it
// re-injects the persona's relevant long-term memory — retrieval is re-run every
// round from the current conversation, not cached from the first turn, so it
// supports cross-round refresh — folds in the per-session working-memory summary
// when present, renders the role/speaker attribution of the conversation (design/22
// §7-5), and applies the optional TransformContext hook last so a caller can
// project/filter what the model actually sees. It returns the input unchanged
// when nothing is configured.
//
// The existing order is preserved exactly: the leading blocks are the per-round
// time, the working-memory summary and the memory block, in that order; the role
// declaration sits right after them (immediately before the conversation it
// describes) and the speaker boundary notes are inserted into the conversation
// itself. A conversation with no speaker label renders no attribution at all, so
// every single-role path keeps the prompt it had before multi-character support.
func (l *Loop) assembleContext(ctx context.Context, messages []model.Message, round int) []model.Message {
	var leading []model.Message
	// Per-round "current time" so the model naturally knows the wall clock (not a
	// first-turn snapshot). Keyed on a configured timezone so tests and no-TZ runs
	// keep their exact leading-message counts. It sits before working memory / the
	// memory block and can still be projected by transformContext.
	if l.timeLocation != nil {
		leading = append(leading, l.currentTimeBlock())
	}
	if l.working != nil {
		if s := l.working.Summarize(); s != "" {
			leading = append(leading, model.TextMessage(model.RoleSystem, "工作记忆: "+s))
		}
	}
	if block := l.memoryBlock(ctx, messages); block != "" {
		leading = append(leading, model.TextMessage(model.RoleSystem, block))
	}
	decl, ok, view := l.attributeConversation(ctx, messages)
	if ok {
		leading = append(leading, decl)
	}
	out := append(leading, view...)
	if l.transformContext != nil {
		out = l.transformContext(ctx, out, round)
	}
	return out
}

// compactIfNeeded compacts the conversation when its estimated context exceeds
// the configured MaxContextTokens. It is a no-op when compaction is disabled
// (non-positive MaxContextTokens), the compactor is nil, the estimate is within
// budget, or the compactor returns an error — a compaction failure is logged and
// the conversation is kept as-is rather than failing the run.
func (l *Loop) compactIfNeeded(ctx context.Context, messages []model.Message) []model.Message {
	if l.maxContextTokens <= 0 || l.compactor == nil {
		return messages
	}
	compacted, err := l.compactor.Compact(ctx, messages, l.maxContextTokens)
	if err != nil {
		l.log.Warn("context compact", "error", err, "tokens", l.maxContextTokens)
		return messages
	}
	if compacted == nil {
		return messages
	}
	return compacted
}

// extractMemory runs the extraction hook (design §6) over the current
// conversation and persists any derived items. Errors are logged, never fatal, so
// a bad extraction cannot abort the run.
//
// In a restricted context (group chat / stranger, !MemWritable) extraction is
// skipped entirely: the conversation may read shared memory read-only and may
// never persist anything that could leak owner-sensitive information (design/02
// §3, design/11 §2).
func (l *Loop) extractMemory(ctx context.Context, messages []model.Message, step int) {
	env := sessionEnvFrom(ctx)
	if env.Restricted() {
		l.log.Debug("memory extract skipped: read-only context", "step", step, "env", env.ActorEnv)
		return
	}
	if l.extractor == nil || l.memory == nil {
		return
	}
	items, err := l.extractor.Extract(ctx, messages)
	if err != nil {
		l.log.Warn("memory extract", "step", step, "error", err)
		return
	}
	l.persistMemories(ctx, items, step)
}

// extractMemoryAsync triggers the extraction hook at the turn boundary without
// blocking the reply (design/03 §1: 每轮对话后后台异步抽取). When the configured
// extractor can extract asynchronously (the LLM extractor) the conversation is
// handed to a background goroutine and this call returns immediately, so a slow
// extraction model never delays the response. A deterministic extractor
// (ManualExtractor) does not implement asyncExtractor and is run inline, which
// keeps its cheap "记住 X" path observable and unchanged.
//
// The same gates as extractMemory apply: a restricted conversation is never
// extracted (read-only isolation, design/02 §3 / design/11 §2), and a missing
// extractor or store is a no-op. Extraction failures and store failures are
// logged with the loop's source and never propagated.
func (l *Loop) extractMemoryAsync(ctx context.Context, messages []model.Message, step int) {
	env := sessionEnvFrom(ctx)
	if env.Restricted() {
		l.log.Debug("memory extract skipped: read-only context", "step", step, "env", env.ActorEnv)
		return
	}
	if l.extractor == nil || l.memory == nil {
		return
	}
	async, ok := l.extractor.(asyncExtractor)
	if !ok {
		l.extractMemory(ctx, messages, step)
		return
	}
	// The turn's context dies with the request while the extraction deliberately
	// outlives it, so the goroutine gets a detached copy of the context (values —
	// the session env among them — are kept) under its own timeout.
	dctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), memoryExtractTimeout)
	// The caller keeps appending to messages as the turn continues, so the
	// goroutine reads a private snapshot instead of racing on the shared slice.
	snapshot := append([]model.Message(nil), messages...)
	results, err := async.ExtractAsync(dctx, snapshot)
	if err != nil {
		cancel()
		l.log.Warn("memory extract dispatch", "step", step, "error", err)
		return
	}
	go func() {
		defer cancel()
		res, ok := <-results
		if !ok {
			return
		}
		if res.Err != nil {
			l.log.Warn("memory extract", "step", step, "error", res.Err)
			return
		}
		l.persistMemories(dctx, res.Items, step)
	}()
}

// persistMemories writes extracted items to the memory store, stamping every item
// with the role that produced the turn. The assignment is unconditional on
// purpose: the extractor is deliberately role-agnostic (it no longer holds a
// persona baked in at construction, cmd/cyrene/extractor.go), so a multi-role
// process would otherwise file every memory under one role. A blank role (no
// persona configured at all) leaves the item as the extractor produced it. Store
// failures are logged and never fatal, so one rejected item cannot abort a turn.
func (l *Loop) persistMemories(ctx context.Context, items []memory.MemoryItem, step int) {
	personaID := l.personaFor(ctx)
	for _, it := range items {
		if personaID != "" {
			it.PersonaID = personaID
		}
		if _, err := l.memory.Add(ctx, it); err != nil {
			l.log.Warn("memory store", "step", step, "error", err)
		}
	}
}

// queryFromMessages derives a retrieval query from the conversation: the last
// user turn's text, falling back to the last message's text, then "".
func queryFromMessages(messages []model.Message) string {
	for i := len(messages) - 1; i >= 0; i-- {
		if messages[i].Role == model.RoleUser {
			return messageText(messages[i])
		}
	}
	if n := len(messages); n > 0 {
		return messageText(messages[n-1])
	}
	return ""
}

// messageText returns the concatenated plain-text payload of a message's text
// parts, which is the shared view used to build a retrieval query.
func messageText(m model.Message) string {
	var b strings.Builder
	for _, part := range m.Content {
		if part.Type != model.ContentTypeText {
			continue
		}
		if b.Len() > 0 {
			b.WriteByte(' ')
		}
		b.WriteString(part.Text)
	}
	return b.String()
}