Agent-System-Backend/internal/ctxasm/finish.go

package ctxasm

import (
	"crypto/sha256"
	"encoding/hex"
	"math"
	"regexp"
	"strconv"
	"strings"
	"sync"

	agentllm "agent-native/runtime/llm"
	"agent-native/runtime/reminders"
	core "agent-system-core"
)

// Injection 是一条在组装末尾才决定是否注入的条件条目。
type Injection struct {
	Keys        []string
	Secondary   string
	Logic       string
	Content     string
	Constant    bool
	ScanDepth   int
	BudgetChars int
}

// Tail 是前缀稳定之后附在末尾的个人指令、会话要求和条件注入。
type Tail struct {
	Instructions string
	Note         string
	Entries      []Injection
	History      []core.Message
	Window       int
}

var calib struct {
	mu     sync.Mutex
	est    int
	actual int
}

// NotePromptUsage 用上一请求的真实 prompt token 校准后续估算。
func NotePromptUsage(est, actual int) {
	if est <= 0 || actual <= 0 {
		return
	}
	calib.mu.Lock()
	calib.est = est
	calib.actual = actual
	calib.mu.Unlock()
}

func scaleTokens(est int) int {
	calib.mu.Lock()
	defer calib.mu.Unlock()
	if est <= 0 || calib.est <= 0 || calib.actual <= 0 {
		return est
	}
	return est * calib.actual / calib.est
}

// EstimateTokens 用 UTF-8 字节数 / 3.35 估算 token。
func EstimateTokens(text string) int {
	n := len([]byte(text))
	if n == 0 {
		return 0
	}
	return int(math.Ceil(float64(n) / 3.35))
}

func visibleHistory(history []core.Message) []core.Message {
	best := map[string]int{}
	for i, m := range history {
		if m.Excluded || m.VariantGroup == "" {
			continue
		}
		prev, ok := best[m.VariantGroup]
		if !ok || m.VariantIndex >= history[prev].VariantIndex {
			best[m.VariantGroup] = i
		}
	}
	out := make([]core.Message, 0, len(history))
	for i, m := range history {
		if m.Excluded {
			continue
		}
		if m.VariantGroup != "" && best[m.VariantGroup] != i {
			continue
		}
		out = append(out, m)
	}
	return out
}

// FinishAssembly 把末尾提醒附上、按窗口丢掉最老的历史，并写分段估算与前缀哈希。
func FinishAssembly(msgs []agentllm.Message, rep *Report, tail Tail) ([]agentllm.Message, *Report) {
	if rep == nil {
		rep = &Report{}
	}
	notes := tailNotes(tail)
	reserve := 0
	for _, n := range notes {
		reserve += EstimateTokens(n) + 8
	}
	msgs = trimOldest(msgs, tail.Window, reserve, rep)
	msgs = reminders.Attach(msgs, notes)
	fillEstimate(msgs, rep)
	return msgs, rep
}

// TailNotes 导出尾注渲染（转录权威模式下由调用方经 Reminders 逐轮注入）。
func TailNotes(tail Tail) []string { return tailNotes(tail) }

func tailNotes(tail Tail) []string {
	var notes []string
	if text := strings.TrimSpace(tail.Instructions); text != "" {
		notes = append(notes, "个人指令：\n"+text)
	}
	if text := strings.TrimSpace(tail.Note); text != "" {
		notes = append(notes, "本会话要求：\n"+text)
	}
	for _, e := range tail.Entries {
		if strings.TrimSpace(e.Content) == "" {
			continue
		}
		if !e.Constant && !entryHit(e, tail.History) {
			continue
		}
		body := e.Content
		if e.BudgetChars > 0 {
			r := []rune(body)
			if len(r) > e.BudgetChars {
				body = string(r[:e.BudgetChars])
			}
		}
		notes = append(notes, body)
	}
	return notes
}

func entryHit(e Injection, history []core.Message) bool {
	if len(e.Keys) == 0 {
		return false
	}
	depth := e.ScanDepth
	if depth <= 0 {
		depth = 4
	}
	start := 0
	if len(history) > depth {
		start = len(history) - depth
	}
	var b strings.Builder
	for _, m := range history[start:] {
		for _, p := range m.Parts {
			if p.Type == core.PartText {
				b.WriteString(p.Text)
				b.WriteByte('\n')
			}
		}
	}
	text := b.String()
	primary := matchAny(e.Keys, text)
	sec := strings.TrimSpace(e.Secondary)
	if sec == "" {
		return primary
	}
	secondary := matchOne(sec, text)
	switch e.Logic {
	case "or":
		return primary || secondary
	case "not":
		return primary && !secondary
	default:
		return primary && secondary
	}
}

