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}, }) }