236 lines
6.9 KiB
Go
236 lines
6.9 KiB
Go
package tools
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"reflect"
|
|
"sort"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
|
|
|
"isphere-ai-bridge/internal/helperclient"
|
|
)
|
|
|
|
var expectedWinHelperToolNames = []string{
|
|
"win_helper_version",
|
|
"win_helper_self_check",
|
|
"win_helper_scan_windows",
|
|
"win_helper_dump_uia",
|
|
}
|
|
|
|
func TestWinHelperToolsRegisterExactlyFourReadOnlyTools(t *testing.T) {
|
|
fake := &fakeHelperCaller{}
|
|
session, cleanup := connectToolsTestSession(t, func(server *mcp.Server) {
|
|
RegisterWinHelperTools(server, fake)
|
|
})
|
|
defer cleanup()
|
|
|
|
result, err := session.ListTools(context.Background(), &mcp.ListToolsParams{})
|
|
if err != nil {
|
|
t.Fatalf("list tools: %v", err)
|
|
}
|
|
|
|
gotNames := make([]string, 0, len(result.Tools))
|
|
for _, tool := range result.Tools {
|
|
gotNames = append(gotNames, tool.Name)
|
|
lower := strings.ToLower(tool.Name)
|
|
forbidden := []string{"send", "search", "file", "login", "conversation", "upload", "download", "receive"}
|
|
for _, word := range forbidden {
|
|
if strings.Contains(lower, word) {
|
|
t.Fatalf("tool %q contains forbidden action word %q", tool.Name, word)
|
|
}
|
|
}
|
|
}
|
|
sort.Strings(gotNames)
|
|
wantNames := append([]string(nil), expectedWinHelperToolNames...)
|
|
sort.Strings(wantNames)
|
|
if !reflect.DeepEqual(gotNames, wantNames) {
|
|
t.Fatalf("registered tool names = %#v, want %#v", gotNames, wantNames)
|
|
}
|
|
}
|
|
|
|
func TestWinHelperToolSchemasExposeOnlyAllowedParameters(t *testing.T) {
|
|
assertJSONFields[VersionArgs](t)
|
|
assertJSONFields[SelfCheckArgs](t)
|
|
assertJSONFields[ScanWindowsArgs](t, "include_all_visible")
|
|
assertJSONFields[DumpUIAArgs](t, "hwnd", "max_depth", "include_text", "max_children")
|
|
|
|
fake := &fakeHelperCaller{}
|
|
session, cleanup := connectToolsTestSession(t, func(server *mcp.Server) {
|
|
RegisterWinHelperTools(server, fake)
|
|
})
|
|
defer cleanup()
|
|
|
|
result, err := session.ListTools(context.Background(), &mcp.ListToolsParams{})
|
|
if err != nil {
|
|
t.Fatalf("list tools: %v", err)
|
|
}
|
|
schemas := map[string][]string{}
|
|
for _, tool := range result.Tools {
|
|
schemas[tool.Name] = schemaPropertyNames(t, tool.InputSchema)
|
|
}
|
|
|
|
assertSameStrings(t, schemas["win_helper_version"], nil)
|
|
assertSameStrings(t, schemas["win_helper_self_check"], nil)
|
|
assertSameStrings(t, schemas["win_helper_scan_windows"], []string{"include_all_visible"})
|
|
assertSameStrings(t, schemas["win_helper_dump_uia"], []string{"hwnd", "max_depth", "include_text", "max_children"})
|
|
}
|
|
|
|
func TestWinHelperHandlersCallHelperClientAbstractionWithCorrectOps(t *testing.T) {
|
|
fake := &fakeHelperCaller{}
|
|
session, cleanup := connectToolsTestSession(t, func(server *mcp.Server) {
|
|
RegisterWinHelperTools(server, fake)
|
|
})
|
|
defer cleanup()
|
|
|
|
calls := []struct {
|
|
tool string
|
|
args map[string]any
|
|
}{
|
|
{tool: "win_helper_version", args: map[string]any{}},
|
|
{tool: "win_helper_self_check", args: map[string]any{}},
|
|
{tool: "win_helper_scan_windows", args: map[string]any{"include_all_visible": true}},
|
|
{tool: "win_helper_dump_uia", args: map[string]any{"hwnd": "0x001A0B2C", "max_depth": 3, "include_text": true, "max_children": 25}},
|
|
}
|
|
|
|
for _, call := range calls {
|
|
if _, err := session.CallTool(context.Background(), &mcp.CallToolParams{Name: call.tool, Arguments: call.args}); err != nil {
|
|
t.Fatalf("call %s: %v", call.tool, err)
|
|
}
|
|
}
|
|
|
|
wantOps := []string{"version", "self_check", "scan_windows", "dump_uia"}
|
|
if len(fake.calls) != len(wantOps) {
|
|
t.Fatalf("helper calls = %#v, want %d calls", fake.calls, len(wantOps))
|
|
}
|
|
for i, wantOp := range wantOps {
|
|
if fake.calls[i].op != wantOp {
|
|
t.Fatalf("call %d op = %q, want %q", i, fake.calls[i].op, wantOp)
|
|
}
|
|
}
|
|
assertSameMap(t, fake.calls[0].args, map[string]any{})
|
|
assertSameMap(t, fake.calls[1].args, map[string]any{})
|
|
assertSameMap(t, fake.calls[2].args, map[string]any{"include_all_visible": true})
|
|
assertSameMap(t, fake.calls[3].args, map[string]any{"hwnd": "0x001A0B2C", "max_depth": float64(3), "include_text": true, "max_children": float64(25)})
|
|
}
|
|
|
|
type fakeHelperCaller struct {
|
|
calls []helperCall
|
|
}
|
|
|
|
type helperCall struct {
|
|
op string
|
|
args map[string]any
|
|
}
|
|
|
|
func (f *fakeHelperCaller) Call(_ context.Context, op string, args map[string]any) (*helperclient.HelperResponse, error) {
|
|
copied := map[string]any{}
|
|
for k, v := range args {
|
|
copied[k] = v
|
|
}
|
|
f.calls = append(f.calls, helperCall{op: op, args: copied})
|
|
return &helperclient.HelperResponse{
|
|
Protocol: helperclient.HelperProtocolV1,
|
|
RequestID: "fake-request",
|
|
Op: op,
|
|
OK: true,
|
|
Data: map[string]json.RawMessage{
|
|
"helper_name": json.RawMessage(`"ISphereWinHelper"`),
|
|
},
|
|
Warnings: []string{},
|
|
}, nil
|
|
}
|
|
|
|
func connectToolsTestSession(t *testing.T, register func(*mcp.Server)) (*mcp.ClientSession, func()) {
|
|
t.Helper()
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
serverTransport, clientTransport := mcp.NewInMemoryTransports()
|
|
server := mcp.NewServer(&mcp.Implementation{Name: "tools-test-server", Version: "0.0.0"}, nil)
|
|
register(server)
|
|
|
|
errCh := make(chan error, 1)
|
|
go func() {
|
|
errCh <- server.Run(ctx, serverTransport)
|
|
}()
|
|
|
|
client := mcp.NewClient(&mcp.Implementation{Name: "tools-test-client", Version: "0.0.0"}, nil)
|
|
session, err := client.Connect(ctx, clientTransport, nil)
|
|
if err != nil {
|
|
cancel()
|
|
t.Fatalf("connect client: %v", err)
|
|
}
|
|
|
|
cleanup := func() {
|
|
_ = session.Close()
|
|
cancel()
|
|
select {
|
|
case <-errCh:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatalf("server did not stop")
|
|
}
|
|
}
|
|
return session, cleanup
|
|
}
|
|
|
|
func assertJSONFields[T any](t *testing.T, want ...string) {
|
|
t.Helper()
|
|
var zero T
|
|
typeOf := reflect.TypeOf(zero)
|
|
if typeOf.Kind() != reflect.Struct {
|
|
t.Fatalf("%T is %s, want struct", zero, typeOf.Kind())
|
|
}
|
|
got := make([]string, 0, typeOf.NumField())
|
|
for i := 0; i < typeOf.NumField(); i++ {
|
|
field := typeOf.Field(i)
|
|
name := strings.Split(field.Tag.Get("json"), ",")[0]
|
|
if name == "" || name == "-" {
|
|
continue
|
|
}
|
|
got = append(got, name)
|
|
}
|
|
assertSameStrings(t, got, want)
|
|
}
|
|
|
|
func schemaPropertyNames(t *testing.T, schema any) []string {
|
|
t.Helper()
|
|
payload, err := json.Marshal(schema)
|
|
if err != nil {
|
|
t.Fatalf("marshal schema: %v", err)
|
|
}
|
|
var decoded map[string]any
|
|
if err := json.Unmarshal(payload, &decoded); err != nil {
|
|
t.Fatalf("unmarshal schema %s: %v", payload, err)
|
|
}
|
|
properties, _ := decoded["properties"].(map[string]any)
|
|
if len(properties) == 0 {
|
|
return nil
|
|
}
|
|
names := make([]string, 0, len(properties))
|
|
for name := range properties {
|
|
names = append(names, name)
|
|
}
|
|
sort.Strings(names)
|
|
return names
|
|
}
|
|
|
|
func assertSameStrings(t *testing.T, got []string, want []string) {
|
|
t.Helper()
|
|
gotCopy := append([]string(nil), got...)
|
|
wantCopy := append([]string(nil), want...)
|
|
sort.Strings(gotCopy)
|
|
sort.Strings(wantCopy)
|
|
if !reflect.DeepEqual(gotCopy, wantCopy) {
|
|
t.Fatalf("strings = %#v, want %#v", gotCopy, wantCopy)
|
|
}
|
|
}
|
|
|
|
func assertSameMap(t *testing.T, got map[string]any, want map[string]any) {
|
|
t.Helper()
|
|
if !reflect.DeepEqual(got, want) {
|
|
t.Fatalf("map = %#v, want %#v", got, want)
|
|
}
|
|
}
|