func matchAny(keys []string, text string) bool {
	for _, k := range keys {
		if matchOne(k, text) {
			return true
		}
	}
	return false
}

func matchOne(key, text string) bool {
	key = strings.TrimSpace(key)
	if key == "" {
		return false
	}
	if len(key) > 2 && strings.HasPrefix(key, "/") && strings.HasSuffix(key, "/") {
		re, err := regexp.Compile(key[1 : len(key)-1])
		if err == nil {
			return re.MatchString(text)
		}
	}
	return strings.Contains(strings.ToLower(text), strings.ToLower(key))
}

func trimOldest(msgs []agentllm.Message, window, reserve int, rep *Report) []agentllm.Message {
	if window <= 0 {
		return msgs
	}
	limit := window - reserve
	if limit < 1 {
		limit = 1
	}
	dropped := 0
	for estimateAll(msgs) > limit {
		start, end := oldestDropRange(msgs)
		if start < 0 {
			break
		}
		dropped += end - start
		msgs = append(msgs[:start], msgs[end:]...)
	}
	if dropped > 0 {
		rep.Truncations = append(rep.Truncations, fmtDropped(dropped))
	}
	return msgs
}

func fmtDropped(n int) string {
	return "按上下文窗口丢掉最早的 " + itoa(n) + " 条历史"
}

func itoa(n int) string { return strconv.Itoa(n) }

// oldestDropRange 返回下一次丢弃的区间 [start, end)（找不到返回 -1,-1）。
// 丢弃按工具组对齐：带 tool_calls 的 assistant 与其连续 tool 结果同退，
// 开头的孤立 tool 结果组整组同退——裁切后既不剩无结果的调用，
// 也不剩无调用的结果（与库层 alignToolTail 同语义、方向相反）。
// 组越过最后一条（当前轮）时不再丢弃：窗口允许略超，配对不可破坏。
func oldestDropRange(msgs []agentllm.Message) (int, int) {
	if len(msgs) < 3 {
		return -1, -1
	}
	// 保留开头的系统前缀和最后一条（当前轮）。
	start := 0
	for start < len(msgs) && msgs[start].Role == agentllm.RoleSystem {
		start++
	}
	if start >= len(msgs)-1 {
		return -1, -1
	}
	end := start + 1
	m := msgs[start]
	switch {
	case m.Role == agentllm.RoleAssistant && len(m.ToolCalls) > 0:
		// assistant 及其连续 tool 结果整组同退。
		for end < len(msgs) && msgs[end].Role == agentllm.RoleTool {
			end++
		}
	case m.Role == agentllm.RoleTool:
		// 开头即孤立 tool 结果（其 assistant 必已被丢或不存在）：整组同退。
		for end < len(msgs) && msgs[end].Role == agentllm.RoleTool {
			end++
		}
	}
	if end > len(msgs)-1 {
		return -1, -1
	}
	return start, end
}

func estimateAll(msgs []agentllm.Message) int {
	n := 0
	for _, m := range msgs {
		n += EstimateTokens(m.Content)
	}
	return n
}

func fillEstimate(msgs []agentllm.Message, rep *Report) {
	sys, hist, cur := 0, 0, 0
	var prefix strings.Builder
	last := len(msgs) - 1
	for i, m := range msgs {
		tok := EstimateTokens(m.Content)
		switch {
		case m.Role == agentllm.RoleSystem && i < last:
			sys += tok
			prefix.WriteString(m.Content)
			prefix.WriteByte('\n')
		case i == last:
			cur += tok
		default:
			hist += tok
		}
	}
	rep.Sections = []Section{
		{Name: "system", EstTokens: scaleTokens(sys)},
		{Name: "history", EstTokens: scaleTokens(hist)},
		{Name: "current", EstTokens: scaleTokens(cur)},
	}
	sum := sha256.Sum256([]byte(prefix.String()))
	rep.PrefixHash = hex.EncodeToString(sum[:8])
	rep.EstChars = scaleTokens(sys + hist + cur)
}

// Section 是一段上下文的估算。
type Section struct {
	Name      string `json:"name"`
	EstTokens int    `json:"estTokens"`
}