Cyrene_v2/internal/capability/mcp/mcp.go

// Package mcp provides a minimal Model Context Protocol (MCP) client and the
// auto-registration glue that surfaces MCP server tools as capabilities.
//
// An MCP server is an external tool provider. This package speaks JSON-RPC 2.0
// to it over stdio (a spawned subprocess) or SSE, discovers the tools it
// exposes via tools/list, maps each one to a pkg/plugins/sdk Tool, and registers
// them into the capability registry so any character can invoke them
// (DevDocs/design/09 插件生态 §1).
//
// It is deliberately dependency-light: only the standard library and the
// project packages (plugins/sdk, capability/registry, internal/logger) are
// used. The wire protocol is implemented directly, so no third-party MCP SDK is
// required.
//
// A caller usually wires it up as:
//
//	reg := registry.New()
//	mcpReg := mcp.New(reg, mcp.MCPConfig{Name: "calc", Command: "mcp-calc"})
//	if err := mcpReg.Register(ctx); err != nil {
//	    log.Fatal(err)
//	}
//	defer mcpReg.Unregister(ctx)
//
// which connects to the server, discovers its tools, and registers them. Call
// sites then invoke the tools by their MCP name through the registry.
package mcp

import (
	"context"
	"errors"
	"io"
	"log/slog"
	"sync"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
)

// protocolVersion is the MCP protocol version this client advertises during the
// initialize handshake.
const protocolVersion = "2024-11-05"

// Sentinel errors returned by this package. Test with errors.Is.
var (
	// ErrNotConnected is returned when an operation requires an established
	// connection but the client is not yet connected or is closed.
	ErrNotConnected = errors.New("mcp: not connected")
	// ErrClosed is returned when operating on a client or transport that has
	// been closed.
	ErrClosed = errors.New("mcp: closed")
	// ErrNoTransport is returned when the configuration selects no transport
	// (neither Command nor URL is set).
	ErrNoTransport = errors.New("mcp: no transport configured")
	// ErrRPC is returned when the MCP server replies with a JSON-RPC error
	// object. The concrete server error is wrapped and accessible via
	// errors.Is/errors.As on an *rpcError.
	ErrRPC = errors.New("mcp: rpc error")
)

// MCPConfig describes a single MCP server connection. Exactly one transport is
// selected: stdio when Command is set, SSE when URL is set.
type MCPConfig struct {
	// Name is the human-readable label of the server, used in logs and as the
	// prefix for registering its tools uniquely.
	Name string `env:"MCP_NAME"`
	// Command is the executable path to spawn for a stdio MCP server.
	Command string `env:"MCP_COMMAND"`
	// Args are the arguments passed to Command.
	Args []string `env:"MCP_ARGS"`
	// URL is the SSE endpoint of an MCP server.
	URL string `env:"MCP_URL"`
	// Enabled reports whether this server should be activated at startup.
	Enabled bool `env:"MCP_ENABLED"`
}

// Config holds the set of configured MCP servers.
type Config struct {
	// Servers is the ordered list of MCP server configurations.
	Servers []MCPConfig
}

// ServerInfo describes the MCP server reported by the initialize handshake.
type ServerInfo struct {
	// Name is the server-provided name.
	Name string
	// Version is the server-provided version string.
	Version string
}

// ToolInfo is a tool exposed by an MCP server.
type ToolInfo struct {
	// Name is the tool identifier used to invoke it.
	Name string
	// Description explains what the tool does.
	Description string
	// InputSchema is the tool's JSON Schema, typically
	// {"type":"object","properties":{...}}.
	InputSchema map[string]any
}

// ContentBlock is a single item in a tool call result (MCP content block).
type ContentBlock struct {
	// Type is the block type, most commonly "text".
	Type string `json:"type"`
	// Text is the textual payload.
	Text string `json:"text"`
}

// ToolResult is the structured result of a tools/call.
type ToolResult struct {
	// Content holds the response content blocks.
	Content []ContentBlock
	// IsError reports whether the server flagged the call as an error.
	IsError bool
}

// Client is a minimal MCP client bound to a single server described by cfg. It
// is safe for concurrent use; all calls are serialized internally.
type Client struct {
	cfg   MCPConfig
	log   *slog.Logger
	tr    Transport
	mu    sync.Mutex
	state state
	info  ServerInfo
}

// state tracks the client lifecycle.
type state int

const (
	stateNew state = iota
	stateConnected
	stateClosed
)

// Transport is the JSON-RPC message exchange with an MCP server over a concrete
// transport (stdio subprocess or SSE). Implementations must be safe for
// concurrent use. This is exported so tests and hosts can inject a fake
// transport.
type Transport interface {
	// Start makes the transport ready to exchange messages.
	Start(ctx context.Context) error
	// Call sends a JSON-RPC request and decodes the matching response into
	// result. Params is marshaled into the request body; result is unmarshaled
	// from the response's result field.
	Call(ctx context.Context, method string, params any, result any) error
	// Close releases every underlying resource (subprocess, streams, sockets).
	Close() error
}

// notifier is implemented by transports that can send a JSON-RPC notification
// (a request with no id and no response).
type notifier interface {
	notify(ctx context.Context, method string, params any) error
}

// Option customizes a Client built by New.
type Option func(*Client)

// WithLogger overrides the client's logger. The given logger is tagged with the
// source capability/mcp.
func WithLogger(l *slog.Logger) Option {
	return func(c *Client) {
		if l != nil {
			c.log = logger.WithSource(l, "capability/mcp")
		}
	}
}

// WithTransport overrides the transport used by the client. This is primarily
// for tests that inject a fake transport instead of spawning a subprocess or
// dialing SSE.
func WithTransport(t Transport) Option {
	return func(c *Client) {
		c.tr = t
	}
}

// New builds a Client for cfg. By default it uses a discard logger and selects a
// stdio or SSE transport from cfg; pass WithLogger and/or WithTransport to
// override.
func New(cfg MCPConfig, opts ...Option) *Client {
	c := &Client{
		cfg:   cfg,
		log:   defaultLogger(),
		state: stateNew,
	}
	for _, opt := range opts {
		opt(c)
	}
	return c
}

// defaultLogger returns a discard logger tagged with the mcp source. Hosts
// should prefer WithLogger so records are observable.
func defaultLogger() *slog.Logger {
	base := slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelWarn}))
	return logger.WithSource(base, "capability/mcp")
}

// transport returns the established transport, or nil when not connected.
func (c *Client) transport() Transport {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.state != stateConnected || c.tr == nil {
		return nil
	}
	return c.tr
}