Cyrene_v2/plugins/tools/io/crop.go

package io

import (
	"bytes"
	"context"
	"encoding/base64"
	"fmt"
	"image"
	"image/draw"
	_ "image/gif"  // register the GIF decoder for image.Decode
	_ "image/jpeg" // register the JPEG decoder for image.Decode
	"image/png"
	"log/slog"
	"net/http"
	"os"
	"strings"

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

// Compile-time assertion that cropTool implements sdk.Tool.
var _ sdk.Tool = (*cropTool)(nil)

// Crop output limits. An upstream vision/OCR service downsamples a large image
// before it sees it (裁决 14.4.5③), so the way to read a region in full detail is
// to crop it and re-read the crop; these caps keep a modelled crop request from
// turning into an unbounded allocation.
const (
	// maxCropSourceBytes caps the decoded source image body.
	maxCropSourceBytes = 50 << 20
	// maxCropPixels caps the requested region's pixel count (~16 MP).
	maxCropPixels = 4096 * 4096
	// cropDataURIMax is the largest crop payload the tool will return inline. A
	// bigger crop is rejected rather than silently truncated: the model must
	// narrow the region, not receive a half image it cannot describe.
	cropDataURIMax = 8 << 20
)

// cropTool crops an axis-aligned region out of an image and returns it as a
// base64 data URI, so the agent can re-read a zoomed region of a large image with
// the vision or OCR tool (裁决 14.4.5③：「图片按区域裁切」，供模型自行裁大图分区
// 精读). Path inputs go through the file tool's sandbox — this tool never widens
// file access; URL inputs use the shared SSRF-guarded client.
type cropTool struct {
	file         *fileTool
	client       *http.Client
	allowedHosts []string
	log          *slog.Logger
}

func newCropTool(t *tools, file *fileTool) *cropTool {
	return &cropTool{
		file:         file,
		client:       t.client,
		allowedHosts: t.allowedHosts,
		log:          t.log,
	}
}

// Name implements sdk.Tool.
func (c *cropTool) Name() string { return "image-crop" }

// Description implements sdk.Tool. It names the downsampling reason, because the
// model decides whether to crop from this text alone.
func (c *cropTool) Description() string {
	return "Crop a rectangular region out of an image and return the crop (PNG data URI) plus its size. " +
		"Use it to read a large image in parts: an image service downsamples big images, so cropping a region " +
		"and re-reading the crop recovers the detail. Coordinates are 0-based pixels from the top-left."
}

// Parameters implements sdk.Tool.
func (c *cropTool) Parameters() (map[string]any, error) {
	return map[string]any{
		"type": "object",
		"properties": map[string]any{
			"image":  map[string]any{"type": "string", "description": "The image: a local path, an http(s) URL, or a base64/data: URI."},
			"x":      map[string]any{"type": "integer", "description": "Left edge of the region, in pixels (0-based)."},
			"y":      map[string]any{"type": "integer", "description": "Top edge of the region, in pixels (0-based)."},
			"width":  map[string]any{"type": "integer", "description": "Region width in pixels."},
			"height": map[string]any{"type": "integer", "description": "Region height in pixels."},
		},
		"required": []any{"image", "x", "y", "width", "height"},
	}, nil
}

// Execute implements sdk.Tool.
func (c *cropTool) Execute(ctx context.Context, args any) (any, error) {
	m, ok := args.(map[string]any)
	if !ok {
		return nil, fmt.Errorf("io: image-crop: args must be an object")
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	src := argString(m, "image")
	if strings.TrimSpace(src) == "" {
		return nil, fmt.Errorf("io: image-crop: image is required")
	}
	region, err := cropRegion(m)
	if err != nil {
		return nil, err
	}
	data, err := c.load(ctx, src)
	if err != nil {
		return nil, err
	}
	img, _, err := image.Decode(bytes.NewReader(data))
	if err != nil {
		return nil, fmt.Errorf("io: image-crop: decode image: %w", err)
	}
	bounds := img.Bounds()
	if err := checkRegion(region, bounds); err != nil {
		return nil, err
	}
	cropped := cropToRGBA(img, region)
	payload, err := encodePNG(cropped)
	if err != nil {
		return nil, err
	}
	c.log.Debug("io image-crop: cropped region",
		"source_width", bounds.Dx(), "source_height", bounds.Dy(),
		"x", region.Min.X, "y", region.Min.Y,
		"width", region.Dx(), "height", region.Dy(), "bytes", len(payload))
	return map[string]any{
		"image":  "data:image/png;base64," + base64.StdEncoding.EncodeToString(payload),
		"format": "image/png",
		"width":  region.Dx(),
		"height": region.Dy(),
		"region": map[string]any{"x": region.Min.X, "y": region.Min.Y, "width": region.Dx(), "height": region.Dy()},
		"source": map[string]any{"width": bounds.Dx(), "height": bounds.Dy()},
	}, nil
}

// cropRegion reads and validates the region arguments. A non-positive width or
// height, or a negative offset, is rejected (never clamped: a clamped crop would
// silently hand the model a different region than it asked for).
func cropRegion(m map[string]any) (image.Rectangle, error) {
	x, ok := argInt(m, "x")
	if !ok {
		return image.Rectangle{}, fmt.Errorf("io: image-crop: x must be an integer")
	}
	y, ok := argInt(m, "y")
	if !ok {
		return image.Rectangle{}, fmt.Errorf("io: image-crop: y must be an integer")
	}
	w, ok := argInt(m, "width")
	if !ok {
		return image.Rectangle{}, fmt.Errorf("io: image-crop: width must be an integer")
	}
	h, ok := argInt(m, "height")
	if !ok {
		return image.Rectangle{}, fmt.Errorf("io: image-crop: height must be an integer")
	}
	if x < 0 || y < 0 {
		return image.Rectangle{}, fmt.Errorf("io: image-crop: x and y must be >= 0, got (%d,%d)", x, y)
	}
	if w <= 0 || h <= 0 {
		return image.Rectangle{}, fmt.Errorf("io: image-crop: width and height must be > 0, got (%d,%d)", w, h)
	}
	if int64(w)*int64(h) > maxCropPixels {
		return image.Rectangle{}, fmt.Errorf("io: image-crop: region of %dx%d exceeds the %d-pixel limit", w, h, maxCropPixels)
	}
	return image.Rect(x, y, x+w, y+h), nil
}

// checkRegion rejects a region that is not fully inside the image.
func checkRegion(r image.Rectangle, bounds image.Rectangle) error {
	if !r.In(bounds) {
		return fmt.Errorf("io: image-crop: region %v is outside the image %v", r, bounds)
	}
	return nil
}

// cropToRGBA copies the region into a fresh RGBA image so the result does not
// alias the (possibly huge) source and always re-encodes losslessly.
func cropToRGBA(img image.Image, r image.Rectangle) *image.RGBA {
	dst := image.NewRGBA(image.Rect(0, 0, r.Dx(), r.Dy()))
	draw.Draw(dst, dst.Bounds(), img, r.Min, draw.Src)
	return dst
}

// encodePNG encodes the crop, refusing a payload above the inline cap so a huge
// crop is reported instead of flooding the tool result.
func encodePNG(img image.Image) ([]byte, error) {
	var buf bytes.Buffer
	if err := png.Encode(&buf, img); err != nil {
		return nil, fmt.Errorf("io: image-crop: encode crop: %w", err)
	}
	if buf.Len() > cropDataURIMax {
		return nil, fmt.Errorf("io: image-crop: crop is %d bytes, above the %d-byte inline limit; crop a smaller region", buf.Len(), cropDataURIMax)
	}
	return buf.Bytes(), nil
}

// load reads the source image: a local path through the file sandbox, an http(s)
// URL through the shared SSRF-guarded client, or a base64/data: URI payload.
func (c *cropTool) load(ctx context.Context, src string) ([]byte, error) {
	switch {
	case strings.HasPrefix(src, "data:"):
		return decodeCropDataURI(src)
	case strings.HasPrefix(src, "http://"), strings.HasPrefix(src, "https://"):
		return c.fetch(ctx, src)
	default:
		return c.readFile(src)
	}
}

// decodeCropDataURI decodes a "data:<mime>;base64,<payload>" source.
func decodeCropDataURI(src string) ([]byte, error) {
	comma := strings.IndexByte(src, ',')
	if comma < 0 || !strings.Contains(src[:comma], ";base64") {
		return nil, fmt.Errorf("io: image-crop: unsupported data URI (expected base64)")
	}
	data, err := base64.StdEncoding.DecodeString(src[comma+1:])
	if err != nil {
		return nil, fmt.Errorf("io: image-crop: decode data URI: %w", err)
	}
	if len(data) > maxCropSourceBytes {
		return nil, fmt.Errorf("io: image-crop: source exceeds the %d-byte limit", maxCropSourceBytes)
	}
	return data, nil
}

// readFile reads a path through the file tool's sandbox so cropping can never
// reach a file the file tool could not read (裁决 13.4.4：文件沙箱唯一实现处).
func (c *cropTool) readFile(p string) ([]byte, error) {
	if c.file == nil {
		return nil, fmt.Errorf("io: image-crop: no sandbox configured")
	}
	real, err := c.file.resolveExisting(p)
	if err != nil {
		return nil, err
	}
	fh, err := os.Open(real)
	if err != nil {
		return nil, fmt.Errorf("io: image-crop: open: %w", err)
	}
	defer fh.Close()
	data, truncated, err := readLimited(fh, maxCropSourceBytes)
	if err != nil {
		return nil, fmt.Errorf("io: image-crop: read: %w", err)
	}
	if truncated {
		return nil, fmt.Errorf("io: image-crop: source exceeds the %d-byte limit", maxCropSourceBytes)
	}
	return data, nil
}

// fetch downloads an http(s) source through the shared client, applying the same
// SSRF guard as the other network tools.
func (c *cropTool) fetch(ctx context.Context, rawURL string) ([]byte, error) {
	if c.client == nil {
		return nil, fmt.Errorf("io: image-crop: no http client configured")
	}
	if _, err := validateRequestURL(ctx, rawURL, c.allowedHosts); err != nil {
		return nil, err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
	if err != nil {
		return nil, fmt.Errorf("io: image-crop: build request: %w", err)
	}
	resp, err := c.client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("io: image-crop: fetch: %w", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("io: image-crop: fetch: status %d", resp.StatusCode)
	}
	data, truncated, err := readLimited(resp.Body, maxCropSourceBytes)
	if err != nil {
		return nil, fmt.Errorf("io: image-crop: read response: %w", err)
	}
	if truncated {
		return nil, fmt.Errorf("io: image-crop: source exceeds the %d-byte limit", maxCropSourceBytes)
	}
	return data, nil
}