Cyrene_v2/internal/capability/mcp/helpers_test.go

package mcp

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"testing"
	"time"

	mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp"
)

// stdioHelperEnv 标记「测试二进制以 stdio MCP server 身份运行」。
const stdioHelperEnv = "CYRENE_MCP_STDIO_HELPER"

// echoToolSpec 返回入参 name（string，默认 "world"）的回显工具定义。
func echoToolSpec() *mcpsdk.Tool {
	return &mcpsdk.Tool{
		Name:        "echo",
		Description: "回显问候语",
		InputSchema: map[string]any{
			"type": "object",
			"properties": map[string]any{
				"name": map[string]any{"type": "string", "default": "world"},
			},
		},
	}
}

// addToolSpec 返回入参 a、b（均为必填 number）的加法工具定义。
func addToolSpec() *mcpsdk.Tool {
	return &mcpsdk.Tool{
		Name:        "add",
		Description: "两数相加",
		InputSchema: map[string]any{
			"type":     "object",
			"required": []string{"a", "b"},
			"properties": map[string]any{
				"a": map[string]any{"type": "number"},
				"b": map[string]any{"type": "number"},
			},
		},
	}
}

// echoHandler 按 schema 解参并回显；arguments 在 SDK 里是 json.RawMessage。
func echoHandler(_ context.Context, req *mcpsdk.CallToolRequest) (*mcpsdk.CallToolResult, error) {
	var in struct {
		Name string `json:"name"`
	}
	if len(req.Params.Arguments) > 0 {
		if err := json.Unmarshal(req.Params.Arguments, &in); err != nil {
			return nil, fmt.Errorf("decode echo args: %w", err)
		}
	}
	return &mcpsdk.CallToolResult{
		Content: []mcpsdk.Content{&mcpsdk.TextContent{Text: "hello " + in.Name}},
	}, nil
}

// addHandler 求和并同时返回文本与结构化结果。
func addHandler(_ context.Context, req *mcpsdk.CallToolRequest) (*mcpsdk.CallToolResult, error) {
	var in struct {
		A float64 `json:"a"`
		B float64 `json:"b"`
	}
	if err := json.Unmarshal(req.Params.Arguments, &in); err != nil {
		return nil, fmt.Errorf("decode add args: %w", err)
	}
	sum := in.A + in.B
	return &mcpsdk.CallToolResult{
		Content:           []mcpsdk.Content{&mcpsdk.TextContent{Text: fmt.Sprintf("%g", sum)}},
		StructuredContent: map[string]any{"sum": sum},
	}, nil
}

// addEchoAndAdd 往 srv 挂上两把测试工具。
func addEchoAndAdd(srv *mcpsdk.Server) {
	srv.AddTool(echoToolSpec(), echoHandler)
	srv.AddTool(addToolSpec(), addHandler)
}

// newFakeServer 构建一台真实的 MCP server（官方 SDK 服务端），挂上 echo/add。
func newFakeServer() *mcpsdk.Server {
	srv := mcpsdk.NewServer(&mcpsdk.Implementation{Name: "fake-mcp", Version: "1.0.0"}, nil)
	addEchoAndAdd(srv)
	return srv
}

// newEmptyServer 构建一台不提供任何工具的服务器（只声明 logging 能力）。
func newEmptyServer() *mcpsdk.Server {
	return mcpsdk.NewServer(&mcpsdk.Implementation{Name: "empty-mcp", Version: "1.0.0"}, nil)
}

// startFakeServer 在 SDK 内存传输上跑 srv，返回客户端侧传输。
func startFakeServer(t *testing.T, srv *mcpsdk.Server) mcpsdk.Transport {
	t.Helper()
	serverT, clientT := mcpsdk.NewInMemoryTransports()
	ctx, cancel := context.WithCancel(context.Background())
	go func() { _ = srv.Run(ctx, serverT) }()
	t.Cleanup(cancel)
	return clientT
}

// connectFakeClient 通过内存传输连上一台 fake server。
func connectFakeClient(t *testing.T, srv *mcpsdk.Server, opts ...Option) *Client {
	t.Helper()
	opts = append(opts, WithTransport(startFakeServer(t, srv)))
	c := New(MCPConfig{Name: "fake"}, opts...)
	if err := c.Connect(context.Background()); err != nil {
		t.Fatalf("Connect: %v", err)
	}
	t.Cleanup(func() { _ = c.Close() })
	return c
}

// toolNames 取出工具名列表。
func toolNames(tools []ToolInfo) []string {
	out := make([]string, 0, len(tools))
	for _, t := range tools {
		out = append(out, t.Name)
	}
	return out
}

// waitFor 轮询 cond 直到成立或超时。
func waitFor(t *testing.T, timeout time.Duration, what string, cond func() bool) {
	t.Helper()
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		if cond() {
			return
		}
		time.Sleep(5 * time.Millisecond)
	}
	t.Fatalf("等待超时：%s", what)
}

// TestMCPStdioHelper 是 stdio 打桩用的辅助进程入口：以真实 MCP server 身份
// 在 stdin/stdout 上讲协议。stdout 是协议通道，因此诊断只写 stderr，且退出时
// 直接 os.Exit（避免测试框架往 stdout 打 PASS/ok 破坏协议帧）。
func TestMCPStdioHelper(t *testing.T) {
	if os.Getenv(stdioHelperEnv) != "1" {
		t.Skip("仅作为 stdio 辅助进程入口")
	}
	srv := newFakeServer()
	if err := srv.Run(context.Background(), &mcpsdk.StdioTransport{}); err != nil {
		fmt.Fprintln(os.Stderr, "stdio helper server:", err)
		os.Exit(1)
	}
	os.Exit(0)
}