Cyrene_v2/internal/runtime/multimodal/vision/shot_test.go

package vision

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"image"
	"image/color"
	"image/png"
	"os"
	"path/filepath"
	"strings"
	"testing"

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

// 本文件钉住 design/28 §7 判据 2 的**视觉侧**（切片 Y4 / B3）：一张已落盘的页面截图
// 到底以什么形态进上下文 —— 主模型能收图时是**真的图**（image 部件 + data: URI），
// 收不了图时是**工具形态的文字**，超内联上限时**先裁再送**，每一条都带如实标注。
//
// 用具名假件而不是真浏览器：本文件的被测对象是「图怎么进上下文」，浏览器那一层由
// `internal/capability/browser` 与 `cmd/cyrene/wire_browser*` 的用例负责，两处不混。

// shotPNG 造一张真实的小 PNG（真字节，因为视觉层对图像**字节嗅探**，假字节会被拒）。
func shotPNG(t *testing.T, w, h int) []byte {
	t.Helper()
	img := image.NewRGBA(image.Rect(0, 0, w, h))
	for y := 0; y < h; y++ {
		for x := 0; x < w; x++ {
			img.Set(x, y, color.RGBA{R: uint8(x), G: uint8(y), B: 0x40, A: 0xff})
		}
	}
	var buf bytes.Buffer
	if err := png.Encode(&buf, img); err != nil {
		t.Fatalf("编码 PNG: %v", err)
	}
	return buf.Bytes()
}

// writeShot 把 PNG 落进临时目录并返回与落盘结论同形的 ShotRef。
func writeShot(t *testing.T, data []byte, w, h int) ShotRef {
	t.Helper()
	path := filepath.Join(t.TempDir(), "shot-test.png")
	if err := os.WriteFile(path, data, 0o600); err != nil {
		t.Fatalf("写截图: %v", err)
	}
	return ShotRef{Path: path, Bytes: len(data), Width: w, Height: h, Format: "png"}
}

// fakeCropper 记录裁切区域并返回一段可控长度的 data: URI。
type fakeCropper struct {
	calls   []string
	uri     string
	failAll bool
}

func (c *fakeCropper) CropRegion(_ context.Context, path string, x, y, w, h int) (string, error) {
	c.calls = append(c.calls, fmt.Sprintf("%s|%d,%d,%d,%d", filepath.Base(path), x, y, w, h))
	if c.failAll {
		return "", errors.New("crop: 假失败")
	}
	return c.uri, nil
}

// visionInvoker 造一个「记录调用参数并返回固定描述」的工具调用器。
func visionInvoker(t *testing.T, calls *[]string, description string) ToolInvoker {
	t.Helper()
	return func(_ context.Context, name string, args map[string]any) (any, error) {
		*calls = append(*calls, fmt.Sprintf("%s|%v", name, args["image"]))
		if name == "ocr" {
			return map[string]any{"text": description}, nil
		}
		return map[string]any{"description": description}, nil
	}
}

// newShotFixture 造一个视觉路由夹具。sight 决定主模型能不能收图。
func newShotFixture(t *testing.T, sight bool, cropper ShotCropper, budget int) (*ShotAttacher, *[]string) {
	t.Helper()
	calls := &[]string{}
	rt := NewRouter(WithVisionTool("vision"), WithOCRTool("ocr"))
	res := NewMessageResolver(rt, func() ModelInfo {
		return ModelInfo{Name: "main", SupportsVision: sight}
	}, visionInvoker(t, calls, "页面上写着「渲染后才出现的文字」"))
	opts := []ShotAttacherOption{WithShotCropper(cropper)}
	if budget > 0 {
		opts = append(opts, WithShotInlineBudget(budget))
	}
	return NewShotAttacher(res, opts...), calls
}

// TestShotAttach_VisionModelGetsRealImage 钉住判据 2 的正面：主模型能收图时，
// 附件是**真的 image 部件**（data: URI），不是一句「图在哪」的路径。
func TestShotAttach_VisionModelGetsRealImage(t *testing.T) {
	t.Parallel()
	data := shotPNG(t, 8, 8)
	shot := writeShot(t, data, 8, 8)
	att, calls := newShotFixture(t, true, nil, 0)

	got, err := att.Attach(context.Background(), shot)
	if err != nil {
		t.Fatalf("Attach: %v", err)
	}
	if got.Route != ShotRouteMain {
		t.Errorf("Route = %q, want %q", got.Route, ShotRouteMain)
	}
	if got.Content.Type != model.ContentTypeImage {
		t.Fatalf("Content.Type = %q, want image（主模型能收图时必须给真图）", got.Content.Type)
	}
	if !strings.HasPrefix(got.Content.ImageURL, "data:image/png;base64,") {
		t.Errorf("内联结果不是 PNG data: URI：%.60s", got.Content.ImageURL)
	}
	if got.Note == "" {
		t.Error("Note 为空：模型无法据此知道自己看到的是哪张图")
	}
	if len(*calls) != 0 {
		t.Errorf("主模型能收图时不该调用 vision/OCR 工具，却调了 %v", *calls)
	}
	if got.Cropped {
		t.Error("小图不该被裁")
	}
}

