Cyrene_v2/internal/capability/mcp/discovery_test.go

package mcp

import (
	"context"
	"errors"
	"net/http"
	"net/http/httptest"
	"os"
	"slices"
	"sync/atomic"
	"testing"
	"time"

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

	"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
	"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
)

// TestClient_DiscoverTools_InMemory 走真实 SDK 内存传输打通「连接 → 工具自动发现」。
func TestClient_DiscoverTools_InMemory(t *testing.T) {
	t.Parallel()
	c := connectFakeClient(t, newFakeServer())

	info := c.ServerInfo()
	if info.Name != "fake-mcp" || info.Version != "1.0.0" {
		t.Fatalf("ServerInfo = %+v, 期望 fake-mcp/1.0.0", info)
	}
	if info.ProtocolVersion == "" {
		t.Fatal("ProtocolVersion 为空，协商结果没有保留")
	}
	if _, ok := info.Capabilities["tools"]; !ok {
		t.Fatalf("服务器声明的 capabilities 没有保留 tools 字段：%v", info.Capabilities)
	}
	if !info.Features.Tools {
		t.Fatalf("Features.Tools = false，服务器已声明 tools：%v", info.Capabilities)
	}

	tools, err := c.ListTools(context.Background())
	if err != nil {
		t.Fatalf("ListTools: %v", err)
	}
	got := toolNames(tools)
	slices.Sort(got)
	if !slices.Equal(got, []string{"add", "echo"}) {
		t.Fatalf("tools = %v, 期望 [add echo]", got)
	}

	discovered, err := c.DiscoverTools(context.Background())
	if err != nil {
		t.Fatalf("DiscoverTools: %v", err)
	}
	if len(discovered) != 2 {
		t.Fatalf("DiscoverTools 返回 %d 把工具，期望 2", len(discovered))
	}
	if c.SchemaCache().Len() != 2 {
		t.Fatalf("schema 缓存条目 = %d，期望 2（发现即缓存）", c.SchemaCache().Len())
	}
}

// TestClient_DiscoverTools_StdioProcess 用真实子进程（stdio 打桩）证明 stdio 传输下
// 工具自动发现能跑通：子进程就是本测试二进制，以 MCP server 身份讲协议。
func TestClient_DiscoverTools_StdioProcess(t *testing.T) {
	t.Parallel()
	self, err := os.Executable()
	if err != nil {
		t.Fatalf("os.Executable: %v", err)
	}
	c := New(MCPConfig{
		Name:    "stdio-fake",
		Command: self,
		Args:    []string{"-test.run=^TestMCPStdioHelper$"},
		Env:     map[string]string{stdioHelperEnv: "1"},
	})
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	if err := c.Connect(ctx); err != nil {
		t.Fatalf("Connect: %v", err)
	}
	defer c.Close()

	tools, err := c.DiscoverTools(ctx)
	if err != nil {
		t.Fatalf("DiscoverTools: %v", err)
	}
	got := toolNames(toToolInfos(tools))
	slices.Sort(got)
	if !slices.Equal(got, []string{"add", "echo"}) {
		t.Fatalf("stdio 发现的工具 = %v, 期望 [add echo]", got)
	}

	res, err := c.CallTool(ctx, "echo", map[string]any{"name": "stdio"})
	if err != nil {
		t.Fatalf("CallTool: %v", err)
	}
	if text := contentText(res.Content); text != "hello stdio" {
		t.Fatalf("工具结果 = %q, 期望 hello stdio", text)
	}
}

// TestClient_DiscoverTools_StreamableHTTP 用 httptest 打桩 Streamable HTTP 端点，
// 验证 URL 型传输 + bearer 鉴权 + 工具自动发现整条链。
func TestClient_DiscoverTools_StreamableHTTP(t *testing.T) {
	t.Parallel()
	srv := newFakeServer()
	handler := mcpsdk.NewStreamableHTTPHandler(func(*http.Request) *mcpsdk.Server { return srv }, nil)
	var authorized atomic.Bool
	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != "Bearer sk-test" {
			w.WriteHeader(http.StatusUnauthorized)
			return
		}
		authorized.Store(true)
		handler.ServeHTTP(w, r)
	}))
	defer ts.Close()

	c := New(MCPConfig{Name: "http-fake", URL: ts.URL, TransportType: "http", APIKey: "sk-test"})
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	if err := c.Connect(ctx); err != nil {
		t.Fatalf("Connect: %v", err)
	}
	defer c.Close()

	tools, err := c.DiscoverTools(ctx)
	if err != nil {
		t.Fatalf("DiscoverTools: %v", err)
	}
	if len(tools) != 2 {
		t.Fatalf("Streamable HTTP 发现 %d 把工具，期望 2", len(tools))
	}
	if !authorized.Load() {
		t.Fatal("服务器没有收到带 bearer 凭据的请求")
	}
	res, err := c.CallTool(ctx, "add", map[string]any{"a": 1, "b": 2})
	if err != nil {
		t.Fatalf("CallTool: %v", err)
	}
	if text := contentText(res.Content); text != "3" {
		t.Fatalf("add 结果 = %q, 期望 3", text)
	}
}

