feat(capability/mcp): MCP 客户端(stdlib JSON-RPC over stdio/SSE)+ Registry 自动注册/Unregister(registry 增 Remove)

This commit is contained in:
Aska committed 2026-09-06 06:28:38 +08:00
1 parent 08223c06a5
commit 488cf9c79c
14 files changed
+2212

No files matched your search

+234
View File
@@ -0,0 +1,234 @@
package mcp
import (
"context"
"fmt"
"log/slog"
"git.yeij.top/AskaEth/Cyrene_v2/pkg/plugins/sdk"
)
// initializeParams is the MCP initialize request body.
type initializeParams struct {
ProtocolVersion string `json:"protocolVersion"`
Capabilities map[string]any `json:"capabilities"`
ClientInfo clientInfo `json:"clientInfo"`
}
// initializeResult is the MCP initialize response body.
type initializeResult struct {
ProtocolVersion string `json:"protocolVersion"`
Capabilities map[string]any `json:"capabilities"`
ServerInfo serverInfo `json:"serverInfo"`
}
// clientInfo identifies the client to the server.
type clientInfo struct {
Name string `json:"name"`
Version string `json:"version"`
}
// serverInfo identifies the server to the client.
type serverInfo struct {
Name string `json:"name"`
Version string `json:"version"`
}
// Connect selects the transport (unless one was injected), starts it, completes
// the MCP initialize handshake, and marks the client connected. It is
// idempotent; a second call while connected returns nil.
func (c *Client) Connect(ctx context.Context) error {
c.mu.Lock()
switch c.state {
case stateClosed:
c.mu.Unlock()
return fmt.Errorf("mcp connect %q: %w", c.cfg.Name, ErrClosed)
case stateConnected:
c.mu.Unlock()
return nil
}
c.mu.Unlock()
tr := c.tr
if tr == nil {
var err error
if tr, err = newTransport(c.cfg, c.log); err != nil {
return fmt.Errorf("mcp connect %q: %w", c.cfg.Name, err)
}
}
if err := tr.Start(ctx); err != nil {
tr.Close()
return fmt.Errorf("mcp connect %q: %w", c.cfg.Name, err)
}
info, err := c.initialize(ctx, tr)
if err != nil {
tr.Close()
return fmt.Errorf("mcp connect %q: %w", c.cfg.Name, err)
}
c.sendInitialized(ctx, tr)
c.mu.Lock()
c.tr = tr
c.info = info
c.state = stateConnected
c.mu.Unlock()
c.log.Debug("mcp connected", "name", c.cfg.Name, "server", info.Name)
return nil
}
// Initialize performs the MCP initialize handshake on an already-connected
// client. It returns the server identity.
func (c *Client) Initialize(ctx context.Context) (ServerInfo, error) {
tr := c.transport()
if tr == nil {
return ServerInfo{}, fmt.Errorf("mcp initialize %q: %w", c.cfg.Name, ErrNotConnected)
}
return c.initialize(ctx, tr)
}
// initialize sends the initialize request over tr and records the server
// identity.
func (c *Client) initialize(ctx context.Context, tr Transport) (ServerInfo, error) {
params := initializeParams{
ProtocolVersion: protocolVersion,
Capabilities: map[string]any{},
ClientInfo: clientInfo{Name: "cyrene", Version: "0.1.0"},
}
var res initializeResult
if err := tr.Call(ctx, "initialize", params, &res); err != nil {
return ServerInfo{}, err
}
return ServerInfo{Name: res.ServerInfo.Name, Version: res.ServerInfo.Version}, nil
}
// sendInitialized best-effort sends the notifications/initialized advisory ack.
func (c *Client) sendInitialized(ctx context.Context, tr Transport) {
n, ok := tr.(notifier)
if !ok {
return
}
if err := n.notify(ctx, "notifications/initialized", map[string]any{}); err != nil {
c.log.Debug("mcp initialized notification failed", "name", c.cfg.Name, "err", err)
}
}
// ListTools returns the tools exposed by the connected MCP server.
func (c *Client) ListTools(ctx context.Context) ([]ToolInfo, error) {
tr := c.transport()
if tr == nil {
return nil, fmt.Errorf("mcp list tools %q: %w", c.cfg.Name, ErrNotConnected)
}
var res listToolsResult
if err := tr.Call(ctx, "tools/list", map[string]any{}, &res); err != nil {
return nil, fmt.Errorf("mcp list tools %q: %w", c.cfg.Name, err)
}
out := make([]ToolInfo, 0, len(res.Tools))
for _, t := range res.Tools {
out = append(out, ToolInfo{Name: t.Name, Description: t.Description, InputSchema: t.InputSchema})
}
return out, nil
}
// CallTool invokes a tool on the connected MCP server.
func (c *Client) CallTool(ctx context.Context, name string, args any) (ToolResult, error) {
tr := c.transport()
if tr == nil {
return ToolResult{}, fmt.Errorf("mcp call tool %q: %w", name, ErrNotConnected)
}
arguments, err := normalizeArgs(args)
if err != nil {
return ToolResult{}, fmt.Errorf("mcp call tool %q: %w", name, err)
}
params := callToolParams{Name: name, Arguments: arguments}
var res callToolResult
if err := tr.Call(ctx, "tools/call", params, &res); err != nil {
return ToolResult{}, fmt.Errorf("mcp call tool %q: %w", name, err)
}
return ToolResult{Content: toContentBlocks(res.Content), IsError: res.IsError}, nil
}
// DiscoverTools lists the server tools and maps each to an sdk.Tool so the MCP
// server appears as a normal tool provider.
func (c *Client) DiscoverTools(ctx context.Context) ([]sdk.Tool, error) {
tools, err := c.ListTools(ctx)
if err != nil {
return nil, err
}
out := make([]sdk.Tool, 0, len(tools))
for _, t := range tools {
out = append(out, &mcpTool{client: c, info: t})
}
return out, nil
}
// Close closes the transport and marks the client closed. It is idempotent.
func (c *Client) Close() error {
c.mu.Lock()
if c.state == stateClosed {
c.mu.Unlock()
return nil
}
c.state = stateClosed
tr := c.tr
c.mu.Unlock()
if tr != nil {
return tr.Close()
}
return nil
}
// newTransport selects the transport matching cfg: stdio when Command is set, SSE
// when URL is set.
func newTransport(cfg MCPConfig, log *slog.Logger) (Transport, error) {
switch {
case cfg.Command != "":
return newStdioTransport(cfg, log), nil
case cfg.URL != "":
return newSSETransport(cfg, log), nil
default:
return nil, fmt.Errorf("mcp: %w", ErrNoTransport)
}
}
// normalizeArgs converts an agent-supplied argument value into the MCP arguments
// object, requiring a JSON object.
func normalizeArgs(args any) (map[string]any, error) {
if args == nil {
return map[string]any{}, nil
}
if m, ok := args.(map[string]any); ok {
return m, nil
}
return nil, fmt.Errorf("arguments must be a JSON object, got %T", args)
}
// listToolsResult is the MCP tools/list response body.
type listToolsResult struct {
Tools []toolSpec `json:"tools"`
}
// toolSpec is a single tool in a tools/list response.
type toolSpec struct {
Name string `json:"name"`
Description string `json:"description"`
InputSchema map[string]any `json:"inputSchema"`
}
// callToolParams is the MCP tools/call request body.
type callToolParams struct {
Name string `json:"name"`
Arguments any `json:"arguments,omitempty"`
}
// callToolResult is the MCP tools/call response body.
type callToolResult struct {
Content []contentBlock `json:"content"`
IsError bool `json:"isError"`
}
// contentBlock is a wire content block in a tools/call response.
type contentBlock struct {
Type string `json:"type"`
Text string `json:"text"`
}
+247
View File
@@ -0,0 +1,247 @@
package mcp
import (
"context"
"errors"
"testing"
)
func TestClient_New(t *testing.T) {
t.Parallel()
c := New(MCPConfig{Name: "calc", Command: "mcp-calc"})
if c.cfg.Name != "calc" {
t.Fatalf("cfg.Name = %q, want calc", c.cfg.Name)
}
if c.state != stateNew {
t.Fatalf("state = %v, want stateNew", c.state)
}
}
func TestClient_Connect_FakeTransport(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 ft.starts != 1 {
t.Fatalf("starts = %d, want 1", ft.starts)
}
if !ft.hasMethod("initialize") {
t.Fatalf("initialize was not called; methods = %v", ft.calledMethods())
}
if c.state != stateConnected {
t.Fatalf("state = %v, want stateConnected", c.state)
}
if c.info.Name != "fake-server" {
t.Fatalf("info.Name = %q, want fake-server", c.info.Name)
}
}
func TestClient_Connect_Idempotent(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("first Connect: %v", err)
}
if err := c.Connect(context.Background()); err != nil {
t.Fatalf("second Connect: %v", err)
}
if ft.starts != 1 {
t.Fatalf("starts = %d, want 1 (second connect must not restart)", ft.starts)
}
}
func TestClient_Connect_NoTransport(t *testing.T) {
t.Parallel()
c := New(MCPConfig{Name: "empty"})
err := c.Connect(context.Background())
if !errors.Is(err, ErrNoTransport) {
t.Fatalf("Connect error = %v, want ErrNoTransport", err)
}
}
func TestClient_Connect_Closed(t *testing.T) {
t.Parallel()
ft := &fakeTransport{}
c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft))
if err := c.Close(); err != nil {
t.Fatalf("Close: %v", err)
}
err := c.Connect(context.Background())
if !errors.Is(err, ErrClosed) {
t.Fatalf("Connect error = %v, want ErrClosed", err)
}
}
func TestClient_Connect_StartError(t *testing.T) {
t.Parallel()
ft := &fakeTransport{startErr: errors.New("boom")}
c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft))
if err := c.Connect(context.Background()); err == nil {
t.Fatal("Connect succeeded, want error")
}
if ft.closes != 1 {
t.Fatalf("closes = %d, want 1 (failed Start must Close the transport)", ft.closes)
}
}
func TestClient_Connect_InitError(t *testing.T) {
t.Parallel()
ft := &fakeTransport{initErr: errors.New("init boom")}
c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft))
if err := c.Connect(context.Background()); err == nil {
t.Fatal("Connect succeeded, want error")
}
if ft.closes != 1 {
t.Fatalf("closes = %d, want 1 (failed initialize must Close the transport)", ft.closes)
}
}
func TestClient_ListTools(t *testing.T) {
t.Parallel()
ft := &fakeTransport{
tools: []ToolInfo{{Name: "add", Description: "Adds", InputSchema: map[string]any{"type": "object"}}},
}
c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft))
if err := c.Connect(context.Background()); err != nil {
t.Fatalf("Connect: %v", err)
}
tools, err := c.ListTools(context.Background())
if err != nil {
t.Fatalf("ListTools: %v", err)
}
if len(tools) != 1 || tools[0].Name != "add" || tools[0].Description != "Adds" {
t.Fatalf("tools = %+v, want [add/Adds]", tools)
}
}
func TestClient_ListTools_NotConnected(t *testing.T) {
t.Parallel()
c := New(MCPConfig{Name: "s", Command: "x"})
_, err := c.ListTools(context.Background())
if !errors.Is(err, ErrNotConnected) {
t.Fatalf("ListTools error = %v, want ErrNotConnected", err)
}
}
func TestClient_CallTool(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)
}
res, err := c.CallTool(context.Background(), "add", map[string]any{"x": 1})
if err != nil {
t.Fatalf("CallTool: %v", err)
}
if res.IsError {
t.Fatal("IsError = true, want false")
}
if len(res.Content) != 1 || res.Content[0].Text != "hello" {
t.Fatalf("Content = %+v, want [hello]", res.Content)
}
}
func TestClient_CallTool_BadArgs(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)
}
_, err := c.CallTool(context.Background(), "add", "not-a-map")
if err == nil {
t.Fatal("CallTool with non-object args succeeded, want error")
}
}
func TestClient_CallTool_NotConnected(t *testing.T) {
t.Parallel()
c := New(MCPConfig{Name: "s", Command: "x"})
_, err := c.CallTool(context.Background(), "add", map[string]any{})
if !errors.Is(err, ErrNotConnected) {
t.Fatalf("CallTool error = %v, want ErrNotConnected", err)
}
}
func TestClient_DiscoverTools_MapsToSDK(t *testing.T) {
t.Parallel()
ft := &fakeTransport{
tools: []ToolInfo{{Name: "add", Description: "Adds", InputSchema: map[string]any{"type": "object"}}},
}
c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft))
if err := c.Connect(context.Background()); err != nil {
t.Fatalf("Connect: %v", err)
}
tools, err := c.DiscoverTools(context.Background())
if err != nil {
t.Fatalf("DiscoverTools: %v", err)
}
if len(tools) != 1 {
t.Fatalf("len = %d, want 1", len(tools))
}
if tools[0].Name() != "add" || tools[0].Description() != "Adds" {
t.Fatalf("tool = %q/%q, want add/Adds", tools[0].Name(), tools[0].Description())
}
params, err := tools[0].Parameters()
if err != nil {
t.Fatalf("Parameters: %v", err)
}
if params["type"] != "object" {
t.Fatalf("params[type] = %v, want object", params["type"])
}
out, err := tools[0].Execute(context.Background(), map[string]any{"x": 1})
if err != nil {
t.Fatalf("Execute: %v", err)
}
if out != "hello" {
t.Fatalf("Execute result = %v, want hello", out)
}
}
func TestClient_DiscoverTools_NilSchema(t *testing.T) {
t.Parallel()
ft := &fakeTransport{tools: []ToolInfo{{Name: "nop", InputSchema: nil}}}
c := New(MCPConfig{Name: "s", Command: "x"}, WithTransport(ft))
if err := c.Connect(context.Background()); err != nil {
t.Fatalf("Connect: %v", err)
}
tools, err := c.DiscoverTools(context.Background())
if err != nil {
t.Fatalf("DiscoverTools: %v", err)
}
params, err := tools[0].Parameters()
if err != nil {
t.Fatalf("Parameters: %v", err)
}
if params["type"] != "object" {
t.Fatalf("params[type] = %v, want object", params["type"])
}
}
func TestClient_Close_ClosesTransportAndForbidsUse(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 err := c.Close(); err != nil {
t.Fatalf("Close: %v", err)
}
if ft.closes != 1 {
t.Fatalf("closes = %d, want 1", ft.closes)
}
if err := c.Close(); err != nil {
t.Fatalf("second Close: %v", err)
}
if ft.closes != 1 {
t.Fatalf("closes = %d, want 1 after idempotent Close", ft.closes)
}
if _, err := c.ListTools(context.Background()); !errors.Is(err, ErrNotConnected) {
t.Fatalf("ListTools after Close = %v, want ErrNotConnected", err)
}
}
+117
View File
@@ -0,0 +1,117 @@
package mcp
import (
"fmt"
"os"
"strconv"
"strings"
"gopkg.in/yaml.v3"
)
// Load reads MCP server configuration from the process environment. It returns
// every configured server; callers filter on Enabled.
func Load() (*Config, error) {
return FromEnv(os.Environ())
}
// FromEnv builds a Config from a KEY=value pair slice such as os.Environ() or a
// test fixture. Named servers are read from MCP_<NAME>_* for indexes; the
// unnumbered MCP_* variables describe a single default server. When N is a
// positive index, one server is read per group and scanning stops at the first
// group with no NAME/COMMAND/URL set.
func FromEnv(pairs []string) (*Config, error) {
env := make(map[string]string, len(pairs))
for _, p := range pairs {
k, v, ok := strings.Cut(p, "=")
if ok {
env[k] = v
}
}
return &Config{Servers: fromEnvMap(env)}, nil
}
// FromYAML parses a YAML document describing MCP servers, for example:
//
// servers:
// - name: calc
// command: mcp-calc
// args: [--stdio]
// enabled: true
func FromYAML(data []byte) (*Config, error) {
var doc struct {
Servers []MCPConfig `yaml:"servers"`
}
if err := yaml.Unmarshal(data, &doc); err != nil {
return nil, fmt.Errorf("mcp config: %w", err)
}
return &Config{Servers: doc.Servers}, nil
}
// fromEnvMap builds the server list from a resolved environment map.
func fromEnvMap(env map[string]string) []MCPConfig {
var servers []MCPConfig
if cfg, ok := readOne(env, "MCP"); ok {
servers = append(servers, cfg)
}
for i := 0; ; i++ {
cfg, ok := readOne(env, "MCP_"+strconv.Itoa(i))
if !ok {
break
}
servers = append(servers, cfg)
}
return servers
}
// readOne reads a single MCPConfig from env using the given variable prefix
// (e.g. "MCP" or "MCP_0"). It reports ok=false when no identifying field is set.
func readOne(env map[string]string, prefix string) (MCPConfig, bool) {
p := prefix
if p == "" {
p = "MCP"
}
name := env[p+"_NAME"]
command := env[p+"_COMMAND"]
url := env[p+"_URL"]
if name == "" && command == "" && url == "" {
return MCPConfig{}, false
}
return MCPConfig{
Name: name,
Command: command,
Args: splitArgs(env[p+"_ARGS"]),
URL: url,
Enabled: parseBool(env[p+"_ENABLED"], false),
}, true
}
// splitArgs splits a comma-separated args value, trimming spaces and empty
// entries.
func splitArgs(raw string) []string {
raw = strings.TrimSpace(raw)
if raw == "" {
return nil
}
parts := strings.Split(raw, ",")
out := make([]string, 0, len(parts))
for _, p := range parts {
if t := strings.TrimSpace(p); t != "" {
out = append(out, t)
}
}
return out
}
// parseBool parses a boolean value, defaulting to def on empty or invalid input.
func parseBool(raw string, def bool) bool {
raw = strings.TrimSpace(raw)
if raw == "" {
return def
}
b, err := strconv.ParseBool(strings.ToLower(raw))
if err != nil {
return def
}
return b
}
+158
View File
@@ -0,0 +1,158 @@
package mcp
import (
"reflect"
"testing"
)
func TestFromEnv_DefaultServer(t *testing.T) {
t.Parallel()
pairs := []string{
"MCP_NAME=calc",
"MCP_COMMAND=mcp-calc",
"MCP_ARGS=--stdio,--transport",
"MCP_URL=",
"MCP_ENABLED=true",
}
cfg, err := FromEnv(pairs)
if err != nil {
t.Fatalf("FromEnv: %v", err)
}
if len(cfg.Servers) != 1 {
t.Fatalf("servers = %d, want 1", len(cfg.Servers))
}
s := cfg.Servers[0]
if s.Name != "calc" || s.Command != "mcp-calc" || !s.Enabled {
t.Fatalf("server = %+v, want calc/mcp-calc/enabled", s)
}
if !reflect.DeepEqual(s.Args, []string{"--stdio", "--transport"}) {
t.Fatalf("Args = %v, want [--stdio --transport]", s.Args)
}
}
func TestFromEnv_MultipleServers(t *testing.T) {
t.Parallel()
pairs := []string{
"MCP_0_NAME=calc",
"MCP_0_COMMAND=mcp-calc",
"MCP_0_ENABLED=true",
"MCP_1_NAME=web",
"MCP_1_URL=http://localhost:9090/sse",
"MCP_1_ENABLED=false",
}
cfg, err := FromEnv(pairs)
if err != nil {
t.Fatalf("FromEnv: %v", err)
}
if len(cfg.Servers) != 2 {
t.Fatalf("servers = %d, want 2", len(cfg.Servers))
}
if cfg.Servers[1].Name != "web" || cfg.Servers[1].URL != "http://localhost:9090/sse" || cfg.Servers[1].Enabled {
t.Fatalf("server[1] = %+v, want web/url/disabled", cfg.Servers[1])
}
}
func TestFromEnv_NoServers(t *testing.T) {
t.Parallel()
cfg, err := FromEnv([]string{"MCP_ENABLED=true"})
if err != nil {
t.Fatalf("FromEnv: %v", err)
}
if len(cfg.Servers) != 0 {
t.Fatalf("servers = %d, want 0 (ENABLED alone is not a server)", len(cfg.Servers))
}
}
func TestFromEnv_Empty(t *testing.T) {
t.Parallel()
cfg, err := FromEnv(nil)
if err != nil {
t.Fatalf("FromEnv: %v", err)
}
if len(cfg.Servers) != 0 {
t.Fatalf("servers = %d, want 0", len(cfg.Servers))
}
}
func TestFromYAML(t *testing.T) {
t.Parallel()
data := []byte(`
servers:
- name: calc
command: mcp-calc
args: [--stdio]
enabled: true
- name: web
url: http://localhost:9090/sse
`)
cfg, err := FromYAML(data)
if err != nil {
t.Fatalf("FromYAML: %v", err)
}
if len(cfg.Servers) != 2 {
t.Fatalf("servers = %d, want 2", len(cfg.Servers))
}
if cfg.Servers[0].Name != "calc" || cfg.Servers[0].Command != "mcp-calc" || !cfg.Servers[0].Enabled {
t.Fatalf("server[0] = %+v, want calc/mcp-calc/enabled", cfg.Servers[0])
}
if !reflect.DeepEqual(cfg.Servers[0].Args, []string{"--stdio"}) {
t.Fatalf("Args = %v, want [--stdio]", cfg.Servers[0].Args)
}
if cfg.Servers[1].URL != "http://localhost:9090/sse" {
t.Fatalf("server[1].URL = %q, want sse url", cfg.Servers[1].URL)
}
}
func TestFromYAML_Invalid(t *testing.T) {
t.Parallel()
if _, err := FromYAML([]byte(`servers: [`)); err == nil {
t.Fatal("FromYAML with malformed YAML succeeded, want error")
}
}
func TestSplitArgs(t *testing.T) {
t.Parallel()
tests := []struct {
in string
want []string
}{
{in: "--stdio,--transport", want: []string{"--stdio", "--transport"}},
{in: " a , b ,, c ", want: []string{"a", "b", "c"}},
{in: "", want: nil},
{in: " ", want: nil},
}
for _, tt := range tests {
tt := tt
t.Run(tt.in, func(t *testing.T) {
t.Parallel()
if got := splitArgs(tt.in); !reflect.DeepEqual(got, tt.want) {
t.Fatalf("splitArgs(%q) = %v, want %v", tt.in, got, tt.want)
}
})
}
}
func TestParseBool(t *testing.T) {
t.Parallel()
tests := []struct {
in string
def bool
want bool
}{
{in: "true", def: false, want: true},
{in: "false", def: true, want: false},
{in: "TRUE", def: false, want: true},
{in: "1", def: false, want: true},
{in: "", def: true, want: true},
{in: "garbage", def: false, want: false},
}
for _, tt := range tests {
tt := tt
t.Run(tt.in, func(t *testing.T) {
t.Parallel()
if got := parseBool(tt.in, tt.def); got != tt.want {
t.Fatalf("parseBool(%q, %v) = %v, want %v", tt.in, tt.def, got, tt.want)
}
})
}
}
+99
View File
@@ -0,0 +1,99 @@
package mcp
import (
"context"
"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.
type fakeTransport struct {
startErr error
closeErr error
initErr error
listErr error
callErr error
tools []ToolInfo
mu sync.Mutex
starts int
closes int
methods []string
}
// Start records a start.
func (f *fakeTransport) Start(_ context.Context) error {
f.mu.Lock()
defer f.mu.Unlock()
f.starts++
return f.startErr
}
// Call answers the three MCP methods the client issues.
func (f *fakeTransport) Call(_ context.Context, method string, _ any, result any) error {
f.mu.Lock()
defer f.mu.Unlock()
f.methods = append(f.methods, method)
switch method {
case "initialize":
if f.initErr != nil {
return f.initErr
}
if r, ok := result.(*initializeResult); ok {
r.ServerInfo = serverInfo{Name: "fake-server", Version: "1.0"}
r.ProtocolVersion = protocolVersion
}
case "tools/list":
if f.listErr != nil {
return f.listErr
}
if r, ok := result.(*listToolsResult); ok {
for _, t := range f.tools {
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"}}
}
}
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
}
// 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
}
+209
View File
@@ -0,0 +1,209 @@
// 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
}
+81
View File
@@ -0,0 +1,81 @@
package mcp
import (
"context"
"errors"
"fmt"
"log/slog"
"sync"
"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
)
// Registry connects an MCP client and auto-registers its discovered tools into
// the shared capability registry, making them available to every character.
type Registry struct {
reg *registry.Registry
mcp *Client
log *slog.Logger
mu sync.Mutex
names []string
closed bool
}
// NewRegistry builds a Registry that registers cfg's tools into reg. Options
// apply to the underlying Client (logger, injected transport). It is named
// NewRegistry rather than New because this package already exports New for
// Client.
func NewRegistry(reg *registry.Registry, cfg MCPConfig, opts ...Option) *Registry {
c := New(cfg, opts...)
return &Registry{reg: reg, mcp: c, log: c.log}
}
// Register connects the MCP server, discovers its tools, and registers each one
// into the underlying registry.
func (r *Registry) Register(ctx context.Context) error {
if r.reg == nil {
return fmt.Errorf("mcp registry: nil registry")
}
if err := r.mcp.Connect(ctx); err != nil {
return fmt.Errorf("mcp registry %q: %w", r.mcp.cfg.Name, err)
}
tools, err := r.mcp.DiscoverTools(ctx)
if err != nil {
return fmt.Errorf("mcp registry %q: %w", r.mcp.cfg.Name, err)
}
r.mu.Lock()
defer r.mu.Unlock()
if r.closed {
return fmt.Errorf("mcp registry %q: %w", r.mcp.cfg.Name, ErrClosed)
}
for _, t := range tools {
if err := r.reg.RegisterTool(t); err != nil {
return fmt.Errorf("mcp register %q: %w", t.Name(), err)
}
r.names = append(r.names, t.Name())
}
r.log.Info("mcp tools registered", "name", r.mcp.cfg.Name, "count", len(tools))
return nil
}
// Unregister removes every tool this registry previously registered and closes
// the underlying client. It is safe to call multiple times.
func (r *Registry) Unregister(ctx context.Context) error {
r.mu.Lock()
names := r.names
r.names = nil
r.closed = true
r.mu.Unlock()
var errs error
for _, name := range names {
if err := r.reg.Remove(name); err != nil && !errors.Is(err, registry.ErrNotFound) {
errs = errors.Join(errs, fmt.Errorf("mcp unregister %q: %w", name, err))
}
}
if err := r.mcp.Close(); err != nil {
errs = errors.Join(errs, fmt.Errorf("mcp close %q: %w", r.mcp.cfg.Name, err))
}
return errs
}
+118
View File
@@ -0,0 +1,118 @@
package mcp
import (
"context"
"errors"
"testing"
"git.yeij.top/AskaEth/Cyrene_v2/internal/capability/registry"
)
func sampleTools() []ToolInfo {
return []ToolInfo{
{Name: "add", Description: "Adds numbers", InputSchema: map[string]any{"type": "object"}},
{Name: "sub", Description: "Subtracts numbers", InputSchema: map[string]any{"type": "object"}},
}
}
func TestRegistry_Register_AutoRegisters(t *testing.T) {
t.Parallel()
ft := &fakeTransport{tools: sampleTools()}
reg := registry.New()
r := NewRegistry(reg, MCPConfig{Name: "calc", Command: "x"}, WithTransport(ft))
if err := r.Register(context.Background()); err != nil {
t.Fatalf("Register: %v", err)
}
for _, name := range []string{"add", "sub"} {
cap, err := reg.Get(name)
if err != nil {
t.Fatalf("Get(%q): %v", name, err)
}
if cap.Kind != registry.KindTool {
t.Fatalf("Kind(%q) = %q, want %q", name, cap.Kind, registry.KindTool)
}
}
if !ft.hasMethod("tools/list") {
t.Fatalf("tools/list was not called; methods = %v", ft.calledMethods())
}
}
func TestRegistry_Register_PropagatesConnectError(t *testing.T) {
t.Parallel()
ft := &fakeTransport{startErr: errors.New("boom")}
reg := registry.New()
r := NewRegistry(reg, MCPConfig{Name: "calc", Command: "x"}, WithTransport(ft))
if err := r.Register(context.Background()); err == nil {
t.Fatal("Register succeeded, want connect error")
}
}
func TestRegistry_Register_DuplicateTool(t *testing.T) {
t.Parallel()
ft := &fakeTransport{tools: []ToolInfo{{Name: "add", InputSchema: map[string]any{}}}}
reg := registry.New()
// Pre-register a tool with the same name so the MCP registration collides.
_ = reg.RegisterTool(&testTool{name: "add"})
r := NewRegistry(reg, MCPConfig{Name: "calc", Command: "x"}, WithTransport(ft))
err := r.Register(context.Background())
if !errors.Is(err, registry.ErrDuplicate) {
t.Fatalf("Register error = %v, want registry.ErrDuplicate", err)
}
}
func TestRegistry_Register_NilRegistry(t *testing.T) {
t.Parallel()
ft := &fakeTransport{}
r := NewRegistry(nil, MCPConfig{Name: "calc", Command: "x"}, WithTransport(ft))
if err := r.Register(context.Background()); err == nil {
t.Fatal("Register with nil registry succeeded, want error")
}
}
func TestRegistry_Unregister_RemovesTools(t *testing.T) {
t.Parallel()
ft := &fakeTransport{tools: sampleTools()}
reg := registry.New()
r := NewRegistry(reg, MCPConfig{Name: "calc", Command: "x"}, WithTransport(ft))
if err := r.Register(context.Background()); err != nil {
t.Fatalf("Register: %v", err)
}
if err := r.Unregister(context.Background()); err != nil {
t.Fatalf("Unregister: %v", err)
}
for _, name := range []string{"add", "sub"} {
if _, err := reg.Get(name); !errors.Is(err, registry.ErrNotFound) {
t.Fatalf("Get(%q) after Unregister = %v, want ErrNotFound", name, err)
}
}
if ft.closes != 1 {
t.Fatalf("closes = %d, want 1 (Unregister must close the client)", ft.closes)
}
}
func TestRegistry_Unregister_Idempotent(t *testing.T) {
t.Parallel()
ft := &fakeTransport{}
reg := registry.New()
r := NewRegistry(reg, MCPConfig{Name: "calc", Command: "x"}, WithTransport(ft))
if err := r.Register(context.Background()); err != nil {
t.Fatalf("Register: %v", err)
}
if err := r.Unregister(context.Background()); err != nil {
t.Fatalf("first Unregister: %v", err)
}
if err := r.Unregister(context.Background()); err != nil {
t.Fatalf("second Unregister: %v", err)
}
}
// testTool is a minimal sdk.Tool used to pre-register a colliding name.
type testTool struct {
name string
}
func (t *testTool) Name() string { return t.name }
func (t *testTool) Description() string { return "desc " + t.name }
func (t *testTool) Parameters() (map[string]any, error) { return map[string]any{}, nil }
func (t *testTool) Execute(_ context.Context, _ any) (any, error) { return nil, nil }
+50
View File
@@ -0,0 +1,50 @@
package mcp
import (
"encoding/json"
"fmt"
)
// rpcRequest is a JSON-RPC 2.0 request. Notifications omit ID so the field is
// tagged omitempty. Params is any so callers pass a struct or a JSON object.
type rpcRequest struct {
JSONRPC string `json:"jsonrpc"`
ID any `json:"id,omitempty"`
Method string `json:"method"`
Params any `json:"params,omitempty"`
}
// rpcResponse is a JSON-RPC 2.0 response. Result is kept raw until the caller
// decodes it into the concrete result type.
type rpcResponse struct {
JSONRPC string `json:"jsonrpc"`
ID any `json:"id,omitempty"`
Result json.RawMessage `json:"result,omitempty"`
Error *rpcError `json:"error,omitempty"`
}
// rpcError is a JSON-RPC error object. It satisfies error so it can be wrapped
// and surfaced with %w.
type rpcError struct {
Code int `json:"code"`
Message string `json:"message"`
Data any `json:"data,omitempty"`
}
// Error implements error.
func (e *rpcError) Error() string {
return fmt.Sprintf("code %d: %s", e.Code, e.Message)
}
// encodeRequest marshals a JSON-RPC request. An id of nil produces a
// notification (no id field).
func encodeRequest(id any, method string, params any) ([]byte, error) {
return json.Marshal(rpcRequest{JSONRPC: "2.0", ID: id, Method: method, Params: params})
}
// decodeIDString reports whether res carries the string id expected for a JSON
// string id request.
func decodeIDString(res rpcResponse, id string) bool {
s, ok := res.ID.(string)
return ok && s == id
}
+250
View File
@@ -0,0 +1,250 @@
package mcp
import (
"bufio"
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"log/slog"
"net/http"
"strconv"
"strings"
"sync"
)
// sseTransport speaks JSON-RPC 2.0 to an MCP server over SSE. It opens a
// long-lived GET stream for server-to-client messages and POSTs client-to-server
// requests to the message endpoint (discovered via an endpoint event, falling
// back to the configured URL).
type sseTransport struct {
cfg MCPConfig
log *slog.Logger
url string
endpoint string
mu sync.Mutex
nextID int
pending map[string]chan []byte
resp *http.Response
body io.ReadCloser
closeCh chan struct{}
closeOnce sync.Once
started bool
}
// maxSSEResponse bounds how much of a POST reply body is read before falling
// back to the SSE stream.
const maxSSEResponse = 1 << 20
// newSSETransport builds an SSE transport for cfg.
func newSSETransport(cfg MCPConfig, log *slog.Logger) *sseTransport {
return &sseTransport{
cfg: cfg,
log: log,
url: cfg.URL,
endpoint: cfg.URL,
pending: make(map[string]chan []byte),
closeCh: make(chan struct{}),
}
}
// Start opens the SSE stream and launches the reader goroutine.
func (s *sseTransport) Start(ctx context.Context) error {
if err := ctx.Err(); err != nil {
return fmt.Errorf("mcp sse start: %w", err)
}
if s.cfg.URL == "" {
return fmt.Errorf("mcp sse: %w", ErrNoTransport)
}
req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, s.cfg.URL, nil)
if err != nil {
return fmt.Errorf("mcp sse get: %w", err)
}
req.Header.Set("Accept", "text/event-stream")
resp, err := http.DefaultClient.Do(req)
if err != nil {
return fmt.Errorf("mcp sse connect %s: %w", s.cfg.URL, err)
}
if resp.StatusCode != http.StatusOK {
resp.Body.Close()
return fmt.Errorf("mcp sse connect %s: status %d", s.cfg.URL, resp.StatusCode)
}
s.resp = resp
s.body = resp.Body
s.started = true
go s.readLoop()
return nil
}
// readLoop parses SSE events and dispatches message events to pending calls.
func (s *sseTransport) readLoop() {
r := bufio.NewReader(s.body)
var event, data string
for {
line, err := r.ReadString('\n')
if err != nil {
return
}
line = strings.TrimRight(line, "\r\n")
switch {
case line == "":
if data != "" {
s.dispatch(event, data)
}
event, data = "", ""
case strings.HasPrefix(line, "event:"):
event = strings.TrimSpace(strings.TrimPrefix(line, "event:"))
case strings.HasPrefix(line, "data:"):
d := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
if data == "" {
data = d
} else {
data += "\n" + d
}
}
}
}
// dispatch handles a single SSE event.
func (s *sseTransport) dispatch(event, data string) {
switch event {
case "endpoint":
if e := strings.TrimSpace(data); e != "" {
s.mu.Lock()
s.endpoint = e
s.mu.Unlock()
}
case "message":
var resp rpcResponse
if err := json.Unmarshal([]byte(data), &resp); err != nil {
return
}
id, ok := resp.ID.(string)
if !ok {
return
}
s.mu.Lock()
ch, has := s.pending[id]
delete(s.pending, id)
s.mu.Unlock()
if has {
ch <- []byte(data)
}
}
}
// Call POSTs a request and waits for a response, either from the POST reply body
// or from the SSE stream.
func (s *sseTransport) Call(ctx context.Context, method string, params any, result any) error {
s.mu.Lock()
if !s.started {
s.mu.Unlock()
return fmt.Errorf("mcp sse %s: %w", method, ErrNotConnected)
}
s.nextID++
id := strconv.Itoa(s.nextID)
ch := make(chan []byte, 1)
s.pending[id] = ch
endpoint := s.endpoint
s.mu.Unlock()
defer func() {
s.mu.Lock()
delete(s.pending, id)
s.mu.Unlock()
}()
reqBody, err := encodeRequest(id, method, params)
if err != nil {
return fmt.Errorf("mcp sse %s: %w", method, err)
}
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(reqBody))
if err != nil {
return fmt.Errorf("mcp sse %s: %w", method, err)
}
httpReq.Header.Set("Content-Type", "application/json")
httpReq.Header.Set("Accept", "application/json")
resp, err := http.DefaultClient.Do(httpReq)
if err != nil {
return fmt.Errorf("mcp sse %s post: %w", method, err)
}
body, _ := io.ReadAll(io.LimitReader(resp.Body, maxSSEResponse))
resp.Body.Close()
var direct rpcResponse
if json.Unmarshal(bytes.TrimSpace(body), &direct) == nil && decodeIDString(direct, id) {
return s.applyResponse(direct, method, result)
}
select {
case data := <-ch:
var rr rpcResponse
if err := json.Unmarshal(data, &rr); err != nil {
return fmt.Errorf("mcp sse %s decode: %w", method, err)
}
return s.applyResponse(rr, method, result)
case <-ctx.Done():
return fmt.Errorf("mcp sse %s: %w", method, ctx.Err())
case <-s.closeCh:
return fmt.Errorf("mcp sse %s: %w", method, ErrClosed)
}
}
// applyResponse decodes a JSON-RPC response into result, surfacing an error
// object if present.
func (s *sseTransport) applyResponse(resp rpcResponse, method string, result any) error {
if resp.Error != nil {
return fmt.Errorf("mcp sse %s: %w: %w", method, ErrRPC, resp.Error)
}
if result != nil {
if err := json.Unmarshal(resp.Result, result); err != nil {
return fmt.Errorf("mcp sse %s decode: %w", method, err)
}
}
return nil
}
// notify POSTs a fire-and-forget JSON-RPC notification.
func (s *sseTransport) notify(ctx context.Context, method string, params any) error {
s.mu.Lock()
started := s.started
endpoint := s.endpoint
s.mu.Unlock()
if !started {
return fmt.Errorf("mcp sse %s: %w", method, ErrNotConnected)
}
reqBody, err := encodeRequest(nil, method, params)
if err != nil {
return fmt.Errorf("mcp notify %s: %w", method, err)
}
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(reqBody))
if err != nil {
return fmt.Errorf("mcp notify %s: %w", method, err)
}
httpReq.Header.Set("Content-Type", "application/json")
resp, err := http.DefaultClient.Do(httpReq)
if err != nil {
return fmt.Errorf("mcp notify %s: %w", method, err)
}
defer resp.Body.Close()
_, _ = io.Copy(io.Discard, resp.Body)
return nil
}
// Close closes the SSE stream and any pending responses.
func (s *sseTransport) Close() error {
s.closeOnce.Do(func() {
close(s.closeCh)
if s.body != nil {
_ = s.body.Close()
}
if s.resp != nil && s.resp.Body != nil {
_ = s.resp.Body.Close()
}
s.mu.Lock()
s.pending = make(map[string]chan []byte)
s.mu.Unlock()
})
return nil
}
+203
View File
@@ -0,0 +1,203 @@
package mcp
import (
"context"
"encoding/json"
"errors"
"io"
"net/http"
"net/http/httptest"
"testing"
"time"
)
// postRespondServer serves a GET stream that ends immediately and answers every
// POST directly with a JSON-RPC response body (the non-streaming reply path).
func postRespondServer(t *testing.T) *httptest.Server {
t.Helper()
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method == http.MethodGet {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
if f, ok := w.(http.Flusher); ok {
f.Flush()
}
return
}
body, _ := io.ReadAll(r.Body)
var req rpcRequest
_ = json.Unmarshal(body, &req)
resp := `{"jsonrpc":"2.0","id":` + idJSON(req.ID) + `,"result":{"content":[{"type":"text","text":"direct"}]}}`
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(resp))
}))
}
// streamServer pushes responses to the POST caller back through the SSE stream.
func streamServer(t *testing.T) *httptest.Server {
t.Helper()
push := make(chan string, 8)
var srv *httptest.Server
srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.Method {
case http.MethodGet:
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("Cache-Control", "no-cache")
f, _ := w.(http.Flusher)
_, _ = w.Write([]byte("event: endpoint\ndata: " + srv.URL + "\n\n"))
f.Flush()
for {
select {
case data := <-push:
_, _ = w.Write([]byte("event: message\ndata: " + data + "\n\n"))
f.Flush()
case <-r.Context().Done():
return
}
}
default:
body, _ := io.ReadAll(r.Body)
var req rpcRequest
_ = json.Unmarshal(body, &req)
resp := `{"jsonrpc":"2.0","id":` + idJSON(req.ID) + `,"result":{"content":[{"type":"text","text":"stream"}]}}`
push <- resp
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte("ok"))
}
}))
return srv
}
// idJSON renders a request id as a JSON string literal when it is a string.
func idJSON(id any) string {
if s, ok := id.(string); ok {
b, _ := json.Marshal(s)
return string(b)
}
return `null`
}
func TestSSE_Call_DirectResponse(t *testing.T) {
t.Parallel()
srv := postRespondServer(t)
defer srv.Close()
tr := newSSETransport(MCPConfig{Name: "s", URL: srv.URL}, defaultLogger())
defer tr.Close()
if err := tr.Start(context.Background()); err != nil {
t.Fatalf("Start: %v", err)
}
var res callToolResult
if err := tr.Call(context.Background(), "tools/call", callToolParams{Name: "add", Arguments: map[string]any{}}, &res); err != nil {
t.Fatalf("Call: %v", err)
}
if len(res.Content) != 1 || res.Content[0].Text != "direct" {
t.Fatalf("Content = %+v, want [direct]", res.Content)
}
}
func TestSSE_Call_MessageEvent(t *testing.T) {
t.Parallel()
srv := streamServer(t)
defer srv.Close()
tr := newSSETransport(MCPConfig{Name: "s", URL: srv.URL}, defaultLogger())
defer tr.Close()
if err := tr.Start(context.Background()); err != nil {
t.Fatalf("Start: %v", err)
}
var res callToolResult
if err := tr.Call(context.Background(), "tools/call", callToolParams{Name: "add", Arguments: map[string]any{}}, &res); err != nil {
t.Fatalf("Call: %v", err)
}
if len(res.Content) != 1 || res.Content[0].Text != "stream" {
t.Fatalf("Content = %+v, want [stream]", res.Content)
}
}
func TestSSE_Start_Non200(t *testing.T) {
t.Parallel()
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusInternalServerError)
}))
defer srv.Close()
tr := newSSETransport(MCPConfig{Name: "s", URL: srv.URL}, defaultLogger())
if err := tr.Start(context.Background()); err == nil {
t.Fatal("Start succeeded on a non-200 response, want error")
}
}
func TestSSE_Call_NotStarted(t *testing.T) {
t.Parallel()
tr := newSSETransport(MCPConfig{Name: "s", URL: "http://example.invalid"}, defaultLogger())
err := tr.Call(context.Background(), "tools/list", nil, nil)
if !errors.Is(err, ErrNotConnected) {
t.Fatalf("Call error = %v, want ErrNotConnected", err)
}
}
func TestSSE_Call_ContextTimeout(t *testing.T) {
t.Parallel()
// A server whose GET stream stays open and whose POST replies with no body
// and never pushes forces the Call to wait on its context deadline.
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method == http.MethodGet {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
if f, ok := w.(http.Flusher); ok {
f.Flush()
}
<-r.Context().Done()
return
}
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte("ok"))
}))
defer srv.Close()
tr := newSSETransport(MCPConfig{Name: "s", URL: srv.URL}, defaultLogger())
if err := tr.Start(context.Background()); err != nil {
t.Fatalf("Start: %v", err)
}
defer tr.Close()
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
defer cancel()
err := tr.Call(ctx, "tools/call", callToolParams{Name: "add"}, &callToolResult{})
if !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("Call error = %v, want context.DeadlineExceeded", err)
}
}
func TestSSE_Start_NoURL(t *testing.T) {
t.Parallel()
tr := newSSETransport(MCPConfig{Name: "s"}, defaultLogger())
err := tr.Start(context.Background())
if !errors.Is(err, ErrNoTransport) {
t.Fatalf("Start error = %v, want ErrNoTransport", err)
}
}
func TestSSE_Close_ClosesPendingAndBody(t *testing.T) {
t.Parallel()
srv := postRespondServer(t)
defer srv.Close()
tr := newSSETransport(MCPConfig{Name: "s", URL: srv.URL}, defaultLogger())
if err := tr.Start(context.Background()); err != nil {
t.Fatalf("Start: %v", err)
}
// Put a sentinel pending call in place.
tr.mu.Lock()
tr.pending["999"] = make(chan []byte, 1)
tr.mu.Unlock()
if err := tr.Close(); err != nil {
t.Fatalf("Close: %v", err)
}
tr.mu.Lock()
empty := len(tr.pending) == 0
tr.mu.Unlock()
if !empty {
t.Fatalf("pending not cleared after Close")
}
if err := tr.Close(); err != nil {
t.Fatalf("second Close: %v", err)
}
}
+184
View File
@@ -0,0 +1,184 @@
package mcp
import (
"bufio"
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"log/slog"
"os/exec"
"strconv"
"strings"
"sync"
)
// stdioTransport runs an MCP server as a child process and speaks
// newline-delimited JSON-RPC 2.0 to its stdin/stdout.
type stdioTransport struct {
cfg MCPConfig
log *slog.Logger
line *lineRPC
cmd *exec.Cmd
stdin io.WriteCloser
closeOnce sync.Once
closed bool
}
// newStdioTransport builds a stdio transport for cfg.
func newStdioTransport(cfg MCPConfig, log *slog.Logger) *stdioTransport {
return &stdioTransport{cfg: cfg, log: log}
}
// Start spawns the subprocess and wires up the request/response streams.
func (s *stdioTransport) Start(ctx context.Context) error {
if err := ctx.Err(); err != nil {
return fmt.Errorf("mcp stdio start: %w", err)
}
if s.cfg.Command == "" {
return fmt.Errorf("mcp stdio: %w", ErrNoTransport)
}
cmd := exec.Command(s.cfg.Command, s.cfg.Args...)
stdin, err := cmd.StdinPipe()
if err != nil {
return fmt.Errorf("mcp stdio stdin: %w", err)
}
stdout, err := cmd.StdoutPipe()
if err != nil {
return fmt.Errorf("mcp stdio stdout: %w", err)
}
stderr, err := cmd.StderrPipe()
if err != nil {
return fmt.Errorf("mcp stdio stderr: %w", err)
}
if err := cmd.Start(); err != nil {
return fmt.Errorf("mcp stdio start %q: %w", s.cfg.Command, err)
}
s.cmd = cmd
s.stdin = stdin
s.line = &lineRPC{in: stdin, out: bufio.NewReader(stdout), log: s.log}
go s.drainStderr(stderr)
return nil
}
// Call sends a request over stdio and decodes the response.
func (s *stdioTransport) Call(ctx context.Context, method string, params any, result any) error {
if s.line == nil {
return fmt.Errorf("mcp stdio %s: %w", method, ErrNotConnected)
}
return s.line.call(ctx, method, params, result)
}
// Close terminates the child process and closes its streams.
func (s *stdioTransport) Close() error {
s.closeOnce.Do(func() {
s.closed = true
if s.stdin != nil {
_ = s.stdin.Close()
}
if s.cmd != nil && s.cmd.Process != nil {
_ = s.cmd.Process.Kill()
}
if s.cmd != nil {
_ = s.cmd.Wait()
}
})
return nil
}
// drainStderr reads and logs the child's standard error at debug level so server
// diagnostics are visible without blocking the process.
func (s *stdioTransport) drainStderr(r io.Reader) {
buf := make([]byte, 4096)
for {
n, err := r.Read(buf)
if n > 0 {
s.log.Debug("mcp server stderr", "name", s.cfg.Name, "msg", strings.TrimSpace(string(buf[:n])))
}
if err != nil {
return
}
}
}
// lineRPC carries newline-delimited JSON-RPC 2.0 over a stream pair. Calls are
// serialized so at most one request is in flight at a time, which keeps id
// matching simple and deterministic.
type lineRPC struct {
in io.Writer
out *bufio.Reader
log *slog.Logger
mu sync.Mutex
nextID int
}
// call writes a request and reads until the matching response arrives.
func (l *lineRPC) call(ctx context.Context, method string, params any, result any) error {
l.mu.Lock()
defer l.mu.Unlock()
if err := ctx.Err(); err != nil {
return fmt.Errorf("mcp %s: %w", method, err)
}
l.nextID++
id := strconv.Itoa(l.nextID)
req, err := encodeRequest(id, method, params)
if err != nil {
return fmt.Errorf("mcp %s: %w", method, err)
}
req = append(req, '\n')
if _, err := l.in.Write(req); err != nil {
return fmt.Errorf("mcp %s write: %w", method, err)
}
return l.readResponse(id, method, result)
}
// readResponse reads lines until a response matching id arrives, skipping blank
// lines, malformed frames and notifications (no id).
func (l *lineRPC) readResponse(id, method string, result any) error {
for {
line, err := l.out.ReadBytes('\n')
if err != nil {
return fmt.Errorf("mcp %s read: %w", method, err)
}
line = bytes.TrimSpace(line)
if len(line) == 0 {
continue
}
var resp rpcResponse
if err := json.Unmarshal(line, &resp); err != nil {
l.log.Debug("mcp ignoring malformed frame", "line", string(line))
continue
}
if resp.ID == nil {
continue
}
if !decodeIDString(resp, id) {
continue
}
if resp.Error != nil {
return fmt.Errorf("mcp %s: %w: %w", method, ErrRPC, resp.Error)
}
if result != nil {
if err := json.Unmarshal(resp.Result, result); err != nil {
return fmt.Errorf("mcp %s decode: %w", method, err)
}
}
return nil
}
}
// notify writes a fire-and-forget JSON-RPC notification (no id).
func (l *lineRPC) notify(ctx context.Context, method string, params any) error {
l.mu.Lock()
defer l.mu.Unlock()
req, err := encodeRequest(nil, method, params)
if err != nil {
return fmt.Errorf("mcp notify %s: %w", method, err)
}
req = append(req, '\n')
if _, err := l.in.Write(req); err != nil {
return fmt.Errorf("mcp notify %s write: %w", method, err)
}
return nil
}
+199
View File
@@ -0,0 +1,199 @@
package mcp
import (
"bufio"
"bytes"
"context"
"errors"
"io"
"strings"
"testing"
)
// newLineRPC builds a lineRPC reading the supplied pre-encoded responses.
func newLineRPC(responses string) *lineRPC {
return &lineRPC{
in: &bytes.Buffer{},
out: bufio.NewReader(bytes.NewReader([]byte(responses))),
log: defaultLogger(),
}
}
func TestLineRPC_Call_Success(t *testing.T) {
t.Parallel()
out := `{"jsonrpc":"2.0","id":"1","result":{"ok":true}}` + "\n"
l := newLineRPC(out)
var res struct {
OK bool `json:"ok"`
}
if err := l.call(context.Background(), "tools/call", map[string]any{}, &res); err != nil {
t.Fatalf("call: %v", err)
}
if !res.OK {
t.Fatal("res.OK = false, want true")
}
}
func TestLineRPC_Call_SkipsNotificationThenMatches(t *testing.T) {
t.Parallel()
out := `{"jsonrpc":"2.0","method":"notifications/initialized"}` + "\n" +
`{"jsonrpc":"2.0","id":"1","result":{}}` + "\n"
l := newLineRPC(out)
var res map[string]any
if err := l.call(context.Background(), "initialize", map[string]any{}, &res); err != nil {
t.Fatalf("call: %v", err)
}
}
func TestLineRPC_Call_SkipsMismatchedID(t *testing.T) {
t.Parallel()
out := `{"jsonrpc":"2.0","id":"2","result":{}}` + "\n" +
`{"jsonrpc":"2.0","id":"1","result":{"ok":true}}` + "\n"
l := newLineRPC(out)
var res struct {
OK bool `json:"ok"`
}
if err := l.call(context.Background(), "tools/call", nil, &res); err != nil {
t.Fatalf("call: %v", err)
}
if !res.OK {
t.Fatal("res.OK = false, want true")
}
}
func TestLineRPC_Call_RPCError(t *testing.T) {
t.Parallel()
out := `{"jsonrpc":"2.0","id":"1","error":{"code":-32000,"message":"boom"}}` + "\n"
l := newLineRPC(out)
err := l.call(context.Background(), "tools/call", nil, nil)
if !errors.Is(err, ErrRPC) {
t.Fatalf("err = %v, want ErrRPC", err)
}
var rpcErr *rpcError
if !errors.As(err, &rpcErr) {
t.Fatalf("err = %v, want *rpcError in chain", err)
}
if rpcErr.Message != "boom" {
t.Fatalf("rpcErr.Message = %q, want boom", rpcErr.Message)
}
}
func TestLineRPC_Call_LastResponseDecodeError(t *testing.T) {
t.Parallel()
// A matching id whose result does not fit the target type must surface a
// decode error.
out := `{"jsonrpc":"2.0","id":"1","result":{"ok":"not-a-bool"}}` + "\n"
l := newLineRPC(out)
var res struct {
OK bool `json:"ok"`
}
err := l.call(context.Background(), "tools/list", nil, &res)
if err == nil {
t.Fatal("call with mismatched result succeeded, want decode error")
}
}
func TestLineRPC_Call_EOF(t *testing.T) {
t.Parallel()
l := newLineRPC("")
err := l.call(context.Background(), "tools/list", nil, nil)
if !errors.Is(err, io.EOF) {
t.Fatalf("err = %v, want io.EOF", err)
}
}
func TestLineRPC_Call_MalformedFrameIgnored(t *testing.T) {
t.Parallel()
out := `not json` + "\n" + `{"jsonrpc":"2.0","id":"1","result":{"ok":true}}` + "\n"
l := newLineRPC(out)
var res struct {
OK bool `json:"ok"`
}
if err := l.call(context.Background(), "tools/call", nil, &res); err != nil {
t.Fatalf("call: %v", err)
}
if !res.OK {
t.Fatal("res.OK = false, want true")
}
}
func TestLineRPC_WriteError(t *testing.T) {
t.Parallel()
l := &lineRPC{in: failWriter{}, out: bufio.NewReader(bytes.NewReader(nil)), log: defaultLogger()}
err := l.call(context.Background(), "tools/list", nil, nil)
if err == nil {
t.Fatal("call with failing writer succeeded, want error")
}
if !strings.Contains(err.Error(), "write") {
t.Fatalf("err = %v, want write error", err)
}
}
func TestLineRPC_Notify_WritesNoID(t *testing.T) {
t.Parallel()
var buf bytes.Buffer
l := &lineRPC{in: &buf, out: bufio.NewReader(bytes.NewReader(nil)), log: defaultLogger()}
if err := l.notify(context.Background(), "notifications/initialized", map[string]any{}); err != nil {
t.Fatalf("notify: %v", err)
}
body := buf.String()
if !strings.Contains(body, "notifications/initialized") {
t.Fatalf("body = %q, want notifications/initialized", body)
}
if strings.Contains(body, `"id"`) {
t.Fatalf("body = %q, must not contain id (notification)", body)
}
}
func TestLineRPC_ContextCanceled(t *testing.T) {
t.Parallel()
ctx, cancel := context.WithCancel(context.Background())
cancel()
l := newLineRPC(`{"jsonrpc":"2.0","id":"1","result":{}}` + "\n")
err := l.call(ctx, "initialize", nil, nil)
if !errors.Is(err, context.Canceled) {
t.Fatalf("err = %v, want context.Canceled", err)
}
}
func TestStdioTransport_Start_NoCommand(t *testing.T) {
t.Parallel()
s := newStdioTransport(MCPConfig{Name: "s"}, defaultLogger())
err := s.Start(context.Background())
if !errors.Is(err, ErrNoTransport) {
t.Fatalf("Start error = %v, want ErrNoTransport", err)
}
}
func TestStdioTransport_Start_InvalidCommand(t *testing.T) {
t.Parallel()
s := newStdioTransport(MCPConfig{Name: "s", Command: "cyrene-no-such-binary-xyz"}, defaultLogger())
if err := s.Start(context.Background()); err == nil {
t.Fatal("Start with bogus command succeeded, want error")
}
}
func TestStdioTransport_Call_NotStarted(t *testing.T) {
t.Parallel()
s := newStdioTransport(MCPConfig{Name: "s", Command: "x"}, defaultLogger())
err := s.Call(context.Background(), "tools/list", nil, nil)
if !errors.Is(err, ErrNotConnected) {
t.Fatalf("Call error = %v, want ErrNotConnected", err)
}
}
func TestStdioTransport_Close_NoCmd(t *testing.T) {
t.Parallel()
s := newStdioTransport(MCPConfig{Name: "s"}, defaultLogger())
if err := s.Close(); err != nil {
t.Fatalf("Close: %v", err)
}
}
// failWriter returns an error on every write.
type failWriter struct{}
// Write always fails.
func (failWriter) Write(_ []byte) (int, error) {
return 0, errors.New("write boom")
}
+63
View File
@@ -0,0 +1,63 @@
package mcp
import (
"context"
"strings"
)
// mcpTool adapts an MCP tool to the pkg/plugins/sdk Tool interface so it can be
// registered and invoked like any built-in tool. It is immutable and safe for
// concurrent use.
type mcpTool struct {
client *Client
info ToolInfo
}
// Name returns the MCP tool identifier.
func (t *mcpTool) Name() string {
return t.info.Name
}
// Description returns the MCP tool description.
func (t *mcpTool) Description() string {
return t.info.Description
}
// Parameters returns the MCP tool input schema, defaulting to an empty object
// schema when the server omits one.
func (t *mcpTool) Parameters() (map[string]any, error) {
if t.info.InputSchema == nil {
return map[string]any{"type": "object", "properties": map[string]any{}}, nil
}
return t.info.InputSchema, nil
}
// Execute invokes the MCP tool and returns its concatenated text content.
func (t *mcpTool) Execute(ctx context.Context, args any) (any, error) {
res, err := t.client.CallTool(ctx, t.info.Name, args)
if err != nil {
return nil, err
}
return contentText(res.Content), nil
}
// contentText concatenates the text blocks of a tool result into a single
// string, ignoring non-text blocks.
func contentText(content []ContentBlock) string {
var b strings.Builder
for _, blk := range content {
if blk.Type == "" || blk.Type == "text" {
b.WriteString(blk.Text)
}
}
return b.String()
}
// toContentBlocks converts wire content blocks to the public content block type.
func toContentBlocks(wire []contentBlock) []ContentBlock {
out := make([]ContentBlock, 0, len(wire))
for _, c := range wire {
out = append(out, ContentBlock{Type: c.Type, Text: c.Text})
}
return out
}