Cyrene_v2/internal/capability/mcp/enhancements_test.go

package mcp

import (
	"context"
	"errors"
	"fmt"
	"reflect"
	"testing"
	"time"
)

func TestClient_Connect_NegotiatesDowngradeOnUnsupportedError(t *testing.T) {
	t.Parallel()
	ft := &fakeTransport{
		initErrors:          []error{unsupportedProtocolError("2024-11-05")},
		initResponseVersion: protocolVersion20241105,
	}
	c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft))
	if err := c.Connect(context.Background()); err != nil {
		t.Fatalf("Connect: %v", err)
	}
	got := ft.initSentVersions()
	if !reflect.DeepEqual(got, []string{"2025-11-25", "2024-11-05"}) {
		t.Fatalf("init versions = %v, want [2025-11-25 2024-11-05]", got)
	}
	if c.protocolVersion() != protocolVersion20241105 {
		t.Fatalf("protocolVersion = %q, want %q", c.protocolVersion(), protocolVersion20241105)
	}
}

func TestClient_Connect_AcceptsServerSupportedVersion(t *testing.T) {
	t.Parallel()
	ft := &fakeTransport{}
	c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft))
	if err := c.Connect(context.Background()); err != nil {
		t.Fatalf("Connect: %v", err)
	}
	if got := ft.initSentVersions(); !reflect.DeepEqual(got, []string{"2025-11-25"}) {
		t.Fatalf("init versions = %v, want [2025-11-25]", got)
	}
	// The fake echoes the legacy version; the client must accept it as supported.
	if c.protocolVersion() != protocolVersion20241105 {
		t.Fatalf("protocolVersion = %q, want %q", c.protocolVersion(), protocolVersion20241105)
	}
}

func TestClient_Connect_NoMutualVersion(t *testing.T) {
	t.Parallel()
	// The server only supports a version the client cannot speak, so negotiation
	// must exhaust and surface ErrUnsupportedProtocolVersion.
	ft := &fakeTransport{initErrors: []error{unsupportedProtocolError("2000-01-01")}}
	c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft))
	err := c.Connect(context.Background())
	if !errors.Is(err, ErrUnsupportedProtocolVersion) {
		t.Fatalf("Connect error = %v, want ErrUnsupportedProtocolVersion", err)
	}
}

func TestClient_Connect_ServerReturnsUnsupportedVersion(t *testing.T) {
	t.Parallel()
	// A server that echoes a version this client does not support (and cannot
	// downgrade to) must fail the handshake with ErrUnsupportedProtocolVersion.
	ft := &fakeTransport{initResponseVersion: "2099-01-01"}
	c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft))
	err := c.Connect(context.Background())
	if !errors.Is(err, ErrUnsupportedProtocolVersion) {
		t.Fatalf("Connect error = %v, want ErrUnsupportedProtocolVersion", err)
	}
}

func TestClient_ToolListChangedInvokesHandler(t *testing.T) {
	t.Parallel()
	called := make(chan struct{}, 1)
	ft := &fakeTransport{tools: []ToolInfo{{Name: "add", InputSchema: map[string]any{}}}}
	c := New(MCPConfig{Name: "s", Command: "x"},
		WithTransport(ft),
		WithToolListChangedHandler(func() { called <- struct{}{} }))
	if err := c.Connect(context.Background()); err != nil {
		t.Fatalf("Connect: %v", err)
	}
	// No handler invoked for unrelated notifications.
	ft.notifyServer("notifications/prompts/list_changed")
	select {
	case <-called:
		t.Fatal("handler invoked for a non-tool notification")
	case <-time.After(50 * time.Millisecond):
	}
	// The real notification must invoke the handler.
	ft.notifyServer(notificationToolsListChanged)
	select {
	case <-called:
	case <-time.After(time.Second):
		t.Fatal("tool list changed handler was not invoked")
	}
}

