// Package mcp is the client side of the Model Context Protocol: it starts the
// servers the config names, asks each what tools it has, and offers those
// tools to the agent as if they were Codey's own.
//
// Only the stdio transport is spoken. A server is a child process with
// JSON-RPC 2.0 on its stdin and stdout, one message per line — the same
// arrangement as the language servers in package diag, save that the LSP
// frames its messages with a Content-Length header and MCP with a newline.
// This file is the sibling of diag/lsp.go for that reason, and where the two
// look alike it is because the same thing goes wrong in both.
package mcp
import (
"bufio"
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"sync"
"time"
)
// protocolVersion is what the client asks for. supportedVersions are the
// answers it will take: every one of them has the same initialize, tools/list
// and tools/call shapes over stdio, and those are the only methods this client
// uses, so a server answering with an older one is a server it can still
// talk to. A version outside the list is one whose messages may mean
// something else, and the spec's instruction is to disconnect rather than
// guess.
const protocolVersion = "2025-06-18"
var supportedVersions = []string{"2025-06-18", "2025-03-26", "2024-11-05"}
// initTimeout bounds the whole handshake: initialize, initialized and the
// first tools/list. A server started with npx may spend most of it
// downloading itself on first run, which is why it is longer than a call. A
// server that has not answered by then is not going to, and every one of its
// tools would otherwise be advertised to the model and then fail.
const initTimeout = 30 * time.Second
// callTimeout is how long one tools/call may take. It matches what the bash
// tool gives a command by default: an MCP tool is a command the model runs
// with different plumbing, and the model's patience is the same either way.
const callTimeout = 60 * time.Second
// maxMessageBytes is the largest line the client will read. Sixteen megabytes
// is far past any tool result the model could use and short of the point
// where a server that streams a file at us shows up in the harness's memory.
// bufio.Scanner enforces it by failing the read, which ends the session; a
// stream that has lost sync cannot be recovered anyway.
const maxMessageBytes = 16 << 20
// errServerDead is what a call gets from a server that has exited or been
// closed. The server is not restarted: a tool that fails this way once fails
// this way for the rest of the session, and the message says so.
var errServerDead = errors.New("MCP server is not running")
// client is one running server and the connection to it.
type client struct {
name string
// stdin is written under writeMu so two tool calls from a parallel batch
// cannot interleave their lines.
writeMu sync.Mutex
stdin io.WriteCloser
// kill ends the process and wait reaps it; both are no-ops for the
// in-process fake the tests use.
kill func()
wait func()
// dead is closed when the connection is gone for any reason; deadErr says
// why, for the error a tool call then gets.
dead chan struct{}
dieOnce sync.Once
deadErr error
mu sync.Mutex
nextID int64
pending map[int64]chan rpcResponse
// serverVersion is what the server answered initialize with, kept for
// error messages: "the server speaks 2024-11-05" is a better complaint
// than "protocol error".
serverVersion string
// hasTools is whether the server declared the tools capability. Without
// it tools/list is a method the server may not have, and the message it
// sends back for that is less clear than saying so ourselves.
hasTools bool
}
// rpcResponse is the part of a response the client uses.
type rpcResponse struct {
Result json.RawMessage
Error *rpcError
}
// rpcError is a JSON-RPC error object.
type rpcError struct {
Code int `json:"code"`
Message string `json:"message"`
}
func (e *rpcError) Error() string { return fmt.Sprintf("%s (%d)", e.Message, e.Code) }
// newClient wires a client over an already-open connection and starts the
// reader. It does not run the handshake: that is the caller's, under the
// caller's deadline. It is the seam the tests use to speak to an in-process
// fake server through pipes.
func newClient(name string, stdin io.WriteCloser, stdout io.Reader, kill, wait func()) *client {
if kill == nil {
kill = func() {}
}
if wait == nil {
wait = func() {}
}
c := &client{
name: name,
stdin: stdin,
kill: kill,
wait: wait,
dead: make(chan struct{}),
pending: map[int64]chan rpcResponse{},
}
go c.readLoop(stdout)
return c
}
// Dead reports whether the connection is gone.
func (c *client) Dead() bool {
select {
case <-c.dead:
return true
default:
return false
}
}
// initialize runs the handshake. The server's protocol version is checked
// against supportedVersions and its capabilities read for tools; the
// initialized notification follows, after which the server may be asked for
// its tools.
func (c *client) initialize(ctx context.Context) error {
params := map[string]any{
"protocolVersion": protocolVersion,
// No roots, no sampling, no elicitation: the client offers the server
// nothing it can ask of Codey. A server that asks anyway is told the
// method is not found, which the spec allows it to expect.
"capabilities": map[string]any{},
"clientInfo": map[string]any{
"name": "codey",
"version": clientVersion,
},
}
raw, err := c.call(ctx, "initialize", params)
if err != nil {
return fmt.Errorf("initialize: %w", err)
}
var res struct {
ProtocolVersion string `json:"protocolVersion"`
Capabilities struct {
Tools json.RawMessage `json:"tools"`
} `json:"capabilities"`
}
if err := json.Unmarshal(raw, &res); err != nil {
return fmt.Errorf("initialize: server sent a malformed result: %w", err)
}
if res.ProtocolVersion == "" {
return fmt.Errorf("initialize: server did not say which protocol version it speaks")
}
if !supported(res.ProtocolVersion) {
return fmt.Errorf("initialize: server speaks MCP %s, which this client does not (it speaks %s)",
res.ProtocolVersion, protocolVersion)
}
c.mu.Lock()
c.serverVersion = res.ProtocolVersion
c.hasTools = len(res.Capabilities.Tools) > 0 && string(res.Capabilities.Tools) != "null"
c.mu.Unlock()
if err := c.notify("notifications/initialized", nil); err != nil {
return fmt.Errorf("initialized: %w", err)
}
return nil
}
func supported(v string) bool {
for _, s := range supportedVersions {
if s == v {
return true
}
}
return false
}
// toolInfo is one entry of tools/list as the server describes it.
type toolInfo struct {
Name string `json:"name"`
Description string `json:"description"`
InputSchema json.RawMessage `json:"inputSchema"`
}
// listTools asks for every tool the server has, following the cursor until
// the server stops sending one. The spec allows a server to page; none of the
// common ones do, but one that did would otherwise silently show up with
// only its first page.
func (c *client) listTools(ctx context.Context) ([]toolInfo, error) {
c.mu.Lock()
hasTools := c.hasTools
c.mu.Unlock()
if !hasTools {
return nil, fmt.Errorf("server offers no tools (it did not declare the tools capability)")
}
var out []toolInfo
cursor := ""
for {
params := map[string]any{}
if cursor != "" {
params["cursor"] = cursor
}
raw, err := c.call(ctx, "tools/list", params)
if err != nil {
return nil, fmt.Errorf("tools/list: %w", err)
}
var res struct {
Tools []toolInfo `json:"tools"`
NextCursor string `json:"nextCursor"`
}
if err := json.Unmarshal(raw, &res); err != nil {
return nil, fmt.Errorf("tools/list: server sent a malformed result: %w", err)
}
out = append(out, res.Tools...)
if res.NextCursor == "" || res.NextCursor == cursor {
return out, nil
}
cursor = res.NextCursor
}
}
// callTool runs one tool and returns the raw result for the adapter to
// render. The deadline is the caller's or callTimeout, whichever comes first.
func (c *client) callTool(ctx context.Context, name string, args json.RawMessage) (json.RawMessage, error) {
ctx, cancel := context.WithTimeout(ctx, callTimeout)
defer cancel()
if len(args) == 0 {
args = json.RawMessage("{}")
}
return c.call(ctx, "tools/call", map[string]any{"name": name, "arguments": args})
}
// Close ends the server. Stdio MCP has no shutdown request: the client
// closes stdin, which is the server's signal to exit, and gives it a moment
// before the process group is killed. The group rather than the leader,
// because npx and uvx wrappers start the real server as a child and a killed
// wrapper leaves that child holding the pipes.
func (c *client) Close() error {
if !c.Dead() {
_ = c.stdin.Close()
select {
case <-c.dead:
case <-time.After(shutdownGrace):
}
}
c.die(errServerDead)
c.wait()
return nil
}
// shutdownGrace is how long a server has to exit on its own after stdin is
// closed. Half a second is what the language servers get, for the same
// reason: one that has not managed it by then is not going to.
const shutdownGrace = 500 * time.Millisecond
// die marks the connection gone, fails every request still waiting on it,
// and ends the process. Everything that can end a client funnels through
// here so it happens once.
func (c *client) die(err error) {
c.dieOnce.Do(func() {
c.mu.Lock()
c.deadErr = err
pending := c.pending
c.pending = map[int64]chan rpcResponse{}
c.mu.Unlock()
for _, ch := range pending {
close(ch)
}
close(c.dead)
c.kill()
_ = c.stdin.Close()
})
}
// deadError is the error a call gets after death: what killed the server,
// wrapped so callers can still test for errServerDead.
func (c *client) deadError() error {
c.mu.Lock()
defer c.mu.Unlock()
if c.deadErr == nil || c.deadErr == errServerDead {
return fmt.Errorf("%w: %s", errServerDead, c.name)
}
return fmt.Errorf("%w: %s: %v", errServerDead, c.name, c.deadErr)
}
// call sends a request and waits for its response, for ctx, or for the
// server to die.
func (c *client) call(ctx context.Context, method string, params any) (json.RawMessage, error) {
c.mu.Lock()
if c.Dead() {
c.mu.Unlock()
return nil, c.deadError()
}
c.nextID++
id := c.nextID
ch := make(chan rpcResponse, 1)
c.pending[id] = ch
c.mu.Unlock()
msg := map[string]any{"jsonrpc": "2.0", "id": id, "method": method}
if params != nil {
msg["params"] = params
}
if err := c.write(msg); err != nil {
c.mu.Lock()
delete(c.pending, id)
c.mu.Unlock()
return nil, err
}
select {
case resp, ok := <-ch:
if !ok {
return nil, c.deadError()
}
if resp.Error != nil {
return nil, resp.Error
}
return resp.Result, nil
case <-ctx.Done():
c.mu.Lock()
delete(c.pending, id)
c.mu.Unlock()
return nil, ctx.Err()
}
}
// notify sends a notification: a request with no id and no reply.
func (c *client) notify(method string, params any) error {
msg := map[string]any{"jsonrpc": "2.0", "method": method}
if params != nil {
msg["params"] = params
}
return c.write(msg)
}
// reply answers a request the server made of us. It writes from its own
// goroutine rather than from the read loop that called it, for the reason
// diag/lsp.go gives: a server blocked writing its next request while we are
// blocked writing this answer is a deadlock the pipe buffer hides until a
// message outgrows it.
func (c *client) reply(id json.RawMessage, result any, err *rpcError) {
msg := map[string]any{"jsonrpc": "2.0", "id": id}
if err != nil {
msg["error"] = err
} else {
msg["result"] = result
}
go func() { _ = c.write(msg) }()
}
// write frames msg as one line and writes it in one call. json.Marshal never
// emits a literal newline — it escapes them inside strings — so the
// trailing one is the only one on the line, which is what makes the framing
// hold.
func (c *client) write(msg any) error {
if c.Dead() {
return c.deadError()
}
body, err := json.Marshal(msg)
if err != nil {
return err
}
body = append(body, '\n')
c.writeMu.Lock()
defer c.writeMu.Unlock()
if _, err := c.stdin.Write(body); err != nil {
// A write that fails is a pipe with nobody on the other end.
c.die(fmt.Errorf("write: %w", err))
return c.deadError()
}
return nil
}
// incoming is the superset of the three JSON-RPC message shapes; which one
// it is depends on which fields are set.
type incoming struct {
ID json.RawMessage `json:"id"`
Method string `json:"method"`
Params json.RawMessage `json:"params"`
Result json.RawMessage `json:"result"`
Error *rpcError `json:"error"`
}
// readLoop reads lines until the connection fails, dispatching responses to
// their callers. A line that is not JSON ends the session: a server that
// prints its logging to stdout instead of stderr has broken the framing for
// good, and the model is better told its tools are gone than left waiting on
// a reply that will never be recognised.
func (c *client) readLoop(r io.Reader) {
sc := bufio.NewScanner(r)
sc.Buffer(make([]byte, 64<<10), maxMessageBytes)
for sc.Scan() {
line := sc.Bytes()
if len(bytes.TrimSpace(line)) == 0 {
continue
}
var msg incoming
if err := json.Unmarshal(line, &msg); err != nil {
c.die(fmt.Errorf("server sent invalid JSON: %w", err))
return
}
switch {
case msg.Method == "" && len(msg.ID) > 0:
c.dispatchResponse(msg)
case msg.Method != "" && len(msg.ID) > 0 && string(msg.ID) != "null":
c.handleRequest(msg)
default:
// Notifications — list_changed, logging, progress — carry nothing
// this client acts on. Tools are listed once at startup.
}
}
if err := sc.Err(); err != nil {
c.die(fmt.Errorf("read: %w", err))
return
}
c.die(errors.New("server exited"))
}
func (c *client) dispatchResponse(msg incoming) {
var id int64
if err := json.Unmarshal(msg.ID, &id); err != nil {
return
}
c.mu.Lock()
ch, ok := c.pending[id]
delete(c.pending, id)
c.mu.Unlock()
if ok {
ch <- rpcResponse{Result: msg.Result, Error: msg.Error}
}
}
// handleRequest answers what a server asks of its client. ping gets its
// empty answer because a server that pings and hears nothing may take the
// silence for a dead client and exit. Everything else — roots, sampling,
// elicitation — is a capability the client did not offer in initialize, and
// the spec's answer for that is method-not-found.
func (c *client) handleRequest(msg incoming) {
switch msg.Method {
case "ping":
c.reply(msg.ID, map[string]any{}, nil)
default:
c.reply(msg.ID, nil, &rpcError{Code: -32601, Message: "method not found"})
}
}