201 lines
8.7 KiB
Go
201 lines
8.7 KiB
Go
package agent
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
|
|
"github.com/ituoke/ituoke/internal/mcpclient"
|
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
|
)
|
|
|
|
func toolTestService(t *testing.T, handler func(context.Context, *mcp.CallToolRequest, map[string]any) (*mcp.CallToolResult, map[string]any, error)) (*mcpclient.Service, string) {
|
|
t.Helper()
|
|
server := mcp.NewServer(&mcp.Implementation{Name: "phone-test", Version: "1"}, nil)
|
|
mcp.AddTool(server, &mcp.Tool{Name: "phone_action", Description: "Test phone action"}, handler)
|
|
httpServer := httptest.NewServer(mcp.NewStreamableHTTPHandler(func(*http.Request) *mcp.Server { return server }, nil))
|
|
t.Cleanup(httpServer.Close)
|
|
// Keep test configuration isolated without touching the user's service file.
|
|
t.Setenv("APPDATA", t.TempDir())
|
|
t.Setenv("XDG_CONFIG_HOME", t.TempDir())
|
|
t.Setenv("HOME", t.TempDir())
|
|
service := mcpclient.NewService()
|
|
view, err := service.SaveServer(mcpclient.ServerInput{Config: mcpclient.ServerConfig{Name: "Phone", Enabled: true, Transport: "http", URL: httpServer.URL, TimeoutSeconds: 3}})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return service, view.Config.ID
|
|
}
|
|
|
|
func TestAgentStreamsToolCallsAndFeedsResultsBack(t *testing.T) {
|
|
var executions atomic.Int32
|
|
mcpService, _ := toolTestService(t, func(ctx context.Context, _ *mcp.CallToolRequest, args map[string]any) (*mcp.CallToolResult, map[string]any, error) {
|
|
executions.Add(1)
|
|
if args["action"] != "open_settings" {
|
|
t.Errorf("arguments lost: %v", args)
|
|
}
|
|
return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: "settings_visible"}}}, nil, nil
|
|
})
|
|
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 `json:"name"`
|
|
} `json:"function"`
|
|
} `json:"tools"`
|
|
Messages []modelMessage `json:"messages"`
|
|
}
|
|
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 {
|
|
if len(body.Tools) != 1 {
|
|
t.Errorf("missing tools: %+v", body.Tools)
|
|
return
|
|
}
|
|
name := body.Tools[0].Function.Name
|
|
fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"先打开设置。\"}}]}\n\n")
|
|
fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call-1\",\"type\":\"function\",\"function\":{\"name\":%q,\"arguments\":\"{\\\"action\\\":\"}}]}}]}\n\n", name)
|
|
fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"\\\"open_settings\\\"}\"}}]},\"finish_reason\":\"tool_calls\"}]}\n\ndata: [DONE]\n\n")
|
|
} else {
|
|
last := body.Messages[len(body.Messages)-1]
|
|
previous := body.Messages[len(body.Messages)-2]
|
|
if previous.ReasoningContent != "先打开设置。" {
|
|
t.Errorf("tool-round reasoning was not passed back: %+v", previous)
|
|
}
|
|
if last.Role != "tool" || last.ToolCallID != "call-1" || !strings.Contains(fmt.Sprint(last.Content), "settings_visible") {
|
|
t.Errorf("tool result not returned: %+v", last)
|
|
}
|
|
fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"reasoning_content\":\"页面验证完成。\",\"content\":\"设置已打开并验证。\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n")
|
|
}
|
|
}))
|
|
defer model.Close()
|
|
s, _ := testService(t)
|
|
s.mcp = mcpService
|
|
if _, err := s.SaveModelConfig(ModelConfigInput{Config: testConfig(model.URL), APIKey: "test-key"}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var events []ToolExecution
|
|
var text strings.Builder
|
|
var thought strings.Builder
|
|
s.emitDelta = func(delta ChatDelta) {
|
|
if delta.Tool != nil {
|
|
events = append(events, *delta.Tool)
|
|
}
|
|
text.WriteString(delta.Delta)
|
|
thought.WriteString(delta.ReasoningDelta)
|
|
}
|
|
reply, err := s.ChatStream(context.Background(), "tool-test", "chat", "打开系统设置")
|
|
if err != nil || reply["content"] != "设置已打开并验证。" || text.String() != reply["content"] {
|
|
t.Fatalf("reply: %v %v", reply, err)
|
|
}
|
|
if executions.Load() != 1 || len(events) != 2 || events[0].Status != "running" || events[1].Status != "complete" {
|
|
t.Fatalf("execution: %d %+v", executions.Load(), events)
|
|
}
|
|
if reply["reasoning"] != "先打开设置。\n\n页面验证完成。" || thought.String() != reply["reasoning"] {
|
|
t.Fatalf("tool-round reasoning lost: %q %v", thought.String(), reply)
|
|
}
|
|
}
|
|
|
|
func TestAgentCancellationStopsLaterTools(t *testing.T) {
|
|
var calls atomic.Int32
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
mcpService, _ := toolTestService(t, func(context.Context, *mcp.CallToolRequest, map[string]any) (*mcp.CallToolResult, map[string]any, error) {
|
|
calls.Add(1)
|
|
return &mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: "done"}}}, nil, nil
|
|
})
|
|
model := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
var body struct {
|
|
Tools []struct {
|
|
Function struct {
|
|
Name string `json:"name"`
|
|
} `json:"function"`
|
|
} `json:"tools"`
|
|
}
|
|
json.NewDecoder(r.Body).Decode(&body)
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"one\",\"type\":\"function\",\"function\":{\"name\":%q,\"arguments\":\"{}\"}},{\"index\":1,\"id\":\"two\",\"type\":\"function\",\"function\":{\"name\":%q,\"arguments\":\"{}\"}}]},\"finish_reason\":\"tool_calls\"}]}\n\ndata: [DONE]\n\n", body.Tools[0].Function.Name, body.Tools[0].Function.Name)
|
|
}))
|
|
defer model.Close()
|
|
s, _ := testService(t)
|
|
s.mcp = mcpService
|
|
s.SaveModelConfig(ModelConfigInput{Config: testConfig(model.URL), APIKey: "key"})
|
|
s.emitDelta = func(delta ChatDelta) {
|
|
if delta.Tool != nil && delta.Tool.Status == "complete" {
|
|
cancel()
|
|
}
|
|
}
|
|
_, err := s.ChatStream(ctx, "cancel-test", "chat", "test")
|
|
if !errors.Is(err, context.Canceled) || calls.Load() != 1 || len(s.history["chat"]) != 0 {
|
|
t.Fatalf("cancellation: %d %v", calls.Load(), err)
|
|
}
|
|
}
|
|
|
|
func TestIncompleteToolCallDoesNotExecute(t *testing.T) {
|
|
for _, stream := range []string{
|
|
`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"x","type":"function","function":{"name":"action","arguments":"{"}}]},"finish_reason":"length"}]}` + "\n\n",
|
|
`data: {"choices":[{"delta":{"tool_calls":[{"index":1,"id":"x","type":"function","function":{"name":"action","arguments":"{}"}}]},"finish_reason":"tool_calls"}]}` + "\n\n",
|
|
} {
|
|
if _, err := readToolStream(strings.NewReader(stream), "", nil); err == nil {
|
|
t.Fatal("accepted incomplete tool call")
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestAgentNonStreamReturnsToolFailureToModel(t *testing.T) {
|
|
mcpService, _ := toolTestService(t, func(context.Context, *mcp.CallToolRequest, map[string]any) (*mcp.CallToolResult, map[string]any, error) {
|
|
return &mcp.CallToolResult{IsError: true, Content: []mcp.Content{&mcp.TextContent{Text: "device disconnected"}}}, nil, nil
|
|
})
|
|
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 `json:"name"`
|
|
} `json:"function"`
|
|
} `json:"tools"`
|
|
Messages []modelMessage `json:"messages"`
|
|
}
|
|
json.NewDecoder(r.Body).Decode(&body)
|
|
w.Header().Set("Content-Type", "application/json")
|
|
if rounds.Add(1) == 1 {
|
|
fmt.Fprintf(w, `{"choices":[{"message":{"content":null,"tool_calls":[{"id":"one","type":"function","function":{"name":%q,"arguments":"{}"}}]},"finish_reason":"tool_calls"}]}`, body.Tools[0].Function.Name)
|
|
} else {
|
|
last := body.Messages[len(body.Messages)-1]
|
|
if last.Role != "tool" || !strings.Contains(fmt.Sprint(last.Content), "device disconnected") || !strings.Contains(fmt.Sprint(last.Content), "失败") {
|
|
t.Errorf("tool failure lost: %+v", last)
|
|
}
|
|
fmt.Fprint(w, `{"choices":[{"message":{"content":"设备断开，操作没有完成。"},"finish_reason":"stop"}]}`)
|
|
}
|
|
}))
|
|
defer model.Close()
|
|
s, _ := testService(t)
|
|
s.mcp = mcpService
|
|
s.SaveModelConfig(ModelConfigInput{Config: testConfig(model.URL), APIKey: "key"})
|
|
result, err := s.Chat("test", "test")
|
|
if err != nil || result["content"] != "设备断开，操作没有完成。" {
|
|
t.Fatalf("non-stream recovery: %v %v", result, err)
|
|
}
|
|
}
|
|
|
|
func TestToolImagesStayOutOfTextAndHistory(t *testing.T) {
|
|
text, images := toolOutput(&mcp.CallToolResult{Content: []mcp.Content{&mcp.TextContent{Text: "screen"}, &mcp.ImageContent{MIMEType: "image/png", Data: []byte("image-bytes")}}})
|
|
if strings.Contains(text, "image-bytes") || len(images) != 1 || !strings.Contains(text, "1 张图片") {
|
|
t.Fatalf("image adapter: %s %v", text, images)
|
|
}
|
|
encoded, _ := json.Marshal(images)
|
|
if !strings.Contains(string(encoded), "data:image/png;base64,") {
|
|
t.Fatal("missing model image input")
|
|
}
|
|
}
|