// TestShotAttach_TextModelGetsToolForm 钉住另一条路：主模型收不了图时走工具形态，
// 图**不进**上下文，进去的是 vision 工具产出的文字。
func TestShotAttach_TextModelGetsToolForm(t *testing.T) {
	t.Parallel()
	shot := writeShot(t, shotPNG(t, 8, 8), 8, 8)
	att, calls := newShotFixture(t, false, nil, 0)

	got, err := att.Attach(context.Background(), shot)
	if err != nil {
		t.Fatalf("Attach: %v", err)
	}
	if got.Route != ShotRouteVisionTool {
		t.Errorf("Route = %q, want %q", got.Route, ShotRouteVisionTool)
	}
	if got.Content.Type != model.ContentTypeText {
		t.Fatalf("Content.Type = %q, want text（图不该进读不了图的主模型）", got.Content.Type)
	}
	if !strings.Contains(got.Content.Text, "渲染后才出现的文字") {
		t.Errorf("工具产出的描述没进内容：%q", got.Content.Text)
	}
	if len(*calls) != 1 || !strings.HasPrefix((*calls)[0], "vision|") {
		t.Fatalf("工具调用 = %v，want 一次 vision 且带截图路径", *calls)
	}
	if !strings.Contains((*calls)[0], shot.Path) {
		t.Errorf("工具拿到的不是截图路径：%v", *calls)
	}
}

// TestShotAttach_OCRToolRoute 钉住 OCR 档（路由器只有 OCR 工具时的路由与字段名不同：
// 读 "text" 而不是 "description"）。
func TestShotAttach_OCRToolRoute(t *testing.T) {
	t.Parallel()
	shot := writeShot(t, shotPNG(t, 8, 8), 8, 8)
	calls := &[]string{}
	res := NewMessageResolver(NewRouter(WithOCRTool("ocr")), func() ModelInfo {
		return ModelInfo{SupportsVision: false}
	}, visionInvoker(t, calls, "OCR 出来的字"))
	att := NewShotAttacher(res)

	got, err := att.Attach(context.Background(), shot)
	if err != nil {
		t.Fatalf("Attach: %v", err)
	}
	if got.Route != ShotRouteOCRTool {
		t.Errorf("Route = %q, want %q", got.Route, ShotRouteOCRTool)
	}
	if !strings.Contains(got.Content.Text, "OCR 出来的字") {
		t.Errorf("OCR 结果没进内容：%q", got.Content.Text)
	}
}

// TestShotAttach_OversizedShotsAreCroppedFirst 钉住主人点名的「大图先过 crop」：
// 超内联上限的截图先按区域裁一块送进去，且**如实标注**保留了哪一块、完整图在哪。
func TestShotAttach_OversizedShotsAreCroppedFirst(t *testing.T) {
	t.Parallel()
	shot := writeShot(t, shotPNG(t, 40, 40), 40, 40) // 真实字节数远大于 1
	cropper := &fakeCropper{uri: "data:image/png;base64,AAAA"}
	att, _ := newShotFixture(t, true, cropper, 1)

	got, err := att.Attach(context.Background(), shot)
	if err != nil {
		t.Fatalf("Attach: %v", err)
	}
	if got.Route != ShotRouteMain || got.Content.Type != model.ContentTypeImage {
		t.Fatalf("裁切后应仍是给主模型的图：route=%q type=%q", got.Route, got.Content.Type)
	}
	if !got.Cropped || got.CropRegion == "" {
		t.Errorf("裁切事实没被如实标注：Cropped=%v Region=%q", got.Cropped, got.CropRegion)
	}
	if len(cropper.calls) != 1 {
		t.Fatalf("裁切调用 = %v，want 恰好一次", cropper.calls)
	}
	// 区域规则：顶部锚定、整宽（x=0,y=0），高度夹在 (0, 原图高] 内。
	want := fmt.Sprintf("shot-test.png|0,0,40,%d", cropKeepHeight(40, 1.0/float64(shot.Bytes)))
	if cropper.calls[0] != want {
		t.Errorf("裁切区域 = %q, want %q", cropper.calls[0], want)
	}
	for _, kw := range []string{shot.Path, "image-crop", "40×40"} {
		if !strings.Contains(got.Note, kw) {
			t.Errorf("标注里缺 %q：%q", kw, got.Note)
		}
	}
}

