Cyrene_v2/plugins/tools/io/crop_test.go

package io

import (
	"bytes"
	"context"
	"encoding/base64"
	"image"
	"image/color"
	"image/png"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
	"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
)

// gradientPNG writes an image whose pixel (x,y) encodes its own coordinates in the
// red/green channels. That makes a crop self-verifying: if the crop's top-left
// pixel is not RGBA{sourceX, sourceY, 0, 255} the wrong region was returned.
func gradientPNG(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), A: 255})
		}
	}
	var buf bytes.Buffer
	if err := png.Encode(&buf, img); err != nil {
		t.Fatalf("png.Encode: %v", err)
	}
	return buf.Bytes()
}

// decodeCropResult pulls the PNG out of a crop tool result and decodes it.
func decodeCropResult(t *testing.T, result any) image.Image {
	t.Helper()
	m, ok := result.(map[string]any)
	if !ok {
		t.Fatalf("result type = %T, want map[string]any", result)
	}
	uri, _ := m["image"].(string)
	if !strings.HasPrefix(uri, "data:image/png;base64,") {
		t.Fatalf("image field = %q, want a PNG data URI", uri)
	}
	raw, err := base64.StdEncoding.DecodeString(strings.TrimPrefix(uri, "data:image/png;base64,"))
	if err != nil {
		t.Fatalf("decode data URI: %v", err)
	}
	img, _, err := image.Decode(bytes.NewReader(raw))
	if err != nil {
		t.Fatalf("decode crop: %v", err)
	}
	return img
}

func newCropToolForTest(t *testing.T, extra ...Option) (*cropTool, string) {
	t.Helper()
	dir := t.TempDir()
	set := buildTools(t, &config.Config{DataDir: dir}, extra...)
	tool, ok := findTool(t, set, "image-crop").(*cropTool)
	if !ok {
		t.Fatalf("image-crop tool has type %T, want *cropTool", findTool(t, set, "image-crop"))
	}
	return tool, dir
}

// TestCropTool_CropsRequestedRegion pins the region arithmetic: the returned crop
// must be exactly the requested region with the requested pixels, not merely the
// right size (裁决 14.4.5③).
func TestCropTool_CropsRequestedRegion(t *testing.T) {
	t.Parallel()
	tool, dir := newCropToolForTest(t)
	path := filepath.Join(dir, "src.png")
	if err := os.WriteFile(path, gradientPNG(t, 100, 80), filePerm); err != nil {
		t.Fatalf("write source: %v", err)
	}

	res, err := tool.Execute(context.Background(), map[string]any{
		"image": path, "x": 10, "y": 20, "width": 30, "height": 40,
	})
	if err != nil {
		t.Fatalf("Execute: %v", err)
	}
	img := decodeCropResult(t, res)
	if img.Bounds().Dx() != 30 || img.Bounds().Dy() != 40 {
		t.Fatalf("crop size = %dx%d, want 30x40", img.Bounds().Dx(), img.Bounds().Dy())
	}
	got := color.NRGBAModel.Convert(img.At(0, 0)).(color.NRGBA)
	if got.R != 10 || got.G != 20 {
		t.Fatalf("top-left pixel = %v, want the source pixel (10,20)", got)
	}
	last := color.NRGBAModel.Convert(img.At(29, 39)).(color.NRGBA)
	if last.R != 39 || last.G != 59 {
		t.Fatalf("bottom-right pixel = %v, want the source pixel (39,59)", last)
	}
}

// TestCropTool_RejectsRegionOutsideImage pins fail-loud instead of clamping: a
// clamp would hand the model a different region than it asked for.
func TestCropTool_RejectsRegionOutsideImage(t *testing.T) {
	t.Parallel()
	tool, dir := newCropToolForTest(t)
	path := filepath.Join(dir, "src.png")
	if err := os.WriteFile(path, gradientPNG(t, 50, 50), filePerm); err != nil {
		t.Fatalf("write source: %v", err)
	}
	cases := []map[string]any{
		{"image": path, "x": 40, "y": 0, "width": 20, "height": 10},
		{"image": path, "x": 0, "y": 40, "width": 10, "height": 20},
		{"image": path, "x": -1, "y": 0, "width": 10, "height": 10},
		{"image": path, "x": 0, "y": 0, "width": 0, "height": 10},
		{"image": path, "x": 0, "y": 0, "width": 10, "height": -5},
	}
	for _, args := range cases {
		if _, err := tool.Execute(context.Background(), args); err == nil {
			t.Errorf("args %v accepted, want an out-of-bounds error", args)
		}
	}
}

