206 lines
8.3 KiB
Go
206 lines
8.3 KiB
Go
package agent
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"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 TestOperationScreenshotsStayInternalUnlessRequested(t *testing.T) {
|
|
for _, tc := range []struct {
|
|
message string
|
|
visible bool
|
|
}{
|
|
{"打开设置", false},
|
|
{"打开设置然后截图给我看", true},
|
|
{"不要显示截图，打开设置", false},
|
|
} {
|
|
t.Run(tc.message, func(t *testing.T) {
|
|
var picture bytes.Buffer
|
|
png.Encode(&picture, image.NewRGBA(image.Rect(0, 0, 20, 40)))
|
|
tools, _ := toolTestService(t, func(context.Context, *mcp.CallToolRequest, map[string]any) (*mcp.CallToolResult, map[string]any, error) {
|
|
return &mcp.CallToolResult{Content: []mcp.Content{&mcp.ImageContent{MIMEType: "image/png", Data: picture.Bytes()}}}, nil, nil
|
|
}, "mobile_take_screenshot")
|
|
var rounds atomic.Int32
|
|
model := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
var body struct {
|
|
Tools []struct{ Function struct{ Name string } }
|
|
Messages []modelMessage
|
|
}
|
|
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
|
t.Error(err)
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
if rounds.Add(1) == 1 {
|
|
fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"capture\",\"type\":\"function\",\"function\":{\"name\":%q,\"arguments\":\"{}\"}}]},\"finish_reason\":\"tool_calls\"}]}\n\ndata: [DONE]\n\n", body.Tools[0].Function.Name)
|
|
} else {
|
|
raw, _ := json.Marshal(body.Messages)
|
|
if !bytes.Contains(raw, []byte("data:image/png;base64,")) {
|
|
t.Error("internal screenshot was not sent to model")
|
|
}
|
|
fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"操作完成\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n")
|
|
}
|
|
}))
|
|
defer model.Close()
|
|
s, _ := testService(t)
|
|
s.mcp = tools
|
|
if _, err := s.SaveModelConfig(ModelConfigInput{Config: testConfig(model.URL), APIKey: "key"}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
conversation, _ := s.CreateConversation()
|
|
var visibleImages int
|
|
s.emitDelta = func(delta ChatDelta) {
|
|
if delta.Tool != nil {
|
|
visibleImages += len(delta.Tool.Images)
|
|
}
|
|
}
|
|
if _, err := s.ChatStream(context.Background(), "operation-shot", conversation.ID, tc.message); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if (visibleImages > 0) != tc.visible {
|
|
t.Fatalf("visible images=%d, expected visible=%t", visibleImages, tc.visible)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestScreenshotDisplayIntentUsesCurrentExplicitRequest(t *testing.T) {
|
|
for _, text := range []string{"截图", "打开微信然后截个图给我", "show a screenshot", "显示手机画面"} {
|
|
if !wantsScreenshotDisplay(text) {
|
|
t.Errorf("missed explicit request: %s", text)
|
|
}
|
|
}
|
|
for _, text := range []string{"打开设置", "继续", "别截图，回到桌面", "截图不要显示", "don't show a screenshot", "我没让显示截图，对话框总显示截图"} {
|
|
if wantsScreenshotDisplay(text) {
|
|
t.Errorf("unexpected image display: %s", text)
|
|
}
|
|
}
|
|
}
|
|
|
|
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()
|
|
if err := mcpSvc.SetEnabled("builtin-filesystem", false); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := mcpSvc.SetEnabled("builtin-mobile-next", false); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
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)
|
|
}
|
|
}
|
|
}
|