internal/mcp: an MCP SDK prototype
This CL contains an minimal prototype of an MCP SDK, using the "new" jsonrpc2_v2 package (similar to x/exp/jsonrpc2). Much is still yet to do, but this initial version addressed the following aspects: - Support for newline delimited JSONRPC2 framing, and message logging. - A generated protocol package, containing type definitions to be used in transport. - A minimal jsonschema package to be used for both reading the MCP spec, and for serving Tool input schemas. - A transport abstraction, and two transport implementations: stdio, and local (for testing). - Client and Server types, to be configured with features and then connected using their Connect methods, which return (respectively) a ServerConnection and ClientConnection object. - A minimal binding API for tools. A catalog of things not yet done is in doc.go, but we should review this CL eagerly, as it contains the fundamental building blocks of an eventual SDK. It is the intention that this package is more-or-less standalone, as it may eventually be carved out into a separate repository. As such, its only dependency within x/tools is jsonrpc2_v2, and the only change required to that package is to expose a callback when the connection is closed. Change-Id: Ib66018961014f85a8884d56ab721d5642c98443c Reviewed-on: https://go-review.googlesource.com/c/tools/+/667036 Reviewed-by: Jonathan Amsterdam <jba@google.com> LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
This commit is contained in:
parent
47e1b0be9d
commit
34cc770852
8 changed files with 1880 additions and 11 deletions
191
mcp/client.go
Normal file
191
mcp/client.go
Normal file
|
|
@ -0,0 +1,191 @@
|
|||
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
|
||||
// Use of this source code is governed by an MIT-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"iter"
|
||||
"slices"
|
||||
"sync"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
|
||||
"github.com/modelcontextprotocol/go-sdk/internal/protocol"
|
||||
)
|
||||
|
||||
// A Client is an MCP client, which may be connected to one or more MCP servers
|
||||
// using the [Client.Connect] method.
|
||||
//
|
||||
// TODO(rfindley): revisit the many-to-one relationship of clients and servers.
|
||||
// It is a bit odd.
|
||||
type Client struct {
|
||||
name string
|
||||
version string
|
||||
|
||||
mu sync.Mutex
|
||||
servers []*ServerConnection
|
||||
}
|
||||
|
||||
// NewClient creates a new Client.
|
||||
//
|
||||
// Use [Client.Connect] to connect it to an MCP server.
|
||||
//
|
||||
// If non-nil, the provided options configure the Client.
|
||||
func NewClient(name, version string, opts *ClientOptions) *Client {
|
||||
return &Client{
|
||||
name: name,
|
||||
version: version,
|
||||
}
|
||||
}
|
||||
|
||||
// Servers returns an iterator that yields the current set of server
|
||||
// connections.
|
||||
func (c *Client) Servers() iter.Seq[*ServerConnection] {
|
||||
c.mu.Lock()
|
||||
clients := slices.Clone(c.servers)
|
||||
c.mu.Unlock()
|
||||
return slices.Values(clients)
|
||||
}
|
||||
|
||||
// ClientOptions configures the behavior of the client, and apply to every
|
||||
// client-server connection created using [Client.Connect].
|
||||
type ClientOptions struct{}
|
||||
|
||||
// bind implements the binder[*ServerConnection] interface, so that Clients can
|
||||
// be connected using [connect].
|
||||
func (c *Client) bind(conn *jsonrpc2.Connection) *ServerConnection {
|
||||
sc := &ServerConnection{
|
||||
conn: conn,
|
||||
client: c,
|
||||
}
|
||||
c.mu.Lock()
|
||||
c.servers = append(c.servers, sc)
|
||||
c.mu.Unlock()
|
||||
return sc
|
||||
}
|
||||
|
||||
// disconnect implements the binder[*ServerConnection] interface, so that
|
||||
// Clients can be connected using [connect].
|
||||
func (c *Client) disconnect(sc *ServerConnection) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.servers = slices.DeleteFunc(c.servers, func(sc2 *ServerConnection) bool {
|
||||
return sc2 == sc
|
||||
})
|
||||
}
|
||||
|
||||
// Connect connects the MCP client over the given transport and initializes an
|
||||
// MCP session.
|
||||
//
|
||||
// It returns a connection object that may be used to query the MCP server,
|
||||
// terminate the connection (with [Connection.Close]), or await server
|
||||
// termination (with [Connection.Wait]).
|
||||
//
|
||||
// Typically, it is the responsibility of the client to close the connection
|
||||
// when it is no longer needed. However, if the connection is closed by the
|
||||
// server, calls or notifications will return an error wrapping
|
||||
// [ErrConnectionClosed].
|
||||
func (c *Client) Connect(ctx context.Context, t *Transport, opts *ConnectionOptions) (sc *ServerConnection, err error) {
|
||||
defer func() {
|
||||
if sc != nil && err != nil {
|
||||
_ = sc.Close()
|
||||
}
|
||||
}()
|
||||
sc, err = connect(ctx, t, opts, c)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
params := &protocol.InitializeParams{
|
||||
ClientInfo: protocol.Implementation{Name: c.name, Version: c.version},
|
||||
}
|
||||
if err := call(ctx, sc.conn, "initialize", params, &sc.initializeResult); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := sc.conn.Notify(ctx, "initialized", &protocol.InitializedParams{}); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return sc, nil
|
||||
}
|
||||
|
||||
// A ServerConnection is a connection with an MCP server.
|
||||
//
|
||||
// It handles messages from the client, and can be used to send messages to the
|
||||
// client. Create a connection by calling [Server.Connect].
|
||||
type ServerConnection struct {
|
||||
conn *jsonrpc2.Connection
|
||||
client *Client
|
||||
initializeResult *protocol.InitializeResult
|
||||
}
|
||||
|
||||
// Close performs a graceful close of the connection, preventing new requests
|
||||
// from being handled, and waiting for ongoing requests to return. Close then
|
||||
// terminates the connection.
|
||||
func (cc *ServerConnection) Close() error {
|
||||
return cc.conn.Close()
|
||||
}
|
||||
|
||||
// Wait waits for the connection to be closed by the server.
|
||||
// Generally, clients should be responsible for closing the connection.
|
||||
func (cc *ServerConnection) Wait() error {
|
||||
return cc.conn.Wait()
|
||||
}
|
||||
|
||||
func (sc *ServerConnection) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) {
|
||||
switch req.Method {
|
||||
}
|
||||
return nil, jsonrpc2.ErrNotHandled
|
||||
}
|
||||
|
||||
// ListTools lists tools that are currently available on the server.
|
||||
func (sc *ServerConnection) ListTools(ctx context.Context) ([]protocol.Tool, error) {
|
||||
var (
|
||||
params = &protocol.ListToolsParams{}
|
||||
result protocol.ListToolsResult
|
||||
)
|
||||
if err := call(ctx, sc.conn, "tools/list", params, &result); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return result.Tools, nil
|
||||
}
|
||||
|
||||
// CallTool calls the tool with the given name and arguments.
|
||||
//
|
||||
// TODO: make the following true:
|
||||
// If the provided arguments do not conform to the schema for the given tool,
|
||||
// the call fails.
|
||||
func (sc *ServerConnection) CallTool(ctx context.Context, name string, args any) (_ []Content, err error) {
|
||||
defer func() {
|
||||
if err != nil {
|
||||
err = fmt.Errorf("calling tool %q: %w", name, err)
|
||||
}
|
||||
}()
|
||||
argJSON, err := json.Marshal(args)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshaling args: %v", err)
|
||||
}
|
||||
var (
|
||||
params = &protocol.CallToolParams{
|
||||
Name: name,
|
||||
Arguments: argJSON,
|
||||
}
|
||||
result protocol.CallToolResult
|
||||
)
|
||||
if err := call(ctx, sc.conn, "tools/call", params, &result); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
content, err := unmarshalContent(result.Content)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unmarshaling tool content: %v", err)
|
||||
}
|
||||
if result.IsError {
|
||||
if len(content) != 1 || !is[TextContent](content[0]) {
|
||||
return nil, errors.New("malformed error content")
|
||||
}
|
||||
return nil, errors.New(content[0].(TextContent).Text)
|
||||
}
|
||||
return content, nil
|
||||
}
|
||||
64
mcp/content.go
Normal file
64
mcp/content.go
Normal file
|
|
@ -0,0 +1,64 @@
|
|||
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
|
||||
// Use of this source code is governed by an MIT-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/internal/protocol"
|
||||
)
|
||||
|
||||
// Content is the abstract result of a Tool call.
|
||||
//
|
||||
// TODO: support all content types.
|
||||
type Content interface {
|
||||
toProtocol() any
|
||||
}
|
||||
|
||||
func marshalContent(content []Content) []json.RawMessage {
|
||||
var msgs []json.RawMessage
|
||||
for _, c := range content {
|
||||
msg, err := json.Marshal(c.toProtocol())
|
||||
if err != nil {
|
||||
panic(fmt.Sprintf("marshaling content: %v", err))
|
||||
}
|
||||
msgs = append(msgs, msg)
|
||||
}
|
||||
return msgs
|
||||
}
|
||||
|
||||
func unmarshalContent(msgs []json.RawMessage) ([]Content, error) {
|
||||
var content []Content
|
||||
for _, msg := range msgs {
|
||||
var allContent struct {
|
||||
Type string `json:"type"`
|
||||
Text json.RawMessage
|
||||
}
|
||||
if err := json.Unmarshal(msg, &allContent); err != nil {
|
||||
return nil, fmt.Errorf("content missing \"type\"")
|
||||
}
|
||||
switch allContent.Type {
|
||||
case "text":
|
||||
var text string
|
||||
if err := json.Unmarshal(allContent.Text, &text); err != nil {
|
||||
return nil, fmt.Errorf("unmarshalling text content: %v", err)
|
||||
}
|
||||
content = append(content, TextContent{Text: text})
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported content type %q", allContent.Type)
|
||||
}
|
||||
}
|
||||
return content, nil
|
||||
}
|
||||
|
||||
// TextContent is a textual content.
|
||||
type TextContent struct {
|
||||
Text string
|
||||
}
|
||||
|
||||
func (c TextContent) toProtocol() any {
|
||||
return protocol.TextContent{Type: "text", Text: c.Text}
|
||||
}
|
||||
23
mcp/mcp.go
Normal file
23
mcp/mcp.go
Normal file
|
|
@ -0,0 +1,23 @@
|
|||
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
|
||||
// Use of this source code is governed by an MIT-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
// The mcp package provides an SDK for writing model context protocol clients
|
||||
// and servers. It is a work-in-progress. As of writing, it is a prototype to
|
||||
// explore the design space of client/server lifecycle and binding.
|
||||
//
|
||||
// To get started, create an MCP client or server with [NewClient] or
|
||||
// [NewServer], then add features to your client or server using Add<feature>
|
||||
// methods, then connect to a peer using a [Transport] instance and a call to
|
||||
// [Client.Connect] or [Server.Connect].
|
||||
//
|
||||
// TODO:
|
||||
// - Support cancellation.
|
||||
// - Support pagination.
|
||||
// - Support all client/server operations.
|
||||
// - Support Streamable HTTP transport.
|
||||
// - Support multiple versions of the spec.
|
||||
// - Implement proper JSON schema support, with both client-side and
|
||||
// server-side validation..
|
||||
// - Support batched JSON messages.
|
||||
package mcp
|
||||
161
mcp/mcp_test.go
Normal file
161
mcp/mcp_test.go
Normal file
|
|
@ -0,0 +1,161 @@
|
|||
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
|
||||
// Use of this source code is governed by an MIT-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package mcp_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"slices"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
"github.com/modelcontextprotocol/go-sdk/internal/jsonschema"
|
||||
"github.com/modelcontextprotocol/go-sdk/internal/protocol"
|
||||
)
|
||||
|
||||
type hiParams struct {
|
||||
Name string
|
||||
}
|
||||
|
||||
func sayHi(_ context.Context, v hiParams) ([]mcp.Content, error) {
|
||||
return []mcp.Content{mcp.TextContent{Text: "hi " + v.Name}}, nil
|
||||
}
|
||||
|
||||
func TestEndToEnd(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
ct, st := mcp.NewLocalTransport()
|
||||
|
||||
s := mcp.NewServer("testServer", "v1.0.0", nil)
|
||||
|
||||
// The 'greet' tool says hi.
|
||||
s.AddTools(mcp.MakeTool("greet", "say hi", sayHi))
|
||||
|
||||
// The 'fail' tool returns this error.
|
||||
failure := errors.New("mcp failure")
|
||||
s.AddTools(mcp.MakeTool("fail", "just fail", func(context.Context, struct{}) ([]mcp.Content, error) {
|
||||
return nil, failure
|
||||
}))
|
||||
|
||||
// Connect the server.
|
||||
cc, err := s.Connect(ctx, st, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := slices.Collect(s.Clients()); len(got) != 1 {
|
||||
t.Errorf("after connection, Clients() has length %d, want 1", len(got))
|
||||
}
|
||||
|
||||
// Wait for the server to exit after the client closes its connection.
|
||||
var clientWG sync.WaitGroup
|
||||
clientWG.Add(1)
|
||||
go func() {
|
||||
if err := cc.Wait(); err != nil {
|
||||
t.Errorf("server failed: %v", err)
|
||||
}
|
||||
clientWG.Done()
|
||||
}()
|
||||
|
||||
c := mcp.NewClient("testClient", "v1.0.0", nil)
|
||||
|
||||
// Connect the client.
|
||||
sc, err := c.Connect(ctx, ct, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := slices.Collect(c.Servers()); len(got) != 1 {
|
||||
t.Errorf("after connection, Servers() has length %d, want 1", len(got))
|
||||
}
|
||||
|
||||
gotTools, err := sc.ListTools(ctx)
|
||||
if err != nil {
|
||||
t.Errorf("tools/list failed: %v", err)
|
||||
}
|
||||
wantTools := []protocol.Tool{{
|
||||
Name: "greet",
|
||||
Description: "say hi",
|
||||
InputSchema: &jsonschema.Schema{
|
||||
Type: "object",
|
||||
Properties: map[string]*jsonschema.Schema{
|
||||
"Name": {Type: "string"},
|
||||
},
|
||||
AdditionalProperties: false,
|
||||
},
|
||||
}, {
|
||||
Name: "fail",
|
||||
Description: "just fail",
|
||||
InputSchema: &jsonschema.Schema{
|
||||
Type: "object",
|
||||
AdditionalProperties: false,
|
||||
},
|
||||
}}
|
||||
if diff := cmp.Diff(wantTools, gotTools); diff != "" {
|
||||
t.Fatalf("tools/list mismatch (-want +got):\n%s", diff)
|
||||
}
|
||||
|
||||
gotHi, err := sc.CallTool(ctx, "greet", hiParams{"user"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
wantHi := []mcp.Content{mcp.TextContent{Text: "hi user"}}
|
||||
if diff := cmp.Diff(wantHi, gotHi); diff != "" {
|
||||
t.Errorf("tools/call 'greet' mismatch (-want +got):\n%s", diff)
|
||||
}
|
||||
|
||||
if _, err := sc.CallTool(ctx, "fail", struct{}{}); err == nil || !strings.Contains(err.Error(), failure.Error()) {
|
||||
t.Errorf("fail returned unexpected error: got %v, want containing %v", err, failure)
|
||||
}
|
||||
|
||||
// Disconnect.
|
||||
sc.Close()
|
||||
clientWG.Wait()
|
||||
|
||||
// After disconnecting, neither client nor server should have any
|
||||
// connections.
|
||||
for range s.Clients() {
|
||||
t.Errorf("unexpected client after disconnection")
|
||||
}
|
||||
for range c.Servers() {
|
||||
t.Errorf("unexpected server after disconnection")
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerClosing(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
ct, st := mcp.NewLocalTransport()
|
||||
|
||||
s := mcp.NewServer("testServer", "v1.0.0", nil)
|
||||
|
||||
// The 'greet' tool says hi.
|
||||
s.AddTools(mcp.MakeTool("greet", "say hi", sayHi))
|
||||
cc, err := s.Connect(ctx, st, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
c := mcp.NewClient("testClient", "v1.0.0", nil)
|
||||
sc, err := c.Connect(ctx, ct, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
if err := sc.Wait(); err != nil {
|
||||
t.Errorf("server connection failed: %v", err)
|
||||
}
|
||||
wg.Done()
|
||||
}()
|
||||
if _, err := sc.CallTool(ctx, "greet", hiParams{"user"}); err != nil {
|
||||
t.Fatalf("after connecting: %v", err)
|
||||
}
|
||||
cc.Close()
|
||||
wg.Wait()
|
||||
if _, err = sc.CallTool(ctx, "greet", hiParams{"user"}); !errors.Is(err, mcp.ErrConnectionClosed) {
|
||||
t.Errorf("after disconnection, got error %v, want EOF", err)
|
||||
}
|
||||
}
|
||||
212
mcp/server.go
Normal file
212
mcp/server.go
Normal file
|
|
@ -0,0 +1,212 @@
|
|||
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
|
||||
// Use of this source code is governed by an MIT-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"iter"
|
||||
"slices"
|
||||
"sync"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
|
||||
"github.com/modelcontextprotocol/go-sdk/internal/protocol"
|
||||
)
|
||||
|
||||
// A Server is an instance of an MCP server.
|
||||
//
|
||||
// Servers expose server-side MCP features, which can serve one or more MCP
|
||||
// sessions by using [Server.Start] or [Server.Run].
|
||||
type Server struct {
|
||||
name string
|
||||
version string
|
||||
opts ServerOptions
|
||||
|
||||
mu sync.Mutex
|
||||
tools []*Tool
|
||||
clients []*ClientConnection
|
||||
}
|
||||
|
||||
// ServerOptions is used to configure behavior of the server.
|
||||
type ServerOptions struct {
|
||||
Instructions string
|
||||
}
|
||||
|
||||
// NewServer creates a new MCP server. The resulting server has no features:
|
||||
// add features using [Server.AddTools]. (TODO: support more features).
|
||||
//
|
||||
// The server can be connected to one or more MCP clients using [Server.Start]
|
||||
// or [Server.Run].
|
||||
//
|
||||
// If non-nil, the provided options is used to configure the server.
|
||||
func NewServer(name, version string, opts *ServerOptions) *Server {
|
||||
if opts == nil {
|
||||
opts = new(ServerOptions)
|
||||
}
|
||||
return &Server{
|
||||
name: name,
|
||||
version: version,
|
||||
opts: *opts,
|
||||
}
|
||||
}
|
||||
|
||||
// AddTools adds the given tools to the server.
|
||||
//
|
||||
// TODO(rfindley): notify connected clients of any changes.
|
||||
func (c *Server) AddTools(tools ...*Tool) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.tools = append(c.tools, tools...)
|
||||
}
|
||||
|
||||
// Clients returns an iterator that yields the current set of client
|
||||
// connections.
|
||||
func (s *Server) Clients() iter.Seq[*ClientConnection] {
|
||||
s.mu.Lock()
|
||||
clients := slices.Clone(s.clients)
|
||||
s.mu.Unlock()
|
||||
return slices.Values(clients)
|
||||
}
|
||||
|
||||
func (c *Server) listTools(_ context.Context, params *protocol.ListToolsParams) (*protocol.ListToolsResult, error) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
res := new(protocol.ListToolsResult)
|
||||
for _, t := range c.tools {
|
||||
res.Tools = append(res.Tools, t.Definition)
|
||||
}
|
||||
return res, nil
|
||||
}
|
||||
|
||||
func (c *Server) callTool(ctx context.Context, params *protocol.CallToolParams) (*protocol.CallToolResult, error) {
|
||||
c.mu.Lock()
|
||||
var tool *Tool
|
||||
if i := slices.IndexFunc(c.tools, func(t *Tool) bool {
|
||||
return t.Definition.Name == params.Name
|
||||
}); i >= 0 {
|
||||
tool = c.tools[i]
|
||||
}
|
||||
c.mu.Unlock()
|
||||
|
||||
if tool == nil {
|
||||
return nil, fmt.Errorf("%s: unknown tool %q", jsonrpc2.ErrInvalidParams, params.Name)
|
||||
}
|
||||
return tool.Handler(ctx, params.Arguments)
|
||||
}
|
||||
|
||||
// Run runs the server over the given transport.
|
||||
//
|
||||
// Run blocks until the client terminates the connection.
|
||||
func (c *Server) Run(ctx context.Context, t *Transport, opts *ConnectionOptions) error {
|
||||
conn, err := c.Connect(ctx, t, opts)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return conn.Wait()
|
||||
}
|
||||
|
||||
// bind implements the binder[*ClientConnection] interface, so that Servers can
|
||||
// be connected using [connect].
|
||||
func (c *Server) bind(conn *jsonrpc2.Connection) *ClientConnection {
|
||||
cc := &ClientConnection{conn: conn, server: c}
|
||||
c.mu.Lock()
|
||||
c.clients = append(c.clients, cc)
|
||||
c.mu.Unlock()
|
||||
return cc
|
||||
}
|
||||
|
||||
// disconnect implements the binder[*ClientConnection] interface, so that
|
||||
// Servers can be connected using [connect].
|
||||
func (c *Server) disconnect(cc *ClientConnection) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.clients = slices.DeleteFunc(c.clients, func(cc2 *ClientConnection) bool {
|
||||
return cc2 == cc
|
||||
})
|
||||
}
|
||||
|
||||
// Connect connects the MCP server over the given transport and starts handling
|
||||
// messages.
|
||||
//
|
||||
// It returns a connection object that may be used to terminate the connection
|
||||
// (with [Connection.Close]), or await client termination (with
|
||||
// [Connection.Wait]).
|
||||
func (c *Server) Connect(ctx context.Context, t *Transport, opts *ConnectionOptions) (*ClientConnection, error) {
|
||||
return connect(ctx, t, opts, c)
|
||||
}
|
||||
|
||||
// A ClientConnection is a connection with an MCP client.
|
||||
//
|
||||
// It handles messages from the client, and can be used to send messages to the
|
||||
// client. Create a connection by calling [Server.Connect].
|
||||
type ClientConnection struct {
|
||||
conn *jsonrpc2.Connection
|
||||
server *Server
|
||||
|
||||
mu sync.Mutex
|
||||
initializeParams *protocol.InitializeParams // set once initialize has been received
|
||||
}
|
||||
|
||||
func (cc *ClientConnection) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) {
|
||||
switch req.Method {
|
||||
case "initialize":
|
||||
return dispatch(ctx, req, cc.initialize)
|
||||
|
||||
// TODO: handle initialized
|
||||
|
||||
case "tools/list":
|
||||
return dispatch(ctx, req, cc.server.listTools)
|
||||
|
||||
case "tools/call":
|
||||
return dispatch(ctx, req, cc.server.callTool)
|
||||
|
||||
case "notifications/initialized":
|
||||
}
|
||||
return nil, jsonrpc2.ErrNotHandled
|
||||
}
|
||||
|
||||
func (cc *ClientConnection) initialize(ctx context.Context, params *protocol.InitializeParams) (*protocol.InitializeResult, error) {
|
||||
cc.mu.Lock()
|
||||
cc.initializeParams = params
|
||||
cc.mu.Unlock()
|
||||
|
||||
return &protocol.InitializeResult{
|
||||
// TODO(rfindley): support multiple protocol versions.
|
||||
ProtocolVersion: "2024-11-05",
|
||||
Capabilities: protocol.ServerCapabilities{
|
||||
Tools: &protocol.ToolCapabilities{
|
||||
ListChanged: true,
|
||||
},
|
||||
},
|
||||
Instructions: cc.server.opts.Instructions,
|
||||
ServerInfo: protocol.Implementation{
|
||||
Name: cc.server.name,
|
||||
Version: cc.server.version,
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Close performs a graceful close of the connection, preventing new requests
|
||||
// from being handled, and waiting for ongoing requests to return. Close then
|
||||
// terminates the connection.
|
||||
func (cc *ClientConnection) Close() error {
|
||||
return cc.conn.Close()
|
||||
}
|
||||
|
||||
// Wait waits for the connection to be closed by the client.
|
||||
func (cc *ClientConnection) Wait() error {
|
||||
return cc.conn.Wait()
|
||||
}
|
||||
|
||||
func dispatch[TParams, TResult any](ctx context.Context, req *jsonrpc2.Request, f func(context.Context, TParams) (TResult, error)) (TResult, error) {
|
||||
var params TParams
|
||||
if err := json.Unmarshal(req.Params, ¶ms); err != nil {
|
||||
var zero TResult
|
||||
return zero, err
|
||||
}
|
||||
return f(ctx, params)
|
||||
}
|
||||
70
mcp/tool.go
Normal file
70
mcp/tool.go
Normal file
|
|
@ -0,0 +1,70 @@
|
|||
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
|
||||
// Use of this source code is governed by an MIT-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/internal/jsonschema"
|
||||
"github.com/modelcontextprotocol/go-sdk/internal/protocol"
|
||||
)
|
||||
|
||||
// A ToolHandler handles a call to tools/call.
|
||||
type ToolHandler func(context.Context, json.RawMessage) (*protocol.CallToolResult, error)
|
||||
|
||||
// A Tool is a tool definition that is bound to a tool handler.
|
||||
type Tool struct {
|
||||
Definition protocol.Tool
|
||||
Handler ToolHandler
|
||||
}
|
||||
|
||||
// MakeTool is a helper to make a tool using reflection on the given handler.
|
||||
//
|
||||
// The input schema for the tool is extracted from the request type for the
|
||||
// handler, and used to unmmarshal and validate requests to the handler.
|
||||
//
|
||||
// It is the caller's responsibility that the handler request type can produce
|
||||
// a valid schema, as documented by [jsonschema.ForType]; otherwise, MakeTool
|
||||
// panics.
|
||||
func MakeTool[TReq any](name, description string, handler func(context.Context, TReq) ([]Content, error)) *Tool {
|
||||
schema, err := jsonschema.ForType[TReq]()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
wrapped := func(ctx context.Context, args json.RawMessage) (*protocol.CallToolResult, error) {
|
||||
var v TReq
|
||||
if err := unmarshalSchema(args, schema, &v); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
content, err := handler(ctx, v)
|
||||
if err != nil {
|
||||
return &protocol.CallToolResult{
|
||||
Content: marshalContent([]Content{TextContent{Text: err.Error()}}),
|
||||
IsError: true,
|
||||
}, nil
|
||||
}
|
||||
res := &protocol.CallToolResult{
|
||||
Content: marshalContent(content),
|
||||
}
|
||||
return res, nil
|
||||
}
|
||||
return &Tool{
|
||||
Definition: protocol.Tool{
|
||||
Name: name,
|
||||
Description: description,
|
||||
InputSchema: schema,
|
||||
},
|
||||
Handler: wrapped,
|
||||
}
|
||||
}
|
||||
|
||||
// unmarshalSchema unmarshals data into v and validates the result according to
|
||||
// the given schema.
|
||||
func unmarshalSchema(data json.RawMessage, _ *jsonschema.Schema, v any) error {
|
||||
// TODO: use reflection to create the struct type to unmarshal into.
|
||||
// Separate validation from assignment.
|
||||
return json.Unmarshal(data, v)
|
||||
}
|
||||
259
mcp/transport.go
Normal file
259
mcp/transport.go
Normal file
|
|
@ -0,0 +1,259 @@
|
|||
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
|
||||
// Use of this source code is governed by an MIT-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net"
|
||||
"os"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
|
||||
)
|
||||
|
||||
// A JSONRPC2 error is an error defined by the JSONRPC2 spec.
|
||||
type JSONRPC2Error = jsonrpc2.WireError
|
||||
|
||||
// ErrConnectionClosed is returned when sending a message to a connection that
|
||||
// is closed or in the process of closing.
|
||||
var ErrConnectionClosed = errors.New("connection closed")
|
||||
|
||||
// A Transport is used to create a bidirectional connection between MCP client
|
||||
// and server.
|
||||
type Transport struct {
|
||||
dialer jsonrpc2.Dialer
|
||||
}
|
||||
|
||||
// ConnectionOptions configures the behavior of an individual client<->server
|
||||
// connection.
|
||||
type ConnectionOptions struct {
|
||||
Logger io.Writer
|
||||
}
|
||||
|
||||
// NewStdIOTransport constructs a transport that communicates over
|
||||
// stdin/stdout.
|
||||
func NewStdIOTransport() *Transport {
|
||||
dialer := dialerFunc(func(ctx context.Context) (io.ReadWriteCloser, error) {
|
||||
return rwc{os.Stdin, os.Stdout}, nil
|
||||
})
|
||||
return &Transport{
|
||||
dialer: dialer,
|
||||
}
|
||||
}
|
||||
|
||||
// NewLocalTransport returns two in-memory transports that connect to
|
||||
// each other, for testing purposes.
|
||||
func NewLocalTransport() (*Transport, *Transport) {
|
||||
c1, c2 := net.Pipe()
|
||||
t1 := &Transport{
|
||||
dialer: dialerFunc(func(ctx context.Context) (io.ReadWriteCloser, error) {
|
||||
return c1, nil
|
||||
}),
|
||||
}
|
||||
t2 := &Transport{
|
||||
dialer: dialerFunc(func(ctx context.Context) (io.ReadWriteCloser, error) {
|
||||
return c2, nil
|
||||
}),
|
||||
}
|
||||
return t1, t2
|
||||
}
|
||||
|
||||
// handler is an unexported version of jsonrpc2.Handler, to be implemented by
|
||||
// [ServerConnection] and [ClientConnection].
|
||||
type handler interface {
|
||||
handle(ctx context.Context, req *jsonrpc2.Request) (result any, err error)
|
||||
comparable
|
||||
}
|
||||
|
||||
type binder[T handler] interface {
|
||||
bind(*jsonrpc2.Connection) T
|
||||
disconnect(T)
|
||||
}
|
||||
|
||||
func connect[H handler](ctx context.Context, t *Transport, opts *ConnectionOptions, b binder[H]) (H, error) {
|
||||
if opts == nil {
|
||||
opts = new(ConnectionOptions)
|
||||
}
|
||||
|
||||
// Frame messages using newline delimited JSON.
|
||||
//
|
||||
// If logging is configured, write message logs.
|
||||
var framer jsonrpc2.Framer = &ndjsonFramer{}
|
||||
if opts.Logger != nil {
|
||||
framer = &loggingFramer{opts.Logger, framer}
|
||||
}
|
||||
|
||||
var h H
|
||||
|
||||
// Bind the server connection.
|
||||
binder := jsonrpc2.BinderFunc(func(_ context.Context, conn *jsonrpc2.Connection) jsonrpc2.ConnectionOptions {
|
||||
h = b.bind(conn)
|
||||
return jsonrpc2.ConnectionOptions{
|
||||
Framer: framer,
|
||||
Handler: jsonrpc2.HandlerFunc(h.handle),
|
||||
OnInternalError: func(err error) {
|
||||
log.Printf("Internal error: %v", err)
|
||||
},
|
||||
}
|
||||
})
|
||||
|
||||
// Clean up the connection when done.
|
||||
onDone := func() {
|
||||
b.disconnect(h)
|
||||
}
|
||||
|
||||
var zero H
|
||||
_, err := jsonrpc2.Dial(ctx, t.dialer, binder, onDone)
|
||||
if err != nil {
|
||||
return zero, err
|
||||
}
|
||||
assert(h != zero, "unbound connection")
|
||||
return h, nil
|
||||
}
|
||||
|
||||
// call executes and awaits a jsonrpc2 call on the given connection,
|
||||
// translating errors into the mcp domain.
|
||||
func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params, result any) error {
|
||||
err := conn.Call(ctx, method, params).Await(ctx, result)
|
||||
switch {
|
||||
case errors.Is(err, jsonrpc2.ErrClientClosing), errors.Is(err, jsonrpc2.ErrServerClosing):
|
||||
return fmt.Errorf("calling %q: %w", method, ErrConnectionClosed)
|
||||
case err != nil:
|
||||
return fmt.Errorf("calling %q: %v", method, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// The helpers below are used to bind transports to jsonrpc2.
|
||||
|
||||
// A dialerFunc implements jsonrpc2.Dialer.Dial.
|
||||
type dialerFunc func(context.Context) (io.ReadWriteCloser, error)
|
||||
|
||||
func (f dialerFunc) Dial(ctx context.Context) (io.ReadWriteCloser, error) {
|
||||
return f(ctx)
|
||||
}
|
||||
|
||||
// A readerFunc implements jsonrpc2.Reader.Read.
|
||||
type readerFunc func(context.Context) (jsonrpc2.Message, int64, error)
|
||||
|
||||
func (f readerFunc) Read(ctx context.Context) (jsonrpc2.Message, int64, error) {
|
||||
return f(ctx)
|
||||
}
|
||||
|
||||
// A writerFunc implements jsonrpc2.Writer.Write.
|
||||
type writerFunc func(context.Context, jsonrpc2.Message) (int64, error)
|
||||
|
||||
func (f writerFunc) Write(ctx context.Context, msg jsonrpc2.Message) (int64, error) {
|
||||
return f(ctx, msg)
|
||||
}
|
||||
|
||||
// A loggingFramer logs jsonrpc2 messages to its enclosed writer.
|
||||
type loggingFramer struct {
|
||||
w io.Writer
|
||||
delegate jsonrpc2.Framer
|
||||
}
|
||||
|
||||
func (f *loggingFramer) Reader(rw io.Reader) jsonrpc2.Reader {
|
||||
delegate := f.delegate.Reader(rw)
|
||||
return readerFunc(func(ctx context.Context) (jsonrpc2.Message, int64, error) {
|
||||
msg, n, err := delegate.Read(ctx)
|
||||
if err != nil {
|
||||
fmt.Fprintf(f.w, "read error: %v", err)
|
||||
} else {
|
||||
data, err := jsonrpc2.EncodeMessage(msg)
|
||||
if err != nil {
|
||||
fmt.Fprintf(f.w, "LoggingFramer: failed to marshal: %v", err)
|
||||
}
|
||||
fmt.Fprintf(f.w, "read: %s", string(data))
|
||||
}
|
||||
return msg, n, err
|
||||
})
|
||||
}
|
||||
|
||||
func (f *loggingFramer) Writer(w io.Writer) jsonrpc2.Writer {
|
||||
delegate := f.delegate.Writer(w)
|
||||
return writerFunc(func(ctx context.Context, msg jsonrpc2.Message) (int64, error) {
|
||||
n, err := delegate.Write(ctx, msg)
|
||||
if err != nil {
|
||||
fmt.Fprintf(f.w, "write error: %v", err)
|
||||
} else {
|
||||
data, err := jsonrpc2.EncodeMessage(msg)
|
||||
if err != nil {
|
||||
fmt.Fprintf(f.w, "LoggingFramer: failed to marshal: %v", err)
|
||||
}
|
||||
fmt.Fprintf(f.w, "write: %s", string(data))
|
||||
}
|
||||
return n, err
|
||||
})
|
||||
}
|
||||
|
||||
// A rwc binds an io.ReadCloser and io.WriteCloser together to create an
|
||||
// io.ReadWriteCloser.
|
||||
type rwc struct {
|
||||
rc io.ReadCloser
|
||||
wc io.WriteCloser
|
||||
}
|
||||
|
||||
func (r rwc) Read(p []byte) (n int, err error) {
|
||||
n, err = r.rc.Read(p)
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (r rwc) Write(p []byte) (n int, err error) {
|
||||
n, err = r.wc.Write(p)
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (r rwc) Close() error {
|
||||
return errors.Join(r.rc.Close(), r.wc.Close())
|
||||
}
|
||||
|
||||
// A ndjsonFramer is a jsonrpc2.Framer that delimits messages with newlines.
|
||||
//
|
||||
// See also https://github.com/ndjson/ndjson-spec.
|
||||
type ndjsonFramer struct{}
|
||||
type rawReader struct{ in *json.Decoder } // relies on json.Decoder message boundaries
|
||||
type ndjsonWriter struct{ out io.Writer } // writes newline message boundaries
|
||||
|
||||
func (ndjsonFramer) Reader(rw io.Reader) jsonrpc2.Reader {
|
||||
return &rawReader{in: json.NewDecoder(rw)}
|
||||
}
|
||||
|
||||
func (ndjsonFramer) Writer(rw io.Writer) jsonrpc2.Writer {
|
||||
return &ndjsonWriter{out: rw}
|
||||
}
|
||||
|
||||
func (r *rawReader) Read(ctx context.Context) (jsonrpc2.Message, int64, error) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, 0, ctx.Err()
|
||||
default:
|
||||
}
|
||||
var raw json.RawMessage
|
||||
if err := r.in.Decode(&raw); err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
msg, err := jsonrpc2.DecodeMessage(raw)
|
||||
return msg, int64(len(raw)), err
|
||||
}
|
||||
|
||||
func (w *ndjsonWriter) Write(ctx context.Context, msg jsonrpc2.Message) (int64, error) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return 0, ctx.Err()
|
||||
default:
|
||||
}
|
||||
data, err := jsonrpc2.EncodeMessage(msg)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("marshaling message: %v", err)
|
||||
}
|
||||
data = append(data, '\n') // newline delimited
|
||||
n, err := w.out.Write(data)
|
||||
return int64(n), err
|
||||
}
|
||||
16
mcp/util.go
Normal file
16
mcp/util.go
Normal file
|
|
@ -0,0 +1,16 @@
|
|||
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
|
||||
// Use of this source code is governed by an MIT-style
|
||||
// license that can be found in the LICENSE file.
|
||||
|
||||
package mcp
|
||||
|
||||
func assert(cond bool, msg string) {
|
||||
if !cond {
|
||||
panic(msg)
|
||||
}
|
||||
}
|
||||
|
||||
func is[T any](v any) bool {
|
||||
_, ok := v.(T)
|
||||
return ok
|
||||
}
|
||||
Loading…
Reference in a new issue