// TestCropTool_RejectsPathOutsideSandbox pins the file sandbox reuse: cropping can
// never reach a file the file tool could not read.
func TestCropTool_RejectsPathOutsideSandbox(t *testing.T) {
	t.Parallel()
	tool, _ := newCropToolForTest(t)
	outside := filepath.Join(t.TempDir(), "outside.png")
	if err := os.WriteFile(outside, gradientPNG(t, 10, 10), filePerm); err != nil {
		t.Fatalf("write outside image: %v", err)
	}
	if _, err := tool.Execute(context.Background(), map[string]any{
		"image": outside, "x": 0, "y": 0, "width": 2, "height": 2,
	}); err == nil {
		t.Fatal("path outside the data dir accepted, want a sandbox rejection")
	}
}

// TestCropTool_AcceptsDataURIAndURL pins the other two sources: a data: URI and an
// http(s) URL both decode into a crop.
func TestCropTool_AcceptsDataURIAndURL(t *testing.T) {
	t.Parallel()
	src := gradientPNG(t, 20, 20)
	uri := "data:image/png;base64," + base64.StdEncoding.EncodeToString(src)

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "image/png")
		_, _ = w.Write(src)
	}))
	t.Cleanup(srv.Close)

	tool, _ := newCropToolForTest(t, WithAllowedHosts(serverHost(srv)))
	cases := []string{uri, srv.URL + "/img.png"}
	for _, source := range cases {
		res, err := tool.Execute(context.Background(), map[string]any{
			"image": source, "x": 5, "y": 5, "width": 4, "height": 3,
		})
		if err != nil {
			t.Fatalf("Execute(%s): %v", source, err)
		}
		if img := decodeCropResult(t, res); img.Bounds().Dx() != 4 || img.Bounds().Dy() != 3 {
			t.Fatalf("source %s crop size = %dx%d, want 4x3", source, img.Bounds().Dx(), img.Bounds().Dy())
		}
	}
}

// TestCropTool_RejectsUnsupportedDataURI pins a non-base64 data URI as an error
// rather than a silent empty decode.
func TestCropTool_RejectsUnsupportedDataURI(t *testing.T) {
	t.Parallel()
	tool, _ := newCropToolForTest(t)
	if _, err := tool.Execute(context.Background(), map[string]any{
		"image": "data:image/png,plaintext", "x": 0, "y": 0, "width": 1, "height": 1,
	}); err == nil {
		t.Fatal("non-base64 data URI accepted, want an error")
	}
}

// TestCropTool_RegisteredInSet pins that the tool ships with the default set and
// advertises a usable schema (the agent only sees it if it is registered).
func TestCropTool_RegisteredInSet(t *testing.T) {
	t.Parallel()
	tool, _ := newCropToolForTest(t)
	if tool.Name() != "image-crop" {
		t.Errorf("Name = %q, want image-crop", tool.Name())
	}
	params, err := tool.Parameters()
	if err != nil {
		t.Fatalf("Parameters: %v", err)
	}
	if params["type"] != "object" {
		t.Errorf("schema type = %v, want object", params["type"])
	}
	var _ sdk.Tool = tool
	// A nil/mis-typed argument set must be an error, not a panic.
	if _, err := tool.Execute(context.Background(), "not-a-map"); err == nil {
		t.Error("non-object args accepted, want an error")
	}
	if _, err := tool.Execute(context.Background(), map[string]any{}); err == nil {
		t.Error("missing image accepted, want an error")
	}
}