feat(capability/mcp): MCP 客户端(stdlib JSON-RPC over stdio/SSE)+ Registry 自动注册/Unregister(registry 增 Remove)
This commit is contained in:
1 parent
08223c06a5
commit
488cf9c79c
14 files changed
+2212
No files matched your search
@@ -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"`
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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 }
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
Reference in new issue
Block a user