Cyrene_v2/cmd/cyrene/wire_document_budget_test.go

package main

import (
	"context"
	"encoding/json"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"testing"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/agent"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/documents"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/loop"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
	"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
	iotools "git.yeij.top/AskaEth/Cyrene_v2/plugins/tools/io"
)

// 本文件是切片 T1 §7.2 的**端到端**成对契约（S4 实测断裂的正面用例）：
// 装配期注入预算与文档读侧共用同一套字符坐标系，因此清单行给的续读游标真的能
// 被 document-read 用起来。预算推导本身的用例在 wire_context_test.go。

// docE2EName 是合法存储名形态（32 位十六进制 + .txt）。
const docE2EName = "0123456789abcdef0123456789abcdef.txt"

// docBudgetFixture 是一份端到端夹具：隔离数据目录 + 已落盘的正文 + 生产解析器 +
// 生产 document-read 工具 + 一个能捕获模型请求的循环。
type docBudgetFixture struct {
	cfg    *config.Config
	body   string
	ad     *compactProbeAdapter
	read   sdk.Tool
	cursor int
}

// newDocBudgetFixture 造夹具并跑**一次真实的回合**，返回模型看到的清单游标。
func newDocBudgetFixture(t *testing.T, injectBudget int) *docBudgetFixture {
	t.Helper()
	dataDir := t.TempDir()
	uploads := filepath.Join(dataDir, "uploads", "documents")
	if err := os.MkdirAll(uploads, 0o700); err != nil {
		t.Fatalf("建上传目录: %v", err)
	}
	fx := &docBudgetFixture{cfg: &config.Config{DataDir: dataDir}, body: strings.Repeat("abcdefghij", 20000)}
	if err := os.WriteFile(filepath.Join(uploads, docE2EName), []byte(fx.body), 0o600); err != nil {
		t.Fatalf("写上传件: %v", err)
	}

	adapter := &compactProbeAdapter{}
	ag, err := agent.New("你是昔涟。", nil, nil, agent.WithModel(adapter, nil))
	if err != nil {
		t.Fatalf("agent.New: %v", err)
	}
	l := loop.New(loop.Options{
		Agent:                  ag,
		DocumentResolver:       documents.NewResolver(documents.Config{DataDir: dataDir, UploadsDir: uploads}),
		DocumentMaxInjectChars: injectBudget,
		OnStream:               func(model.StreamEvent) {},
	})
	msg := model.DocumentMessage(model.RoleUser, "看看这个附件",
		model.Document{Name: docE2EName, MIME: "text/plain", Size: int64(len(fx.body))})
	if _, err := l.Run(context.Background(), "sess-doc", []model.Message{msg}); err != nil {
		t.Fatalf("Run: %v", err)
	}
	if len(adapter.reqs) == 0 {
		t.Fatal("没有捕获到模型请求")
	}
	fx.ad = adapter

	tools, err := iotools.New(fx.cfg)
	if err != nil {
		t.Fatalf("iotools.New: %v", err)
	}
	fx.read = documentReadToolByName(t, tools)

	req := adapter.reqs[len(adapter.reqs)-1]
	prompt := loopText(req.Messages[len(req.Messages)-1])
	fx.cursor, _ = cursorFromManifest(t, prompt)
	return fx
}

