Cyrene_v2/internal/capability/mcp/capabilities_test.go

package mcp

import (
	"context"
	"errors"
	"testing"

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

// TestFeaturesFrom_DeclarationMatrix 覆盖「声明了才算可用」的判定矩阵。
func TestFeaturesFrom_DeclarationMatrix(t *testing.T) {
	t.Parallel()
	host := hostFeatures{sampling: true, elicitation: true, roots: true}
	tests := []struct {
		name         string
		caps         map[string]any
		wantTools    bool
		wantPrompts  bool
		wantSampling bool
	}{
		{name: "未声明任何能力（legacy 服务器）", caps: map[string]any{}, wantTools: true, wantSampling: false},
		{name: "只有 tools", caps: map[string]any{"tools": map[string]any{}}, wantTools: true, wantSampling: true},
		{name: "只有 logging", caps: map[string]any{"logging": map[string]any{}}, wantTools: false, wantSampling: true},
		{name: "tools+prompts", caps: map[string]any{"tools": map[string]any{}, "prompts": map[string]any{}}, wantTools: true, wantPrompts: true, wantSampling: true},
		{name: "tools 为 null 视为未声明", caps: map[string]any{"tools": nil}, wantTools: false, wantSampling: true},
	}
	for _, tt := range tests {
		tt := tt
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()
			got := featuresFrom(tt.caps, host)
			if got.Tools != tt.wantTools {
				t.Fatalf("Tools = %v, 期望 %v（caps=%v）", got.Tools, tt.wantTools, tt.caps)
			}
			if got.Prompts != tt.wantPrompts {
				t.Fatalf("Prompts = %v, 期望 %v", got.Prompts, tt.wantPrompts)
			}
			if got.Sampling != tt.wantSampling {
				t.Fatalf("Sampling = %v, 期望 %v", got.Sampling, tt.wantSampling)
			}
		})
	}
}

// TestFeaturesFrom_HostNotConfigured 宿主没配回调时，即使服务器声明了能力也不可用。
func TestFeaturesFrom_HostNotConfigured(t *testing.T) {
	t.Parallel()
	got := featuresFrom(map[string]any{"tools": map[string]any{}}, hostFeatures{})
	if got.Sampling || got.Elicitation || got.Roots {
		t.Fatalf("宿主未配置任何客户端能力，却推断为可用：%+v", got)
	}
}

// TestClient_SamplingHandler_RefusedWhenNotDeclared 服务器未声明能力族时，
// 包装后的 sampling 回调必须直接拒绝，不放行给宿主。
func TestClient_SamplingHandler_RefusedWhenNotDeclared(t *testing.T) {
	t.Parallel()
	called := false
	c := New(MCPConfig{Name: "s"}, WithSamplingHandler(func(context.Context, *mcpsdk.CreateMessageRequest) (*mcpsdk.CreateMessageResult, error) {
		called = true
		return &mcpsdk.CreateMessageResult{}, nil
	}))
	h := c.wrapSampling()
	if h == nil {
		t.Fatal("配了 sampling 回调，wrapSampling 却返回 nil")
	}
	// 未连接（等价服务器什么都没声明）：必须拒绝。
	if _, err := h(context.Background(), &mcpsdk.CreateMessageRequest{}); !errors.Is(err, ErrCapabilityNotDeclared) {
		t.Fatalf("未声明时 error = %v, 期望 ErrCapabilityNotDeclared", err)
	}
	if called {
		t.Fatal("宿主回调在服务器未声明能力时被调用了")
	}

	// 服务器声明过能力族 → 放行。
	c.mu.Lock()
	c.info.Features = featuresFrom(map[string]any{"tools": map[string]any{}}, c.hostFeatures())
	c.mu.Unlock()
	if _, err := h(context.Background(), &mcpsdk.CreateMessageRequest{}); err != nil {
		t.Fatalf("已声明时 error = %v, 期望放行", err)
	}
	if !called {
		t.Fatal("已声明时宿主回调没有被调用")
	}
}

// TestClient_SamplingHandler_RefusedWhenHostNotConfigured 没配回调就不包装（返回 nil）。
func TestClient_SamplingHandler_RefusedWhenHostNotConfigured(t *testing.T) {
	t.Parallel()
	c := New(MCPConfig{Name: "s"})
	if c.wrapSampling() != nil {
		t.Fatal("未配置 sampling 回调时不应注册 SDK handler")
	}
	if c.wrapElicitation() != nil {
		t.Fatal("未配置 elicitation 回调时不应注册 SDK handler")
	}
}

// TestClient_ClientCapabilities_OnlyConfigured 客户端只宣告宿主真正具备的能力。
func TestClient_ClientCapabilities_OnlyConfigured(t *testing.T) {
	t.Parallel()
	bare := New(MCPConfig{Name: "s"}).clientCapabilities()
	if bare.Sampling != nil || bare.Elicitation != nil || bare.RootsV2 != nil {
		t.Fatalf("默认客户端能力应为空：%+v", bare)
	}
	full := New(MCPConfig{Name: "s"},
		WithSamplingHandler(func(context.Context, *mcpsdk.CreateMessageRequest) (*mcpsdk.CreateMessageResult, error) {
			return nil, nil
		}),
		WithElicitationHandler(func(context.Context, *mcpsdk.ElicitRequest) (*mcpsdk.ElicitResult, error) {
			return nil, nil
		}),
		WithRoots(&mcpsdk.Root{URI: "file:///workspace"}),
	).clientCapabilities()
	if full.Sampling == nil || full.Elicitation == nil || full.RootsV2 == nil {
		t.Fatalf("配置齐全时客户端能力不全：%+v", full)
	}
}