ituoke/internal/agent/screenshots_test.go

package agent

import (
	"bytes"
	"context"
	"encoding/json"
	"image"
	"image/png"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"

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

func TestScreenshotShortcutPersistsImagesWithoutModelOrDirectoryCalls(t *testing.T) {
	var picture bytes.Buffer
	png.Encode(&picture, image.NewRGBA(image.Rect(0, 0, 20, 40)))
	var calls atomic.Int32
	phone := mcp.NewServer(&mcp.Implementation{Name: "phone", Version: "1"}, nil)
	mcp.AddTool(phone, &mcp.Tool{Name: "mobile_list_available_devices"}, func(context.Context, *mcp.CallToolRequest, struct{}) (*mcp.CallToolResult, any, error) {
		calls.Add(1)
		return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: `{"devices":[{"id":"phone-1","state":"online"}]}`}}}, nil, nil
	})
	mcp.AddTool(phone, &mcp.Tool{Name: "mobile_take_screenshot"}, func(_ context.Context, _ *mcp.CallToolRequest, args struct {
		Device  string `json:"device"`
		MaxSize int    `json:"maxSize"`
	}) (*mcp.CallToolResult, any, error) { calls.Add(1); if args.Device != "phone-1" {
		t.Error("wrong device")
	}; return &mcp.CallToolResult{Content: []mcp.Content{&mcp.ImageContent{MIMEType: "image/png", Data: picture.Bytes()}}}, nil, nil })
	remote := httptest.NewServer(mcp.NewStreamableHTTPHandler(func(*http.Request) *mcp.Server { return phone }, nil))
	defer remote.Close()
	t.Setenv("APPDATA", t.TempDir())
	t.Setenv("XDG_CONFIG_HOME", t.TempDir())
	t.Setenv("HOME", t.TempDir())
	mcpSvc := mcpclient.NewService()
	mcpSvc.SaveServer(mcpclient.ServerInput{Config: mcpclient.ServerConfig{Name: "Phone", Transport: "http", URL: remote.URL, Enabled: true, TimeoutSeconds: 3}})
	mcpSvc.SaveServer(mcpclient.ServerInput{Config: mcpclient.ServerConfig{Name: "Unneeded filesystem", Transport: "filesystem", Args: []string{t.TempDir()}, Enabled: true, TimeoutSeconds: 3}})
	var modelRequests atomic.Int32
	model := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		modelRequests.Add(1)
		t.Error("screenshot called model")
		http.Error(w, "unexpected", 500)
	}))
	defer model.Close()
	s, secrets := testService(t)
	s.mcp = mcpSvc
	s.SaveModelConfig(ModelConfigInput{Config: testConfig(model.URL), APIKey: "key"})
	c := memoryConversation(t, s, 4) // Would trigger summary for a normal request.
	var trace []ToolExecution
	s.emitDelta = func(event ChatDelta) {
		if event.Tool != nil && event.Tool.Status == "complete" {
			trace = append(trace, *event.Tool)
		}
	}
	result, err := s.ChatStream(context.Background(), "shot", c.ID, "截图给我看看")
	if err != nil || !strings.Contains(result["content"], "截图已显示") || modelRequests.Load() != 0 || calls.Load() != 2 {
		t.Fatalf("shortcut: %v %v model=%d tools=%d", result, err, modelRequests.Load(), calls.Load())
	}
	if len(trace) != 2 || len(trace[1].Images) != 1 {
		t.Fatalf("missing image event: %+v", trace)
	}
	a := trace[1].Images[0]
	preview, err := s.GetAttachmentPreview(c.ID, a.ID)
	if err != nil || !strings.HasPrefix(preview, "data:image/png;base64,") {
		t.Fatal("preview unavailable", err)
	}
	c.Messages = append(c.Messages, ConversationMessage{ID: "shot-user", Role: "user", Content: "截图"}, ConversationMessage{ID: "shot-answer", Role: "assistant", Content: result["content"], Status: "complete", Tools: trace})
	if _, err := s.SaveConversation(c); err != nil {
		t.Fatal(err)
	}
	restarted := newService(s.configPath, secrets)
	stored, err := restarted.GetConversation(c.ID)
	if err != nil || len(stored.Messages[len(stored.Messages)-1].Tools[1].Images) != 1 {
		t.Fatal("image metadata lost on restart", err)
	}
	if restored, err := restarted.GetAttachmentPreview(c.ID, a.ID); err != nil || restored != preview {
		t.Fatal("persisted preview lost", err)
	}
	other, _ := s.CreateConversation()
	if _, err := s.GetAttachmentPreview(other.ID, a.ID); err == nil {
		t.Fatal("cross-conversation image accessible")
	}
}

func TestScreenshotIntentDoesNotSkipOtherTasks(t *testing.T) {
	for _, text := range []string{"截图", "请帮我截个图给我看看", "手机当前画面截图发给我", "take a screenshot"} {
		if !isScreenshotRequest(text) {
			t.Errorf("simple screenshot missed: %s", text)
		}
	}
	for _, text := range []string{"不要截图", "打开设置然后截图", "截图并分析哪里报错", "将截图保存到 D:/image.png", "别截图，回到桌面"} {
		if isScreenshotRequest(text) {
			t.Errorf("other task discarded: %s", text)
		}
	}
}

func TestQwenToolThinkingCanBeEnabledButDefaultsToFast(t *testing.T) {
	for _, thinking := range []bool{false, true} {
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			var body map[string]any
			json.NewDecoder(r.Body).Decode(&body)
			if body["enable_thinking"] != thinking {
				t.Errorf("thinking setting: %v", body["enable_thinking"])
			}
			memoryReply(w, "done")
		}))
		c := testConfig(server.URL)
		c.Provider = "qwen"
		c.Model = "qwen3.8-max-0902"
		c.ToolThinking = thinking
		_, err := modelToolStep(context.Background(), c, "", nil, []map[string]any{{"type": "function", "function": map[string]any{"name": "test", "parameters": map[string]any{"type": "object"}}}}, nil)
		server.Close()
		if err != nil {
			t.Fatal(err)
		}
	}
}