feat: add go helper client and mcp tools
This commit is contained in:
zhaoyilun
122
internal/tools/winhelper.go
Normal file
122
internal/tools/winhelper.go
Normal file
@@ -0,0 +1,122 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
|
||||
"isphere-ai-bridge/internal/helperclient"
|
||||
)
|
||||
|
||||
const (
|
||||
ToolNameVersion = "win_helper_version"
|
||||
ToolNameSelfCheck = "win_helper_self_check"
|
||||
ToolNameScanWindows = "win_helper_scan_windows"
|
||||
ToolNameDumpUIA = "win_helper_dump_uia"
|
||||
)
|
||||
|
||||
type HelperCaller interface {
|
||||
Call(ctx context.Context, op string, args map[string]any) (*helperclient.HelperResponse, error)
|
||||
}
|
||||
|
||||
type VersionArgs struct{}
|
||||
|
||||
type SelfCheckArgs struct{}
|
||||
|
||||
type ScanWindowsArgs struct {
|
||||
IncludeAllVisible bool `json:"include_all_visible,omitempty" jsonschema:"include all visible desktop windows, not only likely iSphere candidates"`
|
||||
}
|
||||
|
||||
type DumpUIAArgs struct {
|
||||
Hwnd string `json:"hwnd,omitempty" jsonschema:"window handle to dump, for example 0x001A0B2C"`
|
||||
MaxDepth int `json:"max_depth,omitempty" jsonschema:"maximum UI Automation tree depth"`
|
||||
IncludeText bool `json:"include_text,omitempty" jsonschema:"include text values when available"`
|
||||
MaxChildren int `json:"max_children,omitempty" jsonschema:"maximum children per node"`
|
||||
}
|
||||
|
||||
func RegisterWinHelperTools(server *mcp.Server, caller HelperCaller) {
|
||||
if caller == nil {
|
||||
caller = helperclient.Client{}
|
||||
}
|
||||
|
||||
mcp.AddTool[VersionArgs, map[string]any](server, &mcp.Tool{
|
||||
Name: ToolNameVersion,
|
||||
Description: "Return ISphereWinHelper version and protocol metadata.",
|
||||
}, func(ctx context.Context, _ *mcp.CallToolRequest, _ VersionArgs) (*mcp.CallToolResult, map[string]any, error) {
|
||||
return callHelper(ctx, caller, "version", map[string]any{})
|
||||
})
|
||||
|
||||
mcp.AddTool[SelfCheckArgs, map[string]any](server, &mcp.Tool{
|
||||
Name: ToolNameSelfCheck,
|
||||
Description: "Run ISphereWinHelper read-only desktop and UI Automation availability checks.",
|
||||
}, func(ctx context.Context, _ *mcp.CallToolRequest, _ SelfCheckArgs) (*mcp.CallToolResult, map[string]any, error) {
|
||||
return callHelper(ctx, caller, "self_check", map[string]any{})
|
||||
})
|
||||
|
||||
mcp.AddTool[ScanWindowsArgs, map[string]any](server, &mcp.Tool{
|
||||
Name: ToolNameScanWindows,
|
||||
Description: "List visible desktop windows or likely iSphere window candidates in read-only mode.",
|
||||
}, func(ctx context.Context, _ *mcp.CallToolRequest, input ScanWindowsArgs) (*mcp.CallToolResult, map[string]any, error) {
|
||||
return callHelper(ctx, caller, "scan_windows", map[string]any{
|
||||
"include_all_visible": input.IncludeAllVisible,
|
||||
})
|
||||
})
|
||||
|
||||
mcp.AddTool[DumpUIAArgs, map[string]any](server, &mcp.Tool{
|
||||
Name: ToolNameDumpUIA,
|
||||
Description: "Dump a read-only UI Automation tree for a specified window handle.",
|
||||
}, func(ctx context.Context, _ *mcp.CallToolRequest, input DumpUIAArgs) (*mcp.CallToolResult, map[string]any, error) {
|
||||
return callHelper(ctx, caller, "dump_uia", map[string]any{
|
||||
"hwnd": input.Hwnd,
|
||||
"max_depth": float64(input.MaxDepth),
|
||||
"include_text": input.IncludeText,
|
||||
"max_children": float64(input.MaxChildren),
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
func callHelper(ctx context.Context, caller HelperCaller, op string, args map[string]any) (*mcp.CallToolResult, map[string]any, error) {
|
||||
response, err := caller.Call(ctx, op, args)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
return nil, helperResponseToMap(response), nil
|
||||
}
|
||||
|
||||
func helperResponseToMap(response *helperclient.HelperResponse) map[string]any {
|
||||
if response == nil {
|
||||
return map[string]any{"ok": false}
|
||||
}
|
||||
|
||||
out := map[string]any{
|
||||
"protocol": response.Protocol,
|
||||
"request_id": response.RequestID,
|
||||
"op": response.Op,
|
||||
"ok": response.OK,
|
||||
"data": decodeRawData(response.Data),
|
||||
"warnings": response.Warnings,
|
||||
}
|
||||
if response.Error != nil {
|
||||
out["error"] = map[string]any{
|
||||
"code": response.Error.Code,
|
||||
"message": response.Error.Message,
|
||||
}
|
||||
} else {
|
||||
out["error"] = nil
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func decodeRawData(data map[string]json.RawMessage) map[string]any {
|
||||
decoded := map[string]any{}
|
||||
for key, raw := range data {
|
||||
var value any
|
||||
if err := json.Unmarshal(raw, &value); err != nil {
|
||||
decoded[key] = string(raw)
|
||||
continue
|
||||
}
|
||||
decoded[key] = value
|
||||
}
|
||||
return decoded
|
||||
}
|
||||
235
internal/tools/winhelper_test.go
Normal file
235
internal/tools/winhelper_test.go
Normal file
@@ -0,0 +1,235 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user