// TestDocumentCursor_ContinuationReadsBeyondTheInjectionBudget 是切片 T1 §7.2 的
// **端到端**成对契约（S4 实测断裂的正面用例）：正文远大于注入预算时
//
//	①清单行带 truncated + 续读游标（游标 = 实际注入的字符数）；
//	②把清单行里那个游标原样交给 `document-read`，拿到的是注入预算**之外**的正文。
//
// 三个角色都是生产实现：装配期是真实的 loop.resolveDocuments，sidecar 由真实的
// documents.Resolver 写出，读取走真实的 document-read 工具。唯一被单独钉住的是
// 「注入预算 = 1000 字符」，那正是本切片新增的可配置项。
func TestDocumentCursor_ContinuationReadsBeyondTheInjectionBudget(t *testing.T) {
	t.Parallel()
	const injectBudget = 1000
	fx := newDocBudgetFixture(t, injectBudget)

	if fx.cursor != injectBudget {
		t.Fatalf("清单游标 = %d, want %d（= 实际注入的字符数）", fx.cursor, injectBudget)
	}
	prompt := loopText(fx.ad.reqs[len(fx.ad.reqs)-1].Messages[len(fx.ad.reqs[len(fx.ad.reqs)-1].Messages)-1])
	_, injected := cursorFromManifest(t, prompt)
	if injected != fx.body[:injectBudget] {
		t.Fatalf("注入正文 = %d 字符，want 正文的前 %d 字符", len(injected), injectBudget)
	}

	// ② 照游标续读：拿到的是注入预算之外的正文（旧实现这里恒为空 —— 断裂所在）。
	first := callDocumentRead(t, fx.read, map[string]any{"name": docE2EName, "offset": fx.cursor})
	if first["offset"] != float64(fx.cursor) {
		t.Fatalf("续读 offset = %v, want %d", first["offset"], fx.cursor)
	}
	text, _ := first["text"].(string)
	if text == "" {
		t.Fatal("按清单游标续读读到了空正文 —— 清单给的承诺是假的")
	}
	if want := fx.body[fx.cursor : fx.cursor+len(text)]; text != want {
		t.Fatalf("续读正文不是注入预算之外的那一段：\n got %q\nwant %q", text, want)
	}

	// 一路翻页到底：拼起来必须正好是注入预算之外的全文（两边同一坐标系）。
	joined := text
	offset := fx.cursor + len(text)
	for i := 0; i < 100 && first["truncated"] == true; i++ {
		first = callDocumentRead(t, fx.read, map[string]any{"name": docE2EName, "offset": offset})
		page, _ := first["text"].(string)
		joined += page
		offset += len(page)
	}
	if joined != fx.body[fx.cursor:] {
		t.Fatalf("分页续读拼出 %d 字符，want %d（注入预算之外的全文）", len(joined), len(fx.body)-fx.cursor)
	}
	if total, _ := first["totalChars"].(float64); int(total) != len(fx.body) {
		t.Fatalf("totalChars = %v, want 全文 %d 字符", first["totalChars"], len(fx.body))
	}
}

// cursorFromManifest 从模型看到的上下文里取出清单行与续读游标（offset=N）。它刻意
// **解析文本**而不是读内部字段：模型能看到的就只有这一行，游标能不能被模型用起来，
// 正是这条契约要证明的事。返回游标与注入块（清单行 + 换行 + 注入正文）。
func cursorFromManifest(t *testing.T, prompt string) (int, string) {
	t.Helper()
	start := strings.Index(prompt, "[附件 ")
	if start < 0 {
		t.Fatalf("上下文里没有清单行：%q", prompt)
	}
	block := prompt[start:]
	line, injected, ok := strings.Cut(block, "]\n")
	if !ok {
		t.Fatalf("清单行没有闭合或没有正文行：%q", block)
	}
	if !strings.Contains(line, "已截断") {
		t.Fatalf("清单行没有如实标注截断：%q", line)
	}
	i := strings.Index(line, "offset=")
	if i < 0 {
		t.Fatalf("清单行没有续读游标：%q", line)
	}
	digits := line[i+len("offset="):]
	end := 0
	for end < len(digits) && digits[end] >= '0' && digits[end] <= '9' {
		end++
	}
	n, err := strconv.Atoi(digits[:end])
	if err != nil {
		t.Fatalf("续读游标不是整数：%q", digits[:end])
	}
	return n, injected
}

// documentReadToolByName 在真实工具集里找出 document-read 工具。
func documentReadToolByName(t *testing.T, tools []sdk.Tool) sdk.Tool {
	t.Helper()
	for _, tl := range tools {
		if tl.Name() == "document-read" {
			return tl
		}
	}
	t.Fatal("工具集里没有 document-read")
	return nil
}

// callDocumentRead 执行一次 document-read 并把返回的 JSON 文本解回 map。
func callDocumentRead(t *testing.T, tl sdk.Tool, args map[string]any) map[string]any {
	t.Helper()
	raw, err := tl.Execute(context.Background(), args)
	if err != nil {
		t.Fatalf("document-read: %v", err)
	}
	text, ok := raw.(string)
	if !ok {
		t.Fatalf("返回值类型 = %T, want string（JSON 文本）", raw)
	}
	out := map[string]any{}
	if err := json.Unmarshal([]byte(text), &out); err != nil {
		t.Fatalf("返回值不是合法 JSON: %v（%q）", err, text)
	}
	return out
}