// TestClient_ListTools_RefusesUndeclaredCapability 证明「未声明不得调用」：
// 服务器只声明 logging，客户端拒绝发 tools/list。
func TestClient_ListTools_RefusesUndeclaredCapability(t *testing.T) {
	t.Parallel()
	c := connectFakeClient(t, newEmptyServer())
	if c.Features().Tools {
		t.Fatalf("空服务器不应推断出 tools 能力：%v", c.ServerInfo().Capabilities)
	}
	_, err := c.ListTools(context.Background())
	if !errors.Is(err, ErrCapabilityNotDeclared) {
		t.Fatalf("ListTools error = %v, 期望 ErrCapabilityNotDeclared", err)
	}
	if _, err := c.DiscoverTools(context.Background()); !errors.Is(err, ErrCapabilityNotDeclared) {
		t.Fatalf("DiscoverTools error = %v, 期望 ErrCapabilityNotDeclared", err)
	}
}

// TestRegistry_Register_AutoRegisters 打通「自动发现 → 注册进 capability registry
// → 模型可调用」这条链（用 registry 里存的工具值真的发起一次调用）。
func TestRegistry_Register_AutoRegisters(t *testing.T) {
	t.Parallel()
	reg := registry.New()
	tr := startFakeServer(t, newFakeServer())
	r := NewRegistry(reg, MCPConfig{Name: "calc"}, WithTransport(tr))

	if err := r.Register(context.Background()); err != nil {
		t.Fatalf("Register: %v", err)
	}
	names := r.ToolNames()
	slices.Sort(names)
	if !slices.Equal(names, []string{"add", "echo"}) {
		t.Fatalf("ToolNames = %v, 期望 [add echo]", names)
	}

	cap, err := reg.Get("echo")
	if err != nil {
		t.Fatalf("registry.Get(echo): %v", err)
	}
	if cap.Kind != registry.KindTool {
		t.Fatalf("echo 的 Kind = %q, 期望 %q", cap.Kind, registry.KindTool)
	}
	tool, ok := cap.Value.(sdk.Tool)
	if !ok {
		t.Fatalf("echo 的 Value 类型 = %T, 期望 sdk.Tool", cap.Value)
	}

	// 不给参数调用：默认值 world 由 input schema 补上（证明 schema 校验/默认值生效）。
	out, err := tool.Execute(context.Background(), nil)
	if err != nil {
		t.Fatalf("Execute(echo): %v", err)
	}
	if out != "hello world" {
		t.Fatalf("echo 结果 = %v, 期望 hello world（默认值未生效？）", out)
	}

	// 必填参数缺失或类型错误时必须被 input schema 拦下，不把坏参数发给服务器。
	addCap, err := reg.Get("add")
	if err != nil {
		t.Fatalf("registry.Get(add): %v", err)
	}
	addTool := addCap.Value.(sdk.Tool)
	if _, err := addTool.Execute(context.Background(), map[string]any{"a": "not-a-number"}); err == nil {
		t.Fatal("非法参数没有被 input schema 拦下")
	}
}

// TestRegistry_Register_PropagatesConnectError 连接失败要冒泡，不能静默注册空工具集。
func TestRegistry_Register_PropagatesConnectError(t *testing.T) {
	t.Parallel()
	reg := registry.New()
	r := NewRegistry(reg, MCPConfig{Name: "broken", Command: "definitely-not-a-real-binary-xyz"})
	if err := r.Register(context.Background()); err == nil {
		t.Fatal("Register 成功，期望连接错误")
	}
}

// TestRegistry_Register_NilRegistry 空 registry 必须报错。
func TestRegistry_Register_NilRegistry(t *testing.T) {
	t.Parallel()
	r := NewRegistry(nil, MCPConfig{Name: "calc"})
	if err := r.Register(context.Background()); err == nil {
		t.Fatal("Register 成功，期望 nil registry 错误")
	}
}

