Cyrene_v2/internal/capability/mcp/keepalive_test.go

package mcp

import (
	"context"
	"sync"
	"testing"
	"time"

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

// TestClient_KeepAlive_SendsPings keepalive 接线：配了间隔后 ping 真的上了线路。
// ping 本身由 SDK 负责，这里只验我们把配置传导到了 SDK 会话上。
func TestClient_KeepAlive_SendsPings(t *testing.T) {
	t.Parallel()
	sn := &sniffTransport{inner: startFakeServer(t, newFakeServer())}
	c := New(MCPConfig{Name: "fake"},
		WithTransport(sn),
		// 固定到 keepalive 生效的旧协议版本，避免新版协议的订阅路径干扰。
		WithProtocolVersion("2025-11-25"),
		WithKeepAlive(20*time.Millisecond, 1))
	if err := c.Connect(context.Background()); err != nil {
		t.Fatalf("Connect: %v", err)
	}
	defer c.Close()
	if got := c.ServerInfo().ProtocolVersion; got != "2025-11-25" {
		t.Fatalf("协议版本 = %q, 期望 2025-11-25", got)
	}
	waitFor(t, 5*time.Second, "线路上出现 keepalive ping", func() bool { return sn.called("ping") })
}

// TestClient_NoKeepAlive_NoPings 不配 keepalive 时不发 ping（反证：ping 不是 SDK 默认行为）。
func TestClient_NoKeepAlive_NoPings(t *testing.T) {
	t.Parallel()
	sn := &sniffTransport{inner: startFakeServer(t, newFakeServer())}
	c := New(MCPConfig{Name: "fake"}, WithTransport(sn), WithProtocolVersion("2025-11-25"))
	if err := c.Connect(context.Background()); err != nil {
		t.Fatalf("Connect: %v", err)
	}
	defer c.Close()
	if _, err := c.DiscoverTools(context.Background()); err != nil {
		t.Fatalf("DiscoverTools: %v", err)
	}
	time.Sleep(150 * time.Millisecond)
	if sn.called("ping") {
		t.Fatal("没配 keepalive 却发了 ping")
	}
	if !sn.called("tools/list") {
		t.Fatalf("线路上没有 tools/list, 已见方法 = %v", sn.methods())
	}
}

// sniffTransport 包一层 SDK 传输，在线路层记录两个方向的 JSON-RPC 请求方法：
// seen 是客户端发出的（断言客户端确实发了某个请求），fromServer 是服务器发来的
// （断言服务器确实回了某条通知，例如 SEP-2575 的订阅确认）。
type sniffTransport struct {
	inner      mcpsdk.Transport
	mu         sync.Mutex
	seen       []string
	fromServer []string
}

// Connect 实现 mcpsdk.Transport。
func (t *sniffTransport) Connect(ctx context.Context) (mcpsdk.Connection, error) {
	conn, err := t.inner.Connect(ctx)
	if err != nil {
		return nil, err
	}
	return &sniffConn{Connection: conn, owner: t}, nil
}

// called 报告某个方法是否出现过。
func (t *sniffTransport) called(method string) bool {
	t.mu.Lock()
	defer t.mu.Unlock()
	for _, m := range t.seen {
		if m == method {
			return true
		}
	}
	return false
}

// methods 返回已见方法快照。
func (t *sniffTransport) methods() []string {
	t.mu.Lock()
	defer t.mu.Unlock()
	return append([]string(nil), t.seen...)
}

// record 记录一个发出的方法名。
func (t *sniffTransport) record(method string) {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.seen = append(t.seen, method)
}

// received 报告服务器是否发来过某个方法。
func (t *sniffTransport) received(method string) bool {
	t.mu.Lock()
	defer t.mu.Unlock()
	for _, m := range t.fromServer {
		if m == method {
			return true
		}
	}
	return false
}

// recordReceived 记录一个服务器发来的方法名。
func (t *sniffTransport) recordReceived(method string) {
	t.mu.Lock()
	defer t.mu.Unlock()
	t.fromServer = append(t.fromServer, method)
}

// sniffConn 拦截写入的请求消息，记录方法名后交给内层连接。
type sniffConn struct {
	mcpsdk.Connection
	owner *sniffTransport
}

// Write 记录请求方法后转发。
func (c *sniffConn) Write(ctx context.Context, msg sdkjsonrpc.Message) error {
	if req, ok := msg.(*sdkjsonrpc.Request); ok {
		c.owner.record(req.Method)
	}
	return c.Connection.Write(ctx, msg)
}

// Read 记录服务器发来的请求方法后转发。通知也是不带 ID 的请求帧，因此
// 订阅确认、tools/list_changed 这类服务器通知都在这里现形。
func (c *sniffConn) Read(ctx context.Context) (sdkjsonrpc.Message, error) {
	msg, err := c.Connection.Read(ctx)
	if err != nil {
		// 出错时内层消息未被上层使用（jsonrpc2 读循环遇错即退出），原样透传错误。
		return nil, err
	}
	if req, ok := msg.(*sdkjsonrpc.Request); ok {
		c.owner.recordReceived(req.Method)
	}
	return msg, nil
}