func TestClient_KeepAlive_SendsPings(t *testing.T) {
	t.Parallel()
	ft := &fakeTransport{}
	c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft), WithKeepAlive(20*time.Millisecond, 3))
	if err := c.Connect(context.Background()); err != nil {
		t.Fatalf("Connect: %v", err)
	}
	defer c.Close()
	waitFor(t, time.Second, func() bool { return ft.pingCount() >= 1 })
}

func TestClient_KeepAlive_FailureClosesSession(t *testing.T) {
	t.Parallel()
	ft := &fakeTransport{pingErr: errors.New("dead peer")}
	c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft), WithKeepAlive(20*time.Millisecond, 1))
	if err := c.Connect(context.Background()); err != nil {
		t.Fatalf("Connect: %v", err)
	}
	// A single failed ping (threshold 1) must close the client/transport.
	waitFor(t, time.Second, func() bool { return c.transport() == nil })
	waitFor(t, time.Second, func() bool {
		ft.mu.Lock()
		defer ft.mu.Unlock()
		return ft.closes >= 1
	})
}

func TestClient_KeepAlive_StopsOnUnsupportedPing(t *testing.T) {
	t.Parallel()
	ft := &fakeTransport{pingErr: fmt.Errorf("mcp ping: %w: %w", ErrRPC, &rpcError{Code: codeMethodNotFound, Message: "method not found"})}
	c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft), WithKeepAlive(20*time.Millisecond, 1))
	if err := c.Connect(context.Background()); err != nil {
		t.Fatalf("Connect: %v", err)
	}
	defer c.Close()
	// A method-not-found ping must stop keepalive without closing the session.
	waitFor(t, time.Second, func() bool { return ft.pingCount() >= 1 })
	time.Sleep(80 * time.Millisecond)
	if c.transport() == nil {
		t.Fatal("client was closed after an unsupported ping, want it to stay connected")
	}
	if ft.pingCount() > 1 {
		t.Fatalf("keepalive kept pinging after method-not-found; pings = %d", ft.pingCount())
	}
}

func TestNewTransport_SelectsTransportType(t *testing.T) {
	t.Parallel()
	tests := []struct {
		name string
		cfg  MCPConfig
		want any
	}{
		{name: "http", cfg: MCPConfig{URL: "http://x", TransportType: "http"}, want: &streamableHTTPTransport{}},
		{name: "streamable-alias", cfg: MCPConfig{URL: "http://x", TransportType: "streamable"}, want: &streamableHTTPTransport{}},
		{name: "sse", cfg: MCPConfig{URL: "http://x", TransportType: "sse"}, want: &sseTransport{}},
		{name: "stdio", cfg: MCPConfig{Command: "x", TransportType: "stdio"}, want: &stdioTransport{}},
		{name: "infer-stdio", cfg: MCPConfig{Command: "x"}, want: &stdioTransport{}},
		{name: "infer-sse", cfg: MCPConfig{URL: "http://x"}, want: &sseTransport{}},
	}
	for _, tt := range tests {
		tt := tt
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()
			tr, err := newTransport(tt.cfg, defaultLogger())
			if err != nil {
				t.Fatalf("newTransport: %v", err)
			}
			if reflect.TypeOf(tr) != reflect.TypeOf(tt.want) {
				t.Fatalf("transport type = %T, want %T", tr, tt.want)
			}
		})
	}
}

func TestNewTransport_ExplicitTypeSelectsEvenWithoutField(t *testing.T) {
	t.Parallel()
	// Selecting an explicit transport does not validate the field; Start does.
	if tr, err := newTransport(MCPConfig{TransportType: "http"}, defaultLogger()); err != nil {
		t.Fatalf("newTransport(http): %v", err)
	} else if _, ok := tr.(*streamableHTTPTransport); !ok {
		t.Fatalf("transport type = %T, want *streamableHTTPTransport", tr)
	}
}

// waitFor polls cond until it returns true or the deadline passes, failing the
// test on timeout.
func waitFor(t *testing.T, timeout time.Duration, cond func() bool) {
	t.Helper()
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		if cond() {
			return
		}
		time.Sleep(5 * time.Millisecond)
	}
	t.Fatal("condition not met before timeout")
}