// TestShotAttach_OversizedWithoutCropperFallsBackToToolForm 钉住降级方向：
// 没有裁切能力时**不硬塞**（也不假装成功），如实改走工具形态，并在标注里说明原因。
func TestShotAttach_OversizedWithoutCropperFallsBackToToolForm(t *testing.T) {
	t.Parallel()
	shot := writeShot(t, shotPNG(t, 40, 40), 40, 40)
	att, calls := newShotFixture(t, true, nil, 1)

	got, err := att.Attach(context.Background(), shot)
	if err != nil {
		t.Fatalf("Attach: %v", err)
	}
	if got.Route != ShotRouteVisionTool || got.Content.Type != model.ContentTypeText {
		t.Fatalf("没有裁切能力时应走工具形态：route=%q type=%q", got.Route, got.Content.Type)
	}
	if len(*calls) != 1 {
		t.Fatalf("工具调用 = %v，want 一次", *calls)
	}
	if !strings.Contains(got.Note, "无法裁切") && !strings.Contains(got.Note, "内联上限") {
		t.Errorf("降级原因没写在标注里：%q", got.Note)
	}
}

// TestShotAttach_CropStillTooLargeFallsBack 钉住「裁完还超限」这一档：不硬塞半张坏图，
// 降级到工具形态（图仍然被读到了，只是以文字形式）。
func TestShotAttach_CropStillTooLargeFallsBack(t *testing.T) {
	t.Parallel()
	shot := writeShot(t, shotPNG(t, 40, 40), 40, 40)
	cropper := &fakeCropper{uri: "data:image/png;base64," + strings.Repeat("A", MaxInlineImageChars)}
	att, _ := newShotFixture(t, true, cropper, 1)

	got, err := att.Attach(context.Background(), shot)
	if err != nil {
		t.Fatalf("Attach: %v", err)
	}
	if got.Route != ShotRouteVisionTool {
		t.Errorf("裁完仍超限时应降级到工具形态：%q", got.Route)
	}
	if len(cropper.calls) > shotCropMaxAttempts {
		t.Errorf("裁切次数 %d 超过上限 %d", len(cropper.calls), shotCropMaxAttempts)
	}
}

// TestShotAttach_CropSkippedWhenDimensionsUnknown 钉住「不猜区域」：字节超限但尺寸未知时
// **一次都不裁**（任何区域都会是编出来的），直接走工具形态。
func TestShotAttach_CropSkippedWhenDimensionsUnknown(t *testing.T) {
	t.Parallel()
	shot := writeShot(t, shotPNG(t, 40, 40), 0, 0)
	cropper := &fakeCropper{uri: "data:image/png;base64,AAAA"}
	att, _ := newShotFixture(t, true, cropper, 1)

	got, err := att.Attach(context.Background(), shot)
	if err != nil {
		t.Fatalf("Attach: %v", err)
	}
	if len(cropper.calls) != 0 {
		t.Errorf("尺寸未知时不该裁切，却调了 %v", cropper.calls)
	}
	if got.Route != ShotRouteVisionTool {
		t.Errorf("Route = %q, want 工具形态", got.Route)
	}
}

// TestShotAttach_NoToolNoRouteIsHonest 钉住最坏一档：主模型读不了图、也没绑任何工具时
// **如实报错**（绝不返回一个「成功但没有图」的附件 —— 那会让模型以为页面里什么都没有）。
func TestShotAttach_NoToolNoRouteIsHonest(t *testing.T) {
	t.Parallel()
	shot := writeShot(t, shotPNG(t, 8, 8), 8, 8)
	res := NewMessageResolver(NewRouter(), func() ModelInfo {
		return ModelInfo{SupportsVision: false}
	}, nil)
	att := NewShotAttacher(res)

	_, err := att.Attach(context.Background(), shot)
	if !errors.Is(err, ErrUnsupportedSource) {
		t.Fatalf("err = %v, want %v", err, ErrUnsupportedSource)
	}
}

// TestShotAttach_GuardsAndMissingFile 钉住两条入口守卫与一条真实失败：
// nil resolver ⇒ ErrNotConfigured；空路径 ⇒ 报错；路径不存在 ⇒ 内联失败后**不假装**成功。
func TestShotAttach_GuardsAndMissingFile(t *testing.T) {
	t.Parallel()
	if _, err := NewShotAttacher(nil).Attach(context.Background(), ShotRef{Path: "x.png"}); !errors.Is(err, ErrNotConfigured) {
		t.Errorf("nil resolver 的 err = %v, want %v", err, ErrNotConfigured)
	}
	att, _ := newShotFixture(t, true, nil, 0)
	if _, err := att.Attach(context.Background(), ShotRef{Path: "  "}); err == nil {
		t.Error("空路径应当报错")
	}
	missing := ShotRef{Path: filepath.Join(t.TempDir(), "not-there.png"), Bytes: 10, Width: 8, Height: 8}
	got, err := att.Attach(context.Background(), missing)
	if err == nil && got.Content.Type == model.ContentTypeImage {
		t.Error("文件不存在时不该交出一个 image 部件")
	}
	if err == nil && got.Route == ShotRouteMain {
		t.Error("文件不存在时不该报 Route=main")
	}
}