ituoke/internal/agent/screenshots.go

package agent

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"regexp"
	"strings"
	"time"

	"github.com/ituoke/ituoke/internal/mcpclient"
	"github.com/modelcontextprotocol/go-sdk/mcp"
)

var screenshotMention = regexp.MustCompile(`(?i)截图|截屏|截(?:一张|一下|个|张)图|screen\s?shot|显示.{0,6}(?:手机|屏幕|画面)|(?:手机|屏幕)画面.{0,6}(?:发给我|给我看)`)
var screenshotRefusal = regexp.MustCompile(`(?i)(?:不要|不用|无需|别|不需要|不想|没让|没有让).{0,12}(?:截图|截屏|截.{0,3}图|显示|展示)|(?:截图|截屏).{0,8}(?:不要|不用|无需|别).{0,6}(?:显示|展示|发)|(?:don't|do not|without|no need).{0,24}(?:screen\s?shot|show)`)

// Only the current user request may opt screenshots into the visible transcript.
// Tool calls and earlier turns cannot turn internal observations into attachments.
func wantsScreenshotDisplay(text string) bool {
	return screenshotMention.MatchString(text) && !screenshotRefusal.MatchString(text)
}

func (s *Service) saveToolImages(ctx context.Context, sessionID string, result *mcp.CallToolResult) ([]Attachment, error) {
	var files []Attachment
	for _, part := range result.Content {
		img, ok := part.(*mcp.ImageContent)
		if !ok || len(files) >= 2 {
			continue
		}
		ext := map[string]string{"image/png": ".png", "image/jpeg": ".jpg", "image/webp": ".webp"}[img.MIMEType]
		if ext == "" || len(img.Data) == 0 || len(img.Data) > 4*1024*1024 {
			continue
		}
		file, err := s.UploadAttachment(ctx, sessionID, AttachmentInput{Name: fmt.Sprintf("截图-%s-%d%s", time.Now().Format("20060102-150405"), len(files)+1, ext), Data: base64.StdEncoding.EncodeToString(img.Data)})
		if err != nil {
			return files, err
		}
		files = append(files, file)
	}
	return files, nil
}

func isScreenshotRequest(text string) bool {
	text = strings.ToLower(strings.TrimSpace(text))
	if text == "screenshot" || text == "take a screenshot" || text == "show a screenshot" {
		return true
	}
	text = strings.NewReplacer("截一张图", "截图", "截一下图", "截图", "截个图", "截图", "截张图", "截图").Replace(text)
	text = strings.NewReplacer("帮我", "", "麻烦", "", "请", "", "发给我", "", "发我", "", "给我", "", "看一下", "", "看看", "", "看下", "", "看", "", "当前", "", "手机", "", "屏幕", "", "画面", "", "一张", "", "一下", "", " ", "", "\n", "", "，", "", "。", "", "！", "", "!", "", "？", "", "?", "").Replace(text)
	return text == "截图" || text == "截屏"
}

// A plain screenshot request needs no language-model decision or file-manager
// round trip. It uses the same MCP action lease and durable image pipeline.
func (s *Service) quickScreenshot(ctx context.Context, sessionID string, onTool func(ToolExecution)) (string, bool, error) {
	set, err := mcpclient.OpenTools(ctx, s.mcp, nil, func(c mcpclient.ServerConfig) bool { return c.Transport != "filesystem" })
	if err != nil {
		return "", true, err
	}
	defer set.Close()
	var list, capture mcpclient.AgentTool
	for _, tool := range set.Tools() {
		if tool.Tool == "mobile_list_available_devices" {
			list = tool
		}
		if tool.Tool == "mobile_take_screenshot" {
			capture = tool
		}
	}
	if list.Name == "" || capture.Name == "" {
		return "", false, nil
	}
	call := func(tool mcpclient.AgentTool, args string) (*mcp.CallToolResult, error) {
		event := ToolExecution{ID: "quick-" + tool.Tool, Server: tool.Server, Tool: tool.Tool, Status: "running", Arguments: args}
		if onTool != nil {
			onTool(event)
		}
		start := time.Now()
		result, err := set.Call(ctx, tool.Name, args)
		event.Status = "complete"
		if err != nil {
			event.Status = "error"
			event.Result = err.Error()
		} else {
			text, _ := toolOutput(result)
			event.Result = clipBytes(text, 4000)
			if result.IsError {
				event.Status = "error"
				err = fmt.Errorf("%s", clipBytes(text, 1000))
			} else if tool.Tool == "mobile_take_screenshot" {
				event.Images, err = s.saveToolImages(ctx, sessionID, result)
				if err == nil && len(event.Images) == 0 {
					err = fmt.Errorf("截图工具未返回有效图片")
				}
				if err != nil {
					event.Status = "error"
					event.Result = err.Error()
				} else {
					event.Result = "截图已保存并显示在会话中。"
				}
			}
		}
		event.ElapsedMs = time.Since(start).Milliseconds()
		if onTool != nil {
			onTool(event)
		}
		return result, err
	}
	result, err := call(list, "{}")
	if err != nil {
		return "", true, err
	}
	var response struct {
		Devices []struct {
			ID    string `json:"id"`
			State string `json:"state"`
		} `json:"devices"`
	}
	for _, content := range result.Content {
		if text, ok := content.(*mcp.TextContent); ok {
			if json.Unmarshal([]byte(text.Text), &response) == nil && len(response.Devices) > 0 {
				break
			}
		}
	}
	var devices []string
	for _, device := range response.Devices {
		if device.ID != "" && (device.State == "online" || device.State == "") {
			devices = append(devices, device.ID)
		}
	}
	if len(devices) != 1 {
		return "", true, fmt.Errorf("检测到 %d 台在线设备，无法唯一确定截图目标，请明确设备后再试", len(devices))
	}
	if err := ctx.Err(); err != nil {
		return "", true, err
	}
	args, _ := json.Marshal(map[string]any{"device": devices[0], "maxSize": 1600})
	if _, err = call(capture, string(args)); err != nil {
		return "", true, err
	}
	return "截图已显示，点击图片可放大查看。", true, nil
}