// TestRegistry_Unregister_RemovesTools 卸载要真从 registry 移除并关会话，且幂等。
func TestRegistry_Unregister_RemovesTools(t *testing.T) {
	t.Parallel()
	reg := registry.New()
	tr := startFakeServer(t, newFakeServer())
	r := NewRegistry(reg, MCPConfig{Name: "calc"}, WithTransport(tr))
	if err := r.Register(context.Background()); err != nil {
		t.Fatalf("Register: %v", err)
	}
	if err := r.Unregister(context.Background()); err != nil {
		t.Fatalf("Unregister: %v", err)
	}
	if err := r.Unregister(context.Background()); err != nil {
		t.Fatalf("第二次 Unregister: %v", err)
	}
	for _, name := range []string{"add", "echo"} {
		if _, err := reg.Get(name); !errors.Is(err, registry.ErrNotFound) {
			t.Fatalf("Get(%q) 卸载后 = %v, 期望 ErrNotFound", name, err)
		}
	}
}

const (
	// methodSubscriptionsAck 是 SEP-2575 的订阅确认通知方法名（SDK 未导出该常量）。
	methodSubscriptionsAck = "notifications/subscriptions/acknowledged"
	// methodToolListChanged 是工具清单变更通知方法名（SDK 未导出该常量）。
	methodToolListChanged = "notifications/tools/list_changed"
	// sep2575ProtocolVersion 是引入 subscriptions/listen 订阅流的协议版本：自它起
	// 服务器只向**已登记订阅**的会话投递 list_changed 通知（更早的协议版本对全部
	// 会话广播，不存在下面这个窗口）。
	sep2575ProtocolVersion = "2026-07-28"
	// notificationWaitBudget 是等上述事件的兜底时限：只用于避免测试挂死，
	// 事件本身才是断言对象。
	notificationWaitBudget = 5 * time.Second
)

// TestRegistry_ToolListChangedRefreshes 服务器运行中加工具 → 通知 → registry 增量刷新，
// 同时 schema 缓存按服务器整表失效（design/09 §1.1）。
//
// 确定性要点：现代协议（≥ sep2575ProtocolVersion）下服务器**只**把 tools/list_changed
// 投给已登记 subscriptions/listen 订阅的会话，而客户端发订阅是 fire-and-forget
// （Connect 返回 ≠ 服务器已登记），通知本身又是「变更后 10ms 发一次、丢了不补」——
// 订阅登记之前触发的变更会被永久静默丢弃。旧写法（改工具 → 轮询 5s）等于拿这个竞态
// 与机器负载对赌，输的那次等多久都等不到 mul。现在先等服务器回订阅确认（服务器先登记
// 订阅、再回确认）再触发变更，把「通知会不会来」变回确定的事。
func TestRegistry_ToolListChangedRefreshes(t *testing.T) {
	t.Parallel()
	srv := newFakeServer()
	reg := registry.New()
	sn := &sniffTransport{inner: startFakeServer(t, srv)}
	r := NewRegistry(reg, MCPConfig{Name: "calc"}, WithTransport(sn))
	if err := r.Register(context.Background()); err != nil {
		t.Fatalf("Register: %v", err)
	}

	if v := r.mcp.ServerInfo().ProtocolVersion; v >= sep2575ProtocolVersion {
		waitFor(t, notificationWaitBudget, "服务器回订阅确认 "+methodSubscriptionsAck,
			func() bool { return sn.received(methodSubscriptionsAck) })
	}

	_, _ = reg.Get("echo") // 让缓存先有命中记录，便于确认失效计数
	srv.AddTool(&mcpsdk.Tool{
		Name:        "mul",
		Description: "两数相乘",
		InputSchema: addToolSpec().InputSchema,
	}, addHandler)

	waitFor(t, notificationWaitBudget, "registry 出现新工具 mul", func() bool {
		_, err := reg.Get("mul")
		return err == nil
	})
	// 新工具必须由变更通知驱动刷进来，而不是别的路径顺手带进来的。
	if !sn.received(methodToolListChanged) {
		t.Fatalf("registry 出现了 mul，但客户端没收到 %s", methodToolListChanged)
	}
	if _, err := reg.Get("echo"); err != nil {
		t.Fatalf("刷新后 echo 丢失：%v", err)
	}
	if got := r.ToolNames(); !slices.Contains(got, "mul") {
		t.Fatalf("ToolNames = %v, 期望含 mul", got)
	}
}

// toToolInfos 把 sdk.Tool 列表还原成 ToolInfo（仅测试断言工具名用）。
func toToolInfos(tools []sdk.Tool) []ToolInfo {
	out := make([]ToolInfo, 0, len(tools))
	for _, t := range tools {
		info := ToolInfo{Name: t.Name(), Description: t.Description()}
		if t2, ok := t.(*mcpTool); ok {
			info = t2.info
		}
		out = append(out, info)
	}
	return out
}