Cyrene_v2/internal/capability/mcp/helpers_test.go

package mcp

import (
	"context"
	"fmt"
	"sync"
)

// fakeTransport is a scripted Transport for tests. It records starts, closes and
// the methods it was asked to call so assertions can verify the client's usage.
// It also lets tests script the initialize handshake, simulate server→client
// notifications and control the ping result.
type fakeTransport struct {
	startErr error
	closeErr error
	initErr  error
	listErr  error
	callErr  error
	pingErr  error
	tools    []ToolInfo

	// initErrors is a per-call queue of initialize errors. When non-empty one
	// entry is consumed per initialize call; once exhausted the handshake
	// succeeds. initErr is returned for every call when initErrors is empty.
	initErrors []error
	// initVersions records the protocolVersion sent in each initialize call.
	initVersions []string
	// initResponseVersion is the protocolVersion echoed on a successful
	// initialize; default 2024-11-05.
	initResponseVersion string

	mu          sync.Mutex
	starts      int
	closes      int
	methods     []string
	pings       int
	onNotifyFn  func(method string)
	activeTools []ToolInfo
}

// Start records a start.
func (f *fakeTransport) Start(_ context.Context) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.starts++
	f.activeTools = f.tools
	return f.startErr
}

// Call answers the MCP methods the client issues.
func (f *fakeTransport) Call(_ context.Context, method string, params any, result any) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.methods = append(f.methods, method)
	switch method {
	case "initialize":
		if p, ok := params.(initializeParams); ok {
			f.initVersions = append(f.initVersions, p.ProtocolVersion)
		}
		if len(f.initErrors) > 0 {
			err := f.initErrors[0]
			f.initErrors = f.initErrors[1:]
			return err
		}
		if f.initErr != nil {
			return f.initErr
		}
		if r, ok := result.(*initializeResult); ok {
			r.ServerInfo = serverInfo{Name: "fake-server", Version: "1.0"}
			r.ProtocolVersion = f.initResponseVersion
			if r.ProtocolVersion == "" {
				r.ProtocolVersion = protocolVersion20241105
			}
			f.activeTools = f.tools
		}
	case "tools/list":
		if f.listErr != nil {
			return f.listErr
		}
		if r, ok := result.(*listToolsResult); ok {
			for _, t := range f.activeTools {
				r.Tools = append(r.Tools, toolSpec{Name: t.Name, Description: t.Description, InputSchema: t.InputSchema})
			}
		}
	case "tools/call":
		if f.callErr != nil {
			return f.callErr
		}
		if r, ok := result.(*callToolResult); ok {
			r.Content = []contentBlock{{Type: "text", Text: "hello"}}
		}
	case "ping":
		f.pings++
		if f.pingErr != nil {
			return f.pingErr
		}
	}
	return nil
}

// notify records a notification method.
func (f *fakeTransport) notify(_ context.Context, method string, _ any) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.methods = append(f.methods, "notify:"+method)
	return nil
}

// setOnNotification records the client's server→client notification handler so
// tests can deliver one via notifyServer.
func (f *fakeTransport) setOnNotification(fn func(method string)) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.onNotifyFn = fn
}

// notifyServer simulates a server→client notification (e.g.
// notifications/tools/list_changed) delivered to the client's handler.
func (f *fakeTransport) notifyServer(method string) {
	f.mu.Lock()
	fn := f.onNotifyFn
	f.mu.Unlock()
	if fn != nil {
		fn(method)
	}
}

// Close records a close.
func (f *fakeTransport) Close() error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.closes++
	return f.closeErr
}

// calledMethods returns a snapshot of the recorded methods.
func (f *fakeTransport) calledMethods() []string {
	f.mu.Lock()
	defer f.mu.Unlock()
	out := make([]string, len(f.methods))
	copy(out, f.methods)
	return out
}

// hasMethod reports whether method was recorded.
func (f *fakeTransport) hasMethod(method string) bool {
	for _, m := range f.calledMethods() {
		if m == method {
			return true
		}
	}
	return false
}

// initSentVersions returns the protocol versions the client requested across all
// initialize calls.
func (f *fakeTransport) initSentVersions() []string {
	f.mu.Lock()
	defer f.mu.Unlock()
	out := make([]string, len(f.initVersions))
	copy(out, f.initVersions)
	return out
}

// pingCount returns the recorded ping request count.
func (f *fakeTransport) pingCount() int {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.pings
}

// setTools replaces the tool list served on subsequent tools/list calls.
func (f *fakeTransport) setTools(tools []ToolInfo) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.tools = tools
	f.activeTools = tools
}

// unsupportedProtocolError builds a JSON-RPC error that rejects the requested
// protocol version and advertises the supported set, as a real MCP server does.
func unsupportedProtocolError(supported ...string) error {
	return fmt.Errorf("mcp initialize: %w: %w", ErrRPC, &rpcError{
		Code:    codeUnsupportedProtocolVersion,
		Message: "unsupported protocol version",
		Data:    map[string]any{"supported": supported},
	})
}