modelcontextprotocol/go-sdk
Watch
1
0
Fork
You've already forked go-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 2

internal/mcp: implement streamable HTTP transport

Browse source 
Add an (almost) complete implementation of the streamable HTTP
transport, including:
- A StreamableHTTPHandler for serving sessions.
- A StreamableHTTPServerTransport that supports session resumption.
- A minimal StreamableHTTPClientTransport.

Additionally, include the following minor changes:
- Refactor SSE event reading to use an iterator.
- Start to to aliases for JSON-RPC types (incomplete)
- Fix handling for null tool arguments (again)

Still to do:
- The streamable client does not support retry/resumption, or the
  side-channel GET requests.
- Support configurable session and event retention (garbage collection).
- Expose the HTTP client used for HTTP requests (for auth).

Change-Id: Ieba216ef901169b4660d601386fa0c7c7bb4c4b2
Reviewed-on: https://go-review.googlesource.com/c/tools/+/667515
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:
Rob Findley 2025-04-23 10:52:44 +00:00 • committed by Robert Findley
commit d2f2f87bbe
10 changed files with 1570 additions and 149 deletions
Download patch file Download diff file Expand all files Collapse all files

9
design/design.md
View file

@ -58,6 +58,13 @@ Furthermore, the spec [states](https://modelcontextprotocol.io/specification/202
Given the diversity of the transport implementations, they can be challenging to abstract. However, since JSON-RPC requires a bidirectional stream, we can use this to model the MCP transport abstraction:
```go
type (
JSONRPCID = jsonrpc2.ID
JSONRPCMessage = jsonrpc2.Message
JSONRPCRequest = jsonrpc2.Request
JSONRPCResponse = jsonrpc2.Response
)
// A Transport is used to create a bidirectional connection between MCP client
// and server.
type Transport interface {
@ -204,7 +211,7 @@ func (*StreamableHTTPHandler) Close() error
// that created the session. It is the caller's responsibility to delegate
// requests to this session.
type StreamableServerTransport struct { /* ... */ }
func NewStreamableServerTransport(sessionID string, w http.ResponseWriter) *StreamableServerTransport
func NewStreamableServerTransport(sessionID string) *StreamableServerTransport
func (*StreamableServerTransport) ServeHTTP(w http.ResponseWriter, req *http.Request)
func (*StreamableServerTransport) Connect(context.Context) (Stream, error)

18
examples/hello/main.go
View file

@ -8,6 +8,7 @@ import (
"context"
"flag"
"fmt"
"log"
"net/http"
"net/url"
"os"
@ -15,21 +16,23 @@ import (
"github.com/modelcontextprotocol/go-sdk/mcp"
)
var httpAddr = flag.String("http", "", "if set, use SSE HTTP at this address, instead of stdin/stdout")
var httpAddr = flag.String("http", "", "if set, use streamable HTTP at this address, instead of stdin/stdout")
type HiArgs struct {
Name string `json:"name"`
}
func SayHi(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[HiArgs]) (*mcp.CallToolResultFor[any], error) {
return &mcp.CallToolResultFor[any]{
func SayHi(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[HiArgs]) (*mcp.CallToolResultFor[struct{}], error) {
return &mcp.CallToolResultFor[struct{}]{
Content: []*mcp.Content{
mcp.NewTextContent("Hi " + params.Name),
},
}, nil
}
func PromptHi(ctx context.Context, ss *mcp.ServerSession, args *HiArgs, params *mcp.GetPromptParams) (*mcp.GetPromptResult, error) {
// TODO(jba): it should be OK for args to be a pointer, but this fails in
// jsonschema. Needs investigation.
func PromptHi(ctx context.Context, ss *mcp.ServerSession, args HiArgs, params *mcp.GetPromptParams) (*mcp.GetPromptResult, error) {
return &mcp.GetPromptResult{
Description: "Code review prompt",
Messages: []*mcp.PromptMessage{
@ -56,14 +59,15 @@ func main() {
})
if *httpAddr != "" {
handler := mcp.NewSSEHandler(func(*http.Request) *mcp.Server {
handler := mcp.NewStreamableHTTPHandler(func(*http.Request) *mcp.Server {
return server
})
}, nil)
log.Printf("MCP handler listening at %s", *httpAddr)
http.ListenAndServe(*httpAddr, handler)
} else {
t := mcp.NewLoggingTransport(mcp.NewStdIOTransport(), os.Stderr)
if err := server.Run(context.Background(), t); err != nil {
fmt.Fprintf(os.Stderr, "Server failed: %v", err)
log.Printf("Server failed: %v", err)
}
}
}

10
mcp/prompt.go
View file

@ -34,13 +34,13 @@ type ServerPrompt struct {
//
// The handler is passed [GetPromptParams] so it can have access to prompt parameters other than name and arguments.
// At present, there are no such parameters.
func NewPrompt[TReq any](name, description string, handler func(context.Context, *ServerSession, TReq, *GetPromptParams) (*GetPromptResult, error), opts ...PromptOption) *ServerPrompt {
schema, err := jsonschema.For[TReq]()
func NewPrompt[In any](name, description string, handler func(context.Context, *ServerSession, In, *GetPromptParams) (*GetPromptResult, error), opts ...PromptOption) *ServerPrompt {
schema, err := jsonschema.For[In]()
if err != nil {
panic(err)
}
if schema.Type != "object" || !reflect.DeepEqual(schema.AdditionalProperties, &jsonschema.Schema{Not: &jsonschema.Schema{}}) {
panic(fmt.Sprintf("handler request type must be a struct"))
panic(fmt.Sprintf("invalid input type %q: handler input type must be a struct", schema.Type))
}
resolved, err := schema.Resolve(nil)
if err != nil {
@ -58,7 +58,7 @@ func NewPrompt[TReq any](name, description string, handler func(context.Context,
}
for name, prop := range util.Sorted(schema.Properties) {
if prop.Type != "string" {
panic(fmt.Sprintf("handler type must consist only of string fields"))
panic(fmt.Sprintf("invalid property type %q: handler type must consist only of string fields", prop.Type))
}
prompt.Prompt.Arguments = append(prompt.Prompt.Arguments, &PromptArgument{
Name: name,
@ -73,7 +73,7 @@ func NewPrompt[TReq any](name, description string, handler func(context.Context,
if err != nil {
return nil, err
}
var v TReq
var v In
if err := unmarshalSchema(rawArgs, resolved, &v); err != nil {
return nil, err
}

9
mcp/server.go
View file

@ -513,6 +513,8 @@ func (ss *ServerSession) CreateMessage(ctx context.Context, params *CreateMessag
// LoggingMessage sends a logging message to the client.
// The message is not sent if the client has not called SetLevel, or if its level
// is below that of the last SetLevel.
//
// TODO(jba): rename to Log or LogMessage. A logging message is the thing that is sent to logging.
func (ss *ServerSession) LoggingMessage(ctx context.Context, params *LoggingMessageParams) error {
ss.mu.Lock()
logLevel := ss.logLevel
@ -610,9 +612,10 @@ func (ss *ServerSession) handle(ctx context.Context, req *jsonrpc2.Request) (any
return nil, fmt.Errorf("method %q is invalid during session initialization", req.Method)
}
}
// TODO(rfindley): embed the incoming request ID in the client context (or, more likely,
// a wrapper around it), so that we can correlate responses and notifications
// to the handler; this is required for the new session-based transport.
// For the streamable transport, we need the request ID to correlate
// server->client calls and notifications to the incoming request from which
// they originated. See [idContext] for details.
ctx = context.WithValue(ctx, idContextKey{}, req.ID)
return handleReceive(ctx, ss, req)
}

169
mcp/sse.go
View file

@ -8,18 +8,20 @@ import (
"bufio"
"bytes"
"context"
"errors"
"fmt"
"io"
"iter"
"net/http"
"net/url"
"slices"
"strings"
"sync"
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
)
// This file implements support for SSE transport server and client.
// This file implements support for SSE (HTTP with server-sent events)
// transport server and client.
// https://modelcontextprotocol.io/specification/2024-11-05/basic/transports
//
// The transport is simple, at least relative to the new streamable transport
@ -32,24 +34,32 @@ import (
// 3. The client POSTs client->server messages to the session endpoint.
//
// Therefore, the each new GET request hands off its responsewriter to an
// [sseSession] type that abstracts the transport as follows:
// [SSEServerTransport] type that abstracts the transport as follows:
// - Write writes a new event to the responseWriter, or fails if the GET has
// exited.
// - Read reads off a message queue that is pushed to via POST requests.
// - Close causes the hanging GEt to exit.
// - Close causes the hanging GET to exit.
// An event is a server-sent event.
type event struct {
name string
id string
data []byte
}
func (e event) empty() bool {
return e.name == "" && e.id == "" && len(e.data) == 0
}
// writeEvent writes the event to w, and flushes.
func writeEvent(w io.Writer, evt event) (int, error) {
var b bytes.Buffer
if evt.name != "" {
fmt.Fprintf(&b, "event: %s\n", evt.name)
}
if evt.id != "" {
fmt.Fprintf(&b, "id: %s\n", evt.id)
}
fmt.Fprintf(&b, "data: %s\n\n", string(evt.data))
n, err := w.Write(b.Bytes())
if f, ok := w.(http.Flusher); ok {
@ -59,9 +69,9 @@ func writeEvent(w io.Writer, evt event) (int, error) {
}
// SSEHandler is an http.Handler that serves SSE-based MCP sessions as defined by
// the 2024-11-05 version of the MCP protocol:
// the [2024-11-05 version] of the MCP spec.
//
// https://modelcontextprotocol.io/specification/2024-11-05/basic/transports
// [2024-11-05 version]: https://modelcontextprotocol.io/specification/2024-11-05/basic/transports
type SSEHandler struct {
getServer func(request *http.Request) *Server
onConnection func(*ServerSession) // for testing; must not block
@ -82,6 +92,8 @@ type SSEHandler struct {
//
// The SSEHandler also handles requests to the message endpoints, by
// delegating them to the relevant server transport.
//
// TODO(rfindley): add options.
func NewSSEHandler(getServer func(request *http.Request) *Server) *SSEHandler {
return &SSEHandler{
getServer: getServer,
@ -337,56 +349,12 @@ func (c *SSEClientTransport) Connect(ctx context.Context) (Stream, error) {
if err != nil {
return nil, err
}
scanner := bufio.NewScanner(resp.Body)
// TODO: investigate proper behavior when events are out of order, or have
// non-standard names.
var (
eventKey = []byte("event")
dataKey = []byte("data")
)
// nextEvent reads one sse event from the wire.
// https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events#examples
//
// - `key: value` line records.
// - Consecutive `data: ...` fields are joined with newlines.
// - Unrecognized fields are ignored. Since we only care about 'event' and
// 'data', these are the only two we consider.
// - Lines starting with ":" are ignored.
// - Records are terminated with two consecutive newlines.
nextEvent := func() (event, error) {
var (
evt event
lastWasData bool // if set, preceding data field was also data
)
for scanner.Scan() {
line := scanner.Bytes()
if len(line) == 0 && (evt.name != "" || len(evt.data) > 0) {
return evt, nil
}
before, after, found := bytes.Cut(line, []byte{':'})
if !found {
return evt, fmt.Errorf("malformed line in SSE stream: %q", string(line))
}
switch {
case bytes.Equal(before, eventKey):
evt.name = strings.TrimSpace(string(after))
case bytes.Equal(before, dataKey):
data := bytes.TrimSpace(after)
if lastWasData {
evt.data = slices.Concat(evt.data, []byte{'\n'}, data)
} else {
evt.data = data
}
lastWasData = true
}
}
return evt, io.EOF
}
msgEndpoint, err := func() (*url.URL, error) {
evt, err := nextEvent()
var evt event
for evt, err = range scanEvents(resp.Body) {
break
}
if err != nil {
return nil, err
}
@ -413,8 +381,7 @@ func (c *SSEClientTransport) Connect(ctx context.Context) (Stream, error) {
go func() {
defer s.Close() // close the transport when the GET exits
for {
evt, err := nextEvent()
for evt, err := range scanEvents(resp.Body) {
if err != nil {
return
}
@ -429,6 +396,95 @@ func (c *SSEClientTransport) Connect(ctx context.Context) (Stream, error) {
return s, nil
}
// scanEvents iterates SSE events in the given scanner. The iterated error is
// terminal: if encountered, the stream is corrupt or broken and should no
// longer be used.
//
// TODO(rfindley): consider a different API here that makes failure modes more
// apparent.
func scanEvents(r io.Reader) iter.Seq2[event, error] {
scanner := bufio.NewScanner(r)
const maxTokenSize = 1 * 1024 * 1024 // 1 MiB max line size
scanner.Buffer(nil, maxTokenSize)
// TODO: investigate proper behavior when events are out of order, or have
// non-standard names.
var (
eventKey = []byte("event")
idKey = []byte("id")
dataKey = []byte("data")
)
return func(yield func(event, error) bool) {
// iterate event from the wire.
// https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events#examples
//
// - `key: value` line records.
// - Consecutive `data: ...` fields are joined with newlines.
// - Unrecognized fields are ignored. Since we only care about 'event', 'id', and
// 'data', these are the only three we consider.
// - Lines starting with ":" are ignored.
// - Records are terminated with two consecutive newlines.
var (
evt event
dataBuf *bytes.Buffer // if non-nil, preceding field was also data
)
flushData := func() {
if dataBuf != nil {
evt.data = dataBuf.Bytes()
dataBuf = nil
}
}
for scanner.Scan() {
line := scanner.Bytes()
if len(line) == 0 {
flushData()
// \n\n is the record delimiter
if !evt.empty() && !yield(evt, nil) {
return
}
evt = event{}
continue
}
before, after, found := bytes.Cut(line, []byte{':'})
if !found {
yield(event{}, fmt.Errorf("malformed line in SSE stream: %q", string(line)))
return
}
if !bytes.Equal(before, dataKey) {
flushData()
}
switch {
case bytes.Equal(before, eventKey):
evt.name = strings.TrimSpace(string(after))
case bytes.Equal(before, idKey):
evt.id = strings.TrimSpace(string(after))
case bytes.Equal(before, dataKey):
data := bytes.TrimSpace(after)
if dataBuf != nil {
dataBuf.WriteByte('\n')
dataBuf.Write(data)
} else {
dataBuf = new(bytes.Buffer)
dataBuf.Write(data)
}
}
}
if err := scanner.Err(); err != nil {
if errors.Is(err, bufio.ErrTooLong) {
err = fmt.Errorf("event exceeded max line length of %d", maxTokenSize)
}
if !yield(event{}, err) {
return
}
}
flushData()
if !evt.empty() {
yield(evt, nil)
}
}
}
// An sseClientStream is a logical jsonrpc2 stream that implements the client
// half of the SSE protocol:
// - Writes are POSTS to the session endpoint.
@ -460,6 +516,7 @@ func (c *sseClientStream) Read(ctx context.Context) (jsonrpc2.Message, error) {
return nil, io.EOF
case data := <-c.incoming:
// TODO(rfindley): do we really need to check this? We receive from c.done above.
if c.isDone() {
return nil, io.EOF
}

92
mcp/sse_test.go
View file

@ -9,6 +9,7 @@ import (
"fmt"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/google/go-cmp/cmp"
@ -46,7 +47,7 @@ func TestSSEServer(t *testing.T) {
ss := <-conns
gotHi, err := cs.CallTool(ctx, &CallToolParams{
Name: "greet",
Arguments: map[string]any{"name": "user"},
Arguments: map[string]any{"Name": "user"},
})
if err != nil {
t.Fatal(err)
@ -70,3 +71,92 @@ func TestSSEServer(t *testing.T) {
})
}
}
func TestScanEvents(t *testing.T) {
tests := []struct {
name string
input string
want []event
wantErr string
}{
{
name: "simple event",
input: "event: message\nid: 1\ndata: hello\n\n",
want: []event{
{name: "message", id: "1", data: []byte("hello")},
},
},
{
name: "multiple data lines",
input: "data: line 1\ndata: line 2\n\n",
want: []event{
{data: []byte("line 1\nline 2")},
},
},
{
name: "multiple events",
input: "data: first\n\nevent: second\ndata: second\n\n",
want: []event{
{data: []byte("first")},
{name: "second", data: []byte("second")},
},
},
{
name: "no trailing newline",
input: "data: hello",
want: []event{
{data: []byte("hello")},
},
},
{
name: "malformed line",
input: "invalid line\n\n",
wantErr: "malformed line",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
r := strings.NewReader(tt.input)
var got []event
var err error
for e, err2 := range scanEvents(r) {
if err2 != nil {
err = err2
break
}
got = append(got, e)
}
if tt.wantErr != "" {
if err == nil {
t.Fatalf("scanEvents() got nil error, want error containing %q", tt.wantErr)
}
if !strings.Contains(err.Error(), tt.wantErr) {
t.Fatalf("scanEvents() error = %q, want containing %q", err, tt.wantErr)
}
return
}
if err != nil {
t.Fatalf("scanEvents() returned unexpected error: %v", err)
}
if len(got) != len(tt.want) {
t.Fatalf("scanEvents() got %d events, want %d", len(got), len(tt.want))
}
for i := range got {
if g, w := got[i].name, tt.want[i].name; g != w {
t.Errorf("event %d: name = %q, want %q", i, g, w)
}
if g, w := got[i].id, tt.want[i].id; g != w {
t.Errorf("event %d: id = %q, want %q", i, g, w)
}
if g, w := string(got[i].data), string(tt.want[i].data); g != w {
t.Errorf("event %d: data = %q, want %q", i, g, w)
}
}
})
}
}

737
mcp/streamable.go Normal file
View file

@ -0,0 +1,737 @@
// 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 (
"bytes"
"context"
"fmt"
"io"
"net/http"
"strings"
"sync"
"sync/atomic"
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
)
// A StreamableHTTPHandler is an http.Handler that serves streamable MCP
// sessions, as defined by the [MCP spec].
//
// [MCP spec]: https://modelcontextprotocol.io/2025/03/26/streamable-http-transport.html
type StreamableHTTPHandler struct {
getServer func(*http.Request) *Server
sessionsMu sync.Mutex
sessions map[string]*StreamableServerTransport // keyed by IDs (from Mcp-Session-Id header)
}
// StreamableHTTPOptions is a placeholder options struct for future
// configuration of the StreamableHTTP handler.
type StreamableHTTPOptions struct {
// TODO(rfindley): support configurable session ID generation and event
// store, session retention, and event retention.
}
// NewStreamableHTTPHandler returns a new [StreamableHTTPHandler].
//
// The getServer function is used to create or look up servers for new
// sessions. It is OK for getServer to return the same server multiple times.
func NewStreamableHTTPHandler(getServer func(*http.Request) *Server, opts *StreamableHTTPOptions) *StreamableHTTPHandler {
return &StreamableHTTPHandler{
getServer: getServer,
sessions: make(map[string]*StreamableServerTransport),
}
}
// closeAll closes all ongoing sessions.
//
// TODO(rfindley): investigate the best API for callers to configure their
// session lifecycle.
//
// Should we allow passing in a session store? That would allow the handler to
// be stateless.
func (h *StreamableHTTPHandler) closeAll() {
h.sessionsMu.Lock()
defer h.sessionsMu.Unlock()
for _, s := range h.sessions {
s.Close()
}
h.sessions = nil
}
func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
// Allow multiple 'Accept' headers.
// https://developer.mozilla.org/en-US/docs/Web/HTTP/Reference/Headers/Accept#syntax
accept := strings.Split(strings.Join(req.Header.Values("Accept"), ","), ",")
var jsonOK, streamOK bool
for _, c := range accept {
switch strings.TrimSpace(c) {
case "application/json":
jsonOK = true
case "text/event-stream":
streamOK = true
}
}
if req.Method == http.MethodGet {
if !streamOK {
http.Error(w, "Accept must contain 'text/event-stream' for GET requests", http.StatusBadRequest)
return
}
} else if !jsonOK || !streamOK {
http.Error(w, "Accept must contain both 'application/json' and 'text/event-stream'", http.StatusBadRequest)
return
}
var session *StreamableServerTransport
if id := req.Header.Get("Mcp-Session-Id"); id != "" {
h.sessionsMu.Lock()
session, _ = h.sessions[id]
h.sessionsMu.Unlock()
if session == nil {
http.Error(w, "session not found", http.StatusNotFound)
return
}
}
// TODO(rfindley): simplify the locking so that each request has only one
// critical section.
if req.Method == http.MethodDelete {
if session == nil {
// => Mcp-Session-Id was not set; else we'd have returned NotFound above.
http.Error(w, "DELETE requires an Mcp-Session-Id header", http.StatusBadRequest)
return
}
h.sessionsMu.Lock()
delete(h.sessions, session.id)
h.sessionsMu.Unlock()
session.Close()
w.WriteHeader(http.StatusNoContent)
return
}
switch req.Method {
case http.MethodPost, http.MethodGet:
default:
w.Header().Set("Allow", "GET, POST")
http.Error(w, "unsupported method", http.StatusMethodNotAllowed)
return
}
if session == nil {
s := NewStreamableServerTransport(randText())
server := h.getServer(req)
// Pass req.Context() here, to allow middleware to add context values.
// The context is detached in the jsonrpc2 library when handling the
// long-running stream.
if _, err := server.Connect(req.Context(), s); err != nil {
http.Error(w, "failed connection", http.StatusInternalServerError)
return
}
h.sessionsMu.Lock()
h.sessions[s.id] = s
h.sessionsMu.Unlock()
session = s
}
session.ServeHTTP(w, req)
}
// NewStreamableServerTransport returns a new [StreamableServerTransport] with
// the given session ID.
//
// A StreamableServerTransport implements the server-side of the streamable
// transport.
//
// TODO(rfindley): consider adding options here, to configure event storage
// policy.
func NewStreamableServerTransport(sessionID string) *StreamableServerTransport {
return &StreamableServerTransport{
id: sessionID,
incoming: make(chan jsonrpc2.Message, 10),
done: make(chan struct{}),
outgoingMessages: make(map[connID][]*streamableMsg),
signals: make(map[connID]chan struct{}),
requestConns: make(map[JSONRPCID]connID),
outstandingRequests: make(map[connID]map[JSONRPCID]struct{}),
}
}
// A StreamableServerTransport implements the [Transport] interface for a
// single session.
type StreamableServerTransport struct {
nextConnID atomic.Int64 // incrementing next connection ID
id string
incoming chan jsonrpc2.Message // messages from the client to the server
mu sync.Mutex
// Sessions are closed exactly once.
isDone bool
done chan struct{}
// Sessions can have multiple logical connections, corresponding to HTTP
// requests. Additionally, logical sessions may be resumed by subsequent HTTP
// requests, when the session is terminated unexpectedly.
//
// Therefore, we use a logical connection ID to key the connection state, and
// perform the accounting described below when incoming HTTP requests are
// handled.
//
// The accounting is complicated. It is tempting to merge some of the maps
// below, but they each have different lifecycles, as indicated by Lifecycle:
// comments.
//
// TODO(rfindley): simplify.
// outgoingMessages is the collection of outgoingMessages messages, keyed by the logical
// connection ID where they should be delivered.
//
// connID 0 is used for messages that don't correlate with an incoming
// request.
//
// Lifecycle: outgoingMessages persists for the duration of the session.
//
// TODO(rfindley): garbage collect this data. For now, we save all outgoingMessages
// messages for the lifespan of the transport.
outgoingMessages map[connID][]*streamableMsg
// signals maps a logical connection ID to a 1-buffered channel, owned by an
// incoming HTTP request, that signals that there are messages available to
// write into the HTTP response. Signals guarantees that at most one HTTP
// connection can receive messages for a logical connection. After claiming
// the connection, incoming requests should read from outgoing, to ensure
// that no new messages are missed.
//
// Lifecycle: signals persists for the duration of an HTTP POST or GET
// request for the given connID.
signals map[connID]chan struct{}
// requestConns maps incoming requests to their logical connection ID.
//
// Lifecycle: requestConns persists for the duration of the session.
//
// TODO(rfindley): clean up once requests are handled.
requestConns map[JSONRPCID]connID
// outstandingRequests tracks the set of unanswered incoming RPCs for each logical
// connection.
//
// When the server has responded to each request, the connection should be
// closed.
//
// Lifecycle: outstandingRequests values persist as until the requests have been
// replied to by the server. Notably, NOT until they are sent to an HTTP
// response, as delivery is not guaranteed.
outstandingRequests map[connID]map[JSONRPCID]struct{}
}
type connID int64
// a streamableMsg is an SSE event with an index into its logical stream.
type streamableMsg struct {
idx int
event event
}
// Connect implements the [Transport] interface.
//
// TODO(rfindley): Connect should return a new object.
func (s *StreamableServerTransport) Connect(context.Context) (Stream, error) {
return s, nil
}
// We track the incoming request ID inside the handler context using
// idContextValue, so that notifications and server->client calls that occur in
// the course of handling incoming requests are correlated with the incoming
// request that caused them, and can be dispatched as server-sent events to the
// correct HTTP request.
//
// Currently, this is implemented in [ServerSession.handle]. This is not ideal,
// because it means that a user of the MCP package couldn't implement the
// streamable transport, as they'd lack this priviledged access.
//
// If we ever wanted to expose this mechanism, we have a few options:
// 1. Make ServerSession an interface, and provide an implementation of
// ServerSession to handlers that closes over the incoming request ID.
// 2. Expose a 'HandlerTransport' interface that allows transports to provide
// a handler middleware, so that we don't hard-code this behavior in
// ServerSession.handle.
// 3. Add a `func ForRequest(context.Context) JSONRPCID` accessor that lets
// any transport access the incoming request ID.
//
// For now, by giving only the StreamableServerTransport access to the request
// ID, we avoid having to make this API decision.
type idContextKey struct{}
// ServeHTTP handles a single HTTP request for the session.
func (t *StreamableServerTransport) ServeHTTP(w http.ResponseWriter, req *http.Request) {
switch req.Method {
case http.MethodGet:
t.serveGET(w, req)
case http.MethodPost:
t.servePOST(w, req)
default:
// Should not be reached, as this is checked in StreamableHTTPHandler.ServeHTTP.
w.Header().Set("Allow", "GET, POST")
http.Error(w, "unsupported method", http.StatusMethodNotAllowed)
}
}
func (t *StreamableServerTransport) serveGET(w http.ResponseWriter, req *http.Request) {
// connID 0 corresponds to the default GET request.
id, nextIdx := connID(0), 0
if len(req.Header.Values("Last-Event-ID")) > 0 {
eid := req.Header.Get("Last-Event-ID")
var ok bool
id, nextIdx, ok = parseEventID(eid)
if !ok {
http.Error(w, fmt.Sprintf("malformed Last-Event-ID %q", eid), http.StatusBadRequest)
return
}
nextIdx++
}
t.mu.Lock()
if _, ok := t.signals[id]; ok {
http.Error(w, "stream ID conflicts with ongoing stream", http.StatusBadRequest)
t.mu.Unlock()
return
}
signal := make(chan struct{}, 1)
t.signals[id] = signal
t.mu.Unlock()
t.streamResponse(w, req, id, nextIdx, signal)
}
func (t *StreamableServerTransport) servePOST(w http.ResponseWriter, req *http.Request) {
if len(req.Header.Values("Last-Event-ID")) > 0 {
http.Error(w, "can't send Last-Event-ID for POST request", http.StatusBadRequest)
return
}
// Read incoming messages.
body, err := io.ReadAll(req.Body)
if err != nil {
http.Error(w, "failed to read body", http.StatusBadRequest)
return
}
if len(body) == 0 {
http.Error(w, "POST requires a non-empty body", http.StatusBadRequest)
return
}
incoming, _, err := readBatch(body)
if err != nil {
http.Error(w, fmt.Sprintf("malformed payload: %v", err), http.StatusBadRequest)
return
}
var requests = make(map[JSONRPCID]struct{})
for _, msg := range incoming {
if req, ok := msg.(*JSONRPCRequest); ok && req.ID.IsValid() {
requests[req.ID] = struct{}{}
}
}
// Update accounting for this request.
id := connID(t.nextConnID.Add(1))
signal := make(chan struct{}, 1)
t.mu.Lock()
if len(requests) > 0 {
t.outstandingRequests[id] = make(map[JSONRPCID]struct{})
}
for reqID := range requests {
t.requestConns[reqID] = id
t.outstandingRequests[id][reqID] = struct{}{}
}
t.signals[id] = signal
t.mu.Unlock()
// Publish incoming messages.
for _, msg := range incoming {
t.incoming <- msg
}
// TODO(rfindley): consider optimizing for a single incoming request, by
// responding with application/json when there is only a single message in
// the response.
t.streamResponse(w, req, id, 0, signal)
}
func (t *StreamableServerTransport) streamResponse(w http.ResponseWriter, req *http.Request, id connID, nextIndex int, signal chan struct{}) {
defer func() {
t.mu.Lock()
delete(t.signals, id)
t.mu.Unlock()
}()
// Stream resumption: adjust outgoing index based on what the user says
// they've received.
if nextIndex > 0 {
t.mu.Lock()
// Clamp nextIndex to outgoing messages.
outgoing := t.outgoingMessages[id]
if nextIndex > len(outgoing) {
nextIndex = len(outgoing)
}
t.mu.Unlock()
}
w.Header().Set("Mcp-Session-Id", t.id)
w.Header().Set("Content-Type", "text/event-stream") // Accept checked in [StreamableHTTPHandler]
w.Header().Set("Cache-Control", "no-cache, no-transform")
w.Header().Set("Connection", "keep-alive")
writes := 0
stream:
for {
// Send outgoing messages
t.mu.Lock()
outgoing := t.outgoingMessages[id][nextIndex:]
t.mu.Unlock()
for _, msg := range outgoing {
if _, err := writeEvent(w, msg.event); err != nil {
// Connection closed or broken.
return
}
writes++
nextIndex++
}
t.mu.Lock()
nOutstanding := len(t.outstandingRequests[id])
nOutgoing := len(t.outgoingMessages[id])
t.mu.Unlock()
// If all requests have been handled and replied to, we can terminate this
// connection. However, in the case of a sequencing violation from the server
// (a send on the request context after the request has been handled), we
// loop until we've written all messages.
//
// TODO(rfindley): should we instead refuse to send messages after the last
// response? Decide, write a test, and change the behavior.
if nextIndex < nOutgoing {
continue // more to send
}
if req.Method == http.MethodPost && nOutstanding == 0 {
if writes == 0 {
// Spec: If the server accepts the input, the server MUST return HTTP
// status code 202 Accepted with no body.
w.WriteHeader(http.StatusAccepted)
}
return
}
select {
case <-signal:
case <-t.done:
if writes == 0 {
http.Error(w, "session terminated", http.StatusGone)
}
break stream
case <-req.Context().Done():
if writes == 0 {
w.WriteHeader(http.StatusNoContent)
}
break stream
}
}
}
// Event IDs: encode both the logical connection ID and the index, as
// <connID>_<idx>, to be consistent with the typescript implementation.
// formatEventID returns the event ID to use for the logical connection ID
// connID and message index idx.
//
// See also [parseEventID].
func formatEventID(id connID, idx int) string {
return fmt.Sprintf("%d_%d", id, idx)
}
// parseEventID parses a Last-Event-ID value into a logical connection id and
// index.
//
// See also [formatEventID].
func parseEventID(eventID string) (conn connID, idx int, ok bool) {
_, err := fmt.Sscanf(eventID, "%d_%d", &conn, &idx)
if err != nil || conn < 0 || idx < 0 {
return 0, 0, false
}
return conn, idx, true
}
// Read implements the [Stream] interface.
func (t *StreamableServerTransport) Read(ctx context.Context) (jsonrpc2.Message, error) {
select {
case <-ctx.Done():
return nil, ctx.Err()
case msg, ok := <-t.incoming:
if !ok {
return nil, io.EOF
}
return msg, nil
case <-t.done:
return nil, io.EOF
}
}
// Write implements the [Stream] interface.
func (t *StreamableServerTransport) Write(ctx context.Context, msg JSONRPCMessage) error {
// Find the incoming request that this write relates to, if any.
var forRequest, replyTo JSONRPCID
if resp, ok := msg.(*JSONRPCResponse); ok {
// If the message is a response, it relates to its request (of course).
forRequest = resp.ID
replyTo = resp.ID
} else {
// Otherwise, we check to see if it request was made in the context of an
// ongoing request. This may not be the case if the request way made with
// an unrelated context.
if v := ctx.Value(idContextKey{}); v != nil {
forRequest = v.(JSONRPCID)
}
}
// Find the logical connection corresponding to this request.
//
// For messages sent outside of a request context, this is the default
// connection 0.
var forConn connID
if forRequest.IsValid() {
t.mu.Lock()
forConn = t.requestConns[forRequest]
t.mu.Unlock()
}
data, err := jsonrpc2.EncodeMessage(msg)
if err != nil {
return err
}
t.mu.Lock()
defer t.mu.Unlock()
if t.isDone {
return fmt.Errorf("session is closed") // TODO: should this be EOF?
}
if _, ok := t.outstandingRequests[forConn]; !ok && forConn != 0 {
// No outstanding requests for this connection, which means it is logically
// done. This is a sequencing violation from the server, so we should report
// a side-channel error here. Put the message on the general queue to avoid
// dropping messages.
forConn = 0
}
idx := len(t.outgoingMessages[forConn])
t.outgoingMessages[forConn] = append(t.outgoingMessages[forConn], &streamableMsg{
idx: idx,
event: event{
name: "message",
id: formatEventID(forConn, idx),
data: data,
},
})
if replyTo.IsValid() {
// Once we've put the reply on the queue, it's no longer outstanding.
delete(t.outstandingRequests[forConn], replyTo)
if len(t.outstandingRequests[forConn]) == 0 {
delete(t.outstandingRequests, forConn)
}
}
// Signal work.
if c, ok := t.signals[forConn]; ok {
select {
case c <- struct{}{}:
default:
}
}
return nil
}
// Close implements the [Stream] interface.
func (t *StreamableServerTransport) Close() error {
t.mu.Lock()
defer t.mu.Unlock()
if !t.isDone {
t.isDone = true
close(t.done)
}
return nil
}
// A StreamableClientTransport is a [Transport] that can communicate with an MCP
// endpoint serving the streamable HTTP transport defined by the 2025-03-26
// version of the spec.
//
// TODO(rfindley): support retries and resumption tokens.
type StreamableClientTransport struct {
url string
}
// NewStreamableClientTransport returns a new client transport that connects to
// the streamable HTTP server at the provided URL.
func NewStreamableClientTransport(url string) *StreamableClientTransport {
return &StreamableClientTransport{url: url}
}
// Connect implements the [Transport] interface.
//
// The resulting Stream writes messages via POST requests to the transport URL
// with the Mcp-Session-Id header set, and reads messages from hanging
// requests.
//
// When closed, the stream issues a DELETE request to terminate the logical
// session.
func (t *StreamableClientTransport) Connect(ctx context.Context) (Stream, error) {
return &streamableClientStream{
url: t.url,
client: http.DefaultClient,
incoming: make(chan []byte, 100),
done: make(chan struct{}),
}, nil
}
type streamableClientStream struct {
url string
sessionID string
client *http.Client
incoming chan []byte
done chan struct{}
closeOnce sync.Once
closeErr error
mu sync.Mutex
// bodies map[*http.Response]io.Closer
err error
}
// Read implements the [Stream] interface.
func (s *streamableClientStream) Read(ctx context.Context) (jsonrpc2.Message, error) {
select {
case <-ctx.Done():
return nil, ctx.Err()
case <-s.done:
return nil, io.EOF
case data := <-s.incoming:
return jsonrpc2.DecodeMessage(data)
}
}
// Write implements the [Stream] interface.
func (s *streamableClientStream) Write(ctx context.Context, msg jsonrpc2.Message) error {
s.mu.Lock()
if s.err != nil {
s.mu.Unlock()
return s.err
}
sessionID := s.sessionID
if sessionID == "" {
// Hold lock for the first request.
defer s.mu.Unlock()
} else {
s.mu.Unlock()
}
gotSessionID, err := s.postMessage(ctx, sessionID, msg)
if err != nil {
if sessionID != "" {
// unlocked; lock to set err
s.mu.Lock()
defer s.mu.Unlock()
}
if s.err != nil {
s.err = err
}
return err
}
if sessionID == "" {
// locked
s.sessionID = gotSessionID
}
return nil
}
func (s *streamableClientStream) postMessage(ctx context.Context, sessionID string, msg JSONRPCMessage) (string, error) {
data, err := jsonrpc2.EncodeMessage(msg)
if err != nil {
return "", err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, s.url, bytes.NewReader(data))
if err != nil {
return "", err
}
if sessionID != "" {
req.Header.Set("Mcp-Session-Id", sessionID)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "application/json, text/event-stream")
resp, err := s.client.Do(req)
if err != nil {
return "", err
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
// TODO: do a best effort read of the body here, and format it in the error.
resp.Body.Close()
return "", fmt.Errorf("broken session: %v", resp.Status)
}
sessionID = resp.Header.Get("Mcp-Session-Id")
if resp.Header.Get("Content-Type") == "text/event-stream" {
go s.handleSSE(resp)
} else {
resp.Body.Close()
}
return sessionID, nil
}
func (s *streamableClientStream) handleSSE(resp *http.Response) {
defer resp.Body.Close()
done := make(chan struct{})
go func() {
defer close(done)
for evt, err := range scanEvents(resp.Body) {
if err != nil {
// TODO: surface this error; possibly break the stream
return
}
s.incoming <- evt.data
}
}()
select {
case <-s.done:
case <-done:
}
}
// Close implements the [Stream] interface.
func (s *streamableClientStream) Close() error {
s.closeOnce.Do(func() {
close(s.done)
req, err := http.NewRequest(http.MethodDelete, s.url, nil)
if err != nil {
s.closeErr = err
} else {
req.Header.Set("Mcp-Session-Id", s.sessionID)
if _, err := s.client.Do(req); err != nil {
s.closeErr = err
}
}
})
return s.closeErr
}

516
mcp/streamable_test.go Normal file
View file

@ -0,0 +1,516 @@
// 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 (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"net/http/httptest"
"strings"
"sync"
"sync/atomic"
"testing"
"github.com/google/go-cmp/cmp"
"github.com/google/go-cmp/cmp/cmpopts"
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
)
func TestStreamableTransports(t *testing.T) {
// This test checks that the streamable server and client transports can
// communicate.
ctx := context.Background()
// 1. Create a server with a simple "greet" tool.
server := NewServer("testServer", "v1.0.0", nil)
server.AddTools(NewTool("greet", "say hi", sayHi))
// 2. Start an httptest.Server with the StreamableHTTPHandler.
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil)
httpServer := httptest.NewServer(handler)
defer httpServer.Close()
// 3. Create a client and connect it to the server using our StreamableClientTransport.
transport := NewStreamableClientTransport(httpServer.URL)
client := NewClient("testClient", "v1.0.0", nil)
session, err := client.Connect(ctx, transport)
if err != nil {
t.Fatalf("client.Connect() failed: %v", err)
}
defer session.Close()
// 4. The client calls the "greet" tool.
params := &CallToolParams{
Name: "greet",
Arguments: map[string]any{"name": "streamy"},
}
got, err := session.CallTool(ctx, params)
if err != nil {
t.Fatalf("CallTool() failed: %v", err)
}
// 5. Verify that the correct response is received.
want := &CallToolResult{
Content: []*Content{{Type: "text", Text: "hi streamy"}},
}
if diff := cmp.Diff(want, got); diff != "" {
t.Errorf("CallTool() returned unexpected content (-want +got):\n%s", diff)
}
}
func TestStreamableServerTransport(t *testing.T) {
// This test checks detailed behavior of the streamable server transport, by
// faking the behavior of a streamable client using a sequence of HTTP
// requests.
// A step is a single step in the tests below, consisting of a request payload
// and expected response.
type step struct {
// If OnRequest is > 0, this step only executes after a request with the
// given ID is received.
//
// All OnRequest steps must occur before the step that creates the request.
//
// To avoid tests hanging when there's a bug, it's expected that this
// request is received in the course of a *synchronous* request to the
// server (otherwise, we wouldn't be able to terminate the test without
// analyzing a dependency graph).
OnRequest int64
// If set, Async causes the step to run asynchronously to other steps.
// Redundant with OnRequest: all OnRequest steps are asynchronous.
Async bool
Method string // HTTP request method
Send []JSONRPCMessage // messages to send
CloseAfter int // if nonzero, close after receiving this many messages
StatusCode int // expected status code
Recv []JSONRPCMessage // expected messages to receive
}
// JSON-RPC message constructors.
req := func(id int64, method string, params any) *JSONRPCRequest {
r := &JSONRPCRequest{
Method: method,
Params: mustMarshal(t, params),
}
if id > 0 {
r.ID = jsonrpc2.Int64ID(id)
}
return r
}
resp := func(id int64, result any, err error) *JSONRPCResponse {
return &JSONRPCResponse{
ID: jsonrpc2.Int64ID(id),
Result: mustMarshal(t, result),
Error: err,
}
}
// Predefined steps, to avoid repetition below.
initReq := req(1, "initialize", &InitializeParams{})
initResp := resp(1, &InitializeResult{
Capabilities: &serverCapabilities{
Logging: &loggingCapabilities{},
Prompts: &promptCapabilities{ListChanged: true},
Resources: &resourceCapabilities{ListChanged: true},
Tools: &toolCapabilities{ListChanged: true},
},
ProtocolVersion: "2024-11-05",
ServerInfo: &implementation{Name: "testServer", Version: "v1.0.0"},
}, nil)
initializedMsg := req(0, "initialized", &InitializedParams{})
initialize := step{
Method: "POST",
Send: []JSONRPCMessage{initReq},
StatusCode: http.StatusOK,
Recv: []JSONRPCMessage{initResp},
}
initialized := step{
Method: "POST",
Send: []JSONRPCMessage{initializedMsg},
StatusCode: http.StatusAccepted,
}
tests := []struct {
name string
tool func(*testing.T, context.Context, *ServerSession)
steps []step
}{
{
name: "basic",
steps: []step{
initialize,
initialized,
{
Method: "POST",
Send: []JSONRPCMessage{req(2, "tools/call", &CallToolParams{Name: "tool"})},
StatusCode: http.StatusOK,
Recv: []JSONRPCMessage{resp(2, &CallToolResult{}, nil)},
},
},
},
{
name: "tool notification",
tool: func(t *testing.T, ctx context.Context, ss *ServerSession) {
// Send an arbitrary notification.
if err := ss.NotifyProgress(ctx, &ProgressNotificationParams{}); err != nil {
t.Errorf("Notify failed: %v", err)
}
},
steps: []step{
initialize,
initialized,
{
Method: "POST",
Send: []JSONRPCMessage{
req(2, "tools/call", &CallToolParams{Name: "tool"}),
},
StatusCode: http.StatusOK,
Recv: []JSONRPCMessage{
req(0, "notifications/progress", &ProgressNotificationParams{}),
resp(2, &CallToolResult{}, nil),
},
},
},
},
{
name: "tool upcall",
tool: func(t *testing.T, ctx context.Context, ss *ServerSession) {
// Make an arbitrary call.
if _, err := ss.ListRoots(ctx, &ListRootsParams{}); err != nil {
t.Errorf("Call failed: %v", err)
}
},
steps: []step{
initialize,
initialized,
{
Method: "POST",
OnRequest: 1,
Send: []JSONRPCMessage{
resp(1, &ListRootsResult{}, nil),
},
StatusCode: http.StatusAccepted,
},
{
Method: "POST",
Send: []JSONRPCMessage{
req(2, "tools/call", &CallToolParams{Name: "tool"}),
},
StatusCode: http.StatusOK,
Recv: []JSONRPCMessage{
req(1, "roots/list", &ListRootsParams{}),
resp(2, &CallToolResult{}, nil),
},
},
},
},
{
name: "background",
tool: func(t *testing.T, ctx context.Context, ss *ServerSession) {
// Perform operations on a background context, and ensure the client
// receives it.
ctx = context.Background()
if err := ss.NotifyProgress(ctx, &ProgressNotificationParams{}); err != nil {
t.Errorf("Notify failed: %v", err)
}
// TODO(rfindley): finish implementing logging.
// if err := ss.LoggingMessage(ctx, &LoggingMessageParams{}); err != nil {
// t.Errorf("Logging failed: %v", err)
// }
if _, err := ss.ListRoots(ctx, &ListRootsParams{}); err != nil {
t.Errorf("Notify failed: %v", err)
}
},
steps: []step{
initialize,
initialized,
{
Method: "POST",
OnRequest: 1,
Send: []JSONRPCMessage{
resp(1, &ListRootsResult{}, nil),
},
StatusCode: http.StatusAccepted,
},
{
Method: "GET",
Async: true,
StatusCode: http.StatusOK,
CloseAfter: 2,
Recv: []JSONRPCMessage{
req(0, "notifications/progress", &ProgressNotificationParams{}),
req(1, "roots/list", &ListRootsParams{}),
},
},
{
Method: "POST",
Send: []JSONRPCMessage{
req(2, "tools/call", &CallToolParams{Name: "tool"}),
},
StatusCode: http.StatusOK,
Recv: []JSONRPCMessage{
resp(2, &CallToolResult{}, nil),
},
},
},
},
{
name: "errors",
steps: []step{
{
Method: "PUT",
StatusCode: http.StatusMethodNotAllowed,
},
{
Method: "DELETE",
StatusCode: http.StatusBadRequest,
},
{
Method: "POST",
Send: []JSONRPCMessage{req(2, "tools/call", &CallToolParams{Name: "tool"})},
StatusCode: http.StatusOK,
Recv: []JSONRPCMessage{resp(2, nil, &jsonrpc2.WireError{
Message: `method "tools/call" is invalid during session initialization`,
})},
},
},
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
// Create a server containing a single tool, which runs the test tool
// behavior, if any.
server := NewServer("testServer", "v1.0.0", nil)
tool := NewTool("tool", "test tool", func(ctx context.Context, ss *ServerSession, params *CallToolParamsFor[any]) (*CallToolResultFor[any], error) {
if test.tool != nil {
test.tool(t, ctx, ss)
}
return &CallToolResultFor[any]{}, nil
})
server.AddTools(tool)
// Start the streamable handler.
handler := NewStreamableHTTPHandler(func(req *http.Request) *Server { return server }, nil)
defer handler.closeAll()
httpServer := httptest.NewServer(handler)
defer httpServer.Close()
// blocks records request blocks by JSONRPC ID.
//
// When an OnRequest step is encountered, it waits on the corresponding
// block. When a request with that ID is received, the block is closed.
var mu sync.Mutex
blocks := make(map[int64]chan struct{})
for _, step := range test.steps {
if step.OnRequest > 0 {
blocks[step.OnRequest] = make(chan struct{})
}
}
// signal when all synchronous requests have executed, so we can fail
// async requests that are blocked.
syncRequestsDone := make(chan struct{})
// To avoid complicated accounting for session ID, just set the first
// non-empty session ID from a response.
var sessionID atomic.Value
sessionID.Store("")
// doStep executes a single step.
doStep := func(t *testing.T, step step) {
if step.OnRequest > 0 {
// Block the step until we've received the server->client request.
mu.Lock()
block := blocks[step.OnRequest]
mu.Unlock()
select {
case <-block:
case <-syncRequestsDone:
t.Errorf("after all sync requests are complete, request still blocked on %d", step.OnRequest)
return
}
}
// Collect messages received during this request, unblock other steps
// when requests are received.
var got []JSONRPCMessage
out := make(chan JSONRPCMessage)
// Cancel the step if we encounter a request that isn't going to be
// handled.
ctx, cancel := context.WithCancel(context.Background())
var wg sync.WaitGroup
wg.Add(1)
go func() {
defer wg.Done()
for m := range out {
if req, ok := m.(*JSONRPCRequest); ok && req.ID.IsValid() {
// Encountered a server->client request. We should have a
// response queued. Otherwise, we may deadlock.
mu.Lock()
if block, ok := blocks[req.ID.Raw().(int64)]; ok {
close(block)
} else {
t.Errorf("no queued response for %v", req.ID)
cancel()
}
mu.Unlock()
}
got = append(got, m)
if step.CloseAfter > 0 && len(got) == step.CloseAfter {
cancel()
}
}
}()
gotSessionID, gotStatusCode, err := streamingRequest(ctx,
httpServer.URL, sessionID.Load().(string), step.Method, step.Send, out)
// Don't fail on cancelled requests: error (if any) is handled
// elsewhere.
if err != nil && ctx.Err() == nil {
t.Fatal(err)
}
if gotStatusCode != step.StatusCode {
t.Errorf("got status %d, want %d", gotStatusCode, step.StatusCode)
}
wg.Wait()
transform := cmpopts.AcyclicTransformer("jsonrpcid", func(id JSONRPCID) any { return id.Raw() })
if diff := cmp.Diff(step.Recv, got, transform); diff != "" {
t.Errorf("received unexpected messages (-want +got):\n%s", diff)
}
sessionID.CompareAndSwap("", gotSessionID)
}
var wg sync.WaitGroup
for _, step := range test.steps {
if step.Async || step.OnRequest > 0 {
wg.Add(1)
go func() {
defer wg.Done()
doStep(t, step)
}()
} else {
doStep(t, step)
}
}
// Fail any blocked responses if they weren't needed by a synchronous
// request.
close(syncRequestsDone)
wg.Wait()
})
}
}
// streamingRequest makes a request to the given streamable server with the
// given url, sessionID, and method.
//
// If provided, the in messages are encoded in the request body. A single
// message is encoded as a JSON object. Multiple messages are batched as a JSON
// array.
//
// Any received messages are sent to the out channel, which is closed when the
// request completes.
//
// Returns the sessionID and http status code from the response. If an error is
// returned, sessionID and status code may still be set if the error occurs
// after the response headers have been received.
func streamingRequest(ctx context.Context, serverURL, sessionID, method string, in []JSONRPCMessage, out chan<- JSONRPCMessage) (string, int, error) {
defer close(out)
var body []byte
if len(in) == 1 {
data, err := jsonrpc2.EncodeMessage(in[0])
if err != nil {
return "", 0, fmt.Errorf("encoding message: %w", err)
}
body = data
} else {
var rawMsgs []json.RawMessage
for _, msg := range in {
data, err := jsonrpc2.EncodeMessage(msg)
if err != nil {
return "", 0, fmt.Errorf("encoding message: %w", err)
}
rawMsgs = append(rawMsgs, data)
}
data, err := json.Marshal(rawMsgs)
if err != nil {
return "", 0, fmt.Errorf("marshaling batch: %w", err)
}
body = data
}
req, err := http.NewRequestWithContext(ctx, method, serverURL, bytes.NewReader(body))
if err != nil {
return "", 0, fmt.Errorf("creating request: %w", err)
}
if sessionID != "" {
req.Header.Set("Mcp-Session-Id", sessionID)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Add("Accept", "text/plain") // ensure multiple accept headers are allowed
req.Header.Add("Accept", "application/json, text/event-stream")
resp, err := http.DefaultClient.Do(req)
if err != nil {
return "", 0, fmt.Errorf("request failed: %v", err)
}
defer resp.Body.Close()
newSessionID := resp.Header.Get("Mcp-Session-Id")
if strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream") {
for evt, err := range scanEvents(resp.Body) {
if err != nil {
return newSessionID, resp.StatusCode, fmt.Errorf("reading events: %v", err)
}
// TODO(rfindley): do we need to check evt.name?
// Does the MCP spec say anything about this?
msg, err := jsonrpc2.DecodeMessage(evt.data)
if err != nil {
return newSessionID, resp.StatusCode, fmt.Errorf("decoding message: %w", err)
}
out <- msg
}
} else if strings.HasPrefix(resp.Header.Get("Content-Type"), "application/json") {
data, err := io.ReadAll(resp.Body)
if err != nil {
return newSessionID, resp.StatusCode, fmt.Errorf("reading json body: %w", err)
}
msg, err := jsonrpc2.DecodeMessage(data)
if err != nil {
return newSessionID, resp.StatusCode, fmt.Errorf("decoding message: %w", err)
}
out <- msg
}
return newSessionID, resp.StatusCode, nil
}
func mustMarshal(t *testing.T, v any) json.RawMessage {
if v == nil {
return nil
}
t.Helper()
data, err := json.Marshal(v)
if err != nil {
t.Fatal(err)
}
return data
}

28
mcp/tool.go
View file

@ -20,8 +20,8 @@ import (
// Perhaps this should be an alias for ToolHandlerFor[map[string]any, map[string]any].
type ToolHandler func(context.Context, *ServerSession, *CallToolParamsFor[map[string]any]) (*CallToolResult, error)
// A ToolHandlerFor[Treq, Tres] handles a call to tools/call with typed arguments and results.
type ToolHandlerFor[Treq, Tres any] func(context.Context, *ServerSession, *CallToolParamsFor[Treq]) (*CallToolResultFor[Tres], error)
// A ToolHandlerFor handles a call to tools/call with typed arguments and results.
type ToolHandlerFor[In, Out any] func(context.Context, *ServerSession, *CallToolParamsFor[In]) (*CallToolResultFor[Out], error)
// A rawToolHandler is like a ToolHandler, but takes the arguments as as json.RawMessage.
type rawToolHandler = func(context.Context, *ServerSession, *CallToolParamsFor[json.RawMessage]) (*CallToolResult, error)
@ -44,17 +44,17 @@ type ServerTool struct {
// 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. This
// schema may be customized using the [Input] option.
func NewTool[TReq, TRes any](name, description string, handler ToolHandlerFor[TReq, TRes], opts ...ToolOption) *ServerTool {
st, err := newToolErr[TReq, TRes](name, description, handler, opts...)
func NewTool[In, Out any](name, description string, handler ToolHandlerFor[In, Out], opts ...ToolOption) *ServerTool {
st, err := newToolErr[In, Out](name, description, handler, opts...)
if err != nil {
panic(fmt.Errorf("NewTool(%q): %w", name, err))
}
return st
}
func newToolErr[TReq, TRes any](name, description string, handler ToolHandlerFor[TReq, TRes], opts ...ToolOption) (*ServerTool, error) {
func newToolErr[In, Out any](name, description string, handler ToolHandlerFor[In, Out], opts ...ToolOption) (*ServerTool, error) {
// TODO: check that TReq is a struct.
ischema, err := jsonschema.For[TReq]()
ischema, err := jsonschema.For[In]()
if err != nil {
return nil, err
}
@ -77,12 +77,14 @@ func newToolErr[TReq, TRes any](name, description string, handler ToolHandlerFor
}
t.rawHandler = func(ctx context.Context, ss *ServerSession, rparams *CallToolParamsFor[json.RawMessage]) (*CallToolResult, error) {
var args TReq
if err := unmarshalSchema(rparams.Arguments, t.inputResolved, &args); err != nil {
return nil, err
var args In
if rparams.Arguments != nil {
if err := unmarshalSchema(rparams.Arguments, t.inputResolved, &args); err != nil {
return nil, err
}
}
// TODO(jba): future-proof this copy.
params := &CallToolParamsFor[TReq]{
params := &CallToolParamsFor[In]{
Meta: rparams.Meta,
Name: rparams.Name,
Arguments: args,
@ -114,8 +116,10 @@ func newRawHandler(st *ServerTool) rawToolHandler {
return func(ctx context.Context, ss *ServerSession, rparams *CallToolParamsFor[json.RawMessage]) (*CallToolResult, error) {
// Unmarshal the args into what should be a map.
var args map[string]any
if err := unmarshalSchema(rparams.Arguments, st.inputResolved, &args); err != nil {
return nil, err
if rparams.Arguments != nil {
if err := unmarshalSchema(rparams.Arguments, st.inputResolved, &args); err != nil {
return nil, err
}
}
// TODO: generate copy
params := &CallToolParamsFor[map[string]any]{

133
mcp/transport.go
View file

@ -30,15 +30,28 @@ var ErrConnectionClosed = errors.New("connection closed")
type Transport interface {
// Connect returns the logical stream.
//
// It is called exactly once by [Connect].
// It is called exactly once by [Server.Connect] or [Client.Connect].
Connect(ctx context.Context) (Stream, error)
}
type (
// JSONRPCID is a JSON-RPC request ID.
JSONRPCID = jsonrpc2.ID
// JSONRPCHandler is a JSON-RPC handler function.
JSONRPCHandler = func(ctx context.Context, req *JSONRPCRequest) (result any, err error)
// JSONRPCMessage is a JSON-RPC message.
JSONRPCMessage = jsonrpc2.Message
// JSONRPCRequest is a JSON-RPC request.
JSONRPCRequest = jsonrpc2.Request
// JSONRPCResponse is a JSON-RPC response.
JSONRPCResponse = jsonrpc2.Response
)
// A Stream is a bidirectional jsonrpc2 Stream.
type Stream interface {
jsonrpc2.Reader
jsonrpc2.Writer
io.Closer
Read(context.Context) (JSONRPCMessage, error)
Write(context.Context, JSONRPCMessage) error
Close() error // may be called concurrently by both peers
}
// A StdIOTransport is a [Transport] that communicates over stdin/stdout using
@ -76,16 +89,15 @@ func NewInMemoryTransports() (*InMemoryTransport, *InMemoryTransport) {
return &InMemoryTransport{ioTransport{c1}}, &InMemoryTransport{ioTransport{c2}}
}
// handler is an unexported version of jsonrpc2.Handler.
type handler interface {
handle(ctx context.Context, req *jsonrpc2.Request) (result any, err error)
}
type binder[T handler] interface {
bind(*jsonrpc2.Connection) T
disconnect(T)
}
type handler interface {
handle(ctx context.Context, req *jsonrpc2.Request) (any, error)
}
func connect[H handler](ctx context.Context, t Transport, b binder[H]) (H, error) {
var zero H
stream, err := t.Connect(ctx)
@ -360,73 +372,64 @@ func (t *ioStream) read(ctx context.Context, in *json.Decoder) (jsonrpc2.Message
t.queue = t.queue[1:]
return next, nil
}
var raw json.RawMessage
if err := in.Decode(&raw); err != nil {
return nil, err
}
var rawBatch []json.RawMessage
if err := json.Unmarshal(raw, &rawBatch); err == nil {
msg, err := t.readBatch(rawBatch)
if err != nil {
return nil, err
}
return msg, nil
msgs, batch, err := readBatch(raw)
if err != nil {
return nil, err
}
msg, err := jsonrpc2.DecodeMessage(raw)
return msg, err
t.queue = msgs[1:]
if batch {
var respBatch *msgBatch // track incoming requests in the batch
for _, msg := range msgs {
if req, ok := msg.(*jsonrpc2.Request); ok {
if respBatch == nil {
respBatch = &msgBatch{
unresolved: make(map[jsonrpc2.ID]int),
}
}
if _, ok := respBatch.unresolved[req.ID]; ok {
return nil, fmt.Errorf("duplicate message ID %q", req.ID)
}
respBatch.unresolved[req.ID] = len(respBatch.responses)
respBatch.responses = append(respBatch.responses, nil)
}
}
if respBatch != nil {
// The batch contains one or more incoming requests to track.
if err := t.addBatch(respBatch); err != nil {
return nil, err
}
}
}
return msgs[0], err
}
// readBatch reads a batch of jsonrpc2 messages, and records the batch
// in the framer so that responses can be collected and send back together.
func (t *ioStream) readBatch(rawBatch []json.RawMessage) (jsonrpc2.Message, error) {
if len(rawBatch) == 0 {
return nil, fmt.Errorf("empty batch")
}
// From the spec:
// "If the batch rpc call itself fails to be recognized as an valid JSON or
// as an Array with at least one value, the response from the Server MUST be
// a single Response object. If there are no Response objects contained
// within the Response array as it is to be sent to the client, the server
// MUST NOT return an empty Array and should return nothing at all."
//
// In our case, an error actually breaks the jsonrpc2 connection entirely,
// but defensively we collect batch information before recording it, so that
// we don't leave the framer in an inconsistent state.
var (
first jsonrpc2.Message // first message, to return
queue []jsonrpc2.Message // remaining messages
respBatch *msgBatch // tracks incoming requests in the batch
)
for i, raw := range rawBatch {
msg, err := jsonrpc2.DecodeMessage(raw)
if err != nil {
return nil, err
// readBatch reads batch data, which may be either a single JSON-RPC message,
// or an array of JSON-RPC messages.
func readBatch(data []byte) (msgs []jsonrpc2.Message, isBatch bool, _ error) {
// Try to read an array of messages first.
var rawBatch []json.RawMessage
if err := json.Unmarshal(data, &rawBatch); err == nil {
if len(rawBatch) == 0 {
return nil, true, fmt.Errorf("empty batch")
}
if i == 0 {
first = msg
} else {
queue = append(queue, msg)
}
if req, ok := msg.(*jsonrpc2.Request); ok {
if respBatch == nil {
respBatch = &msgBatch{
unresolved: make(map[jsonrpc2.ID]int),
}
for _, raw := range rawBatch {
msg, err := jsonrpc2.DecodeMessage(raw)
if err != nil {
return nil, true, err
}
respBatch.unresolved[req.ID] = len(respBatch.responses)
respBatch.responses = append(respBatch.responses, nil)
msgs = append(msgs, msg)
}
return msgs, true, nil
}
if respBatch != nil {
// The batch contains one or more incoming requests to track.
if err := t.addBatch(respBatch); err != nil {
return nil, err
}
}
t.queue = append(t.queue, queue...)
return first, nil
// Try again with a single message.
msg, err := jsonrpc2.DecodeMessage(data)
return []jsonrpc2.Message{msg}, false, err
}
func (t *ioStream) Write(ctx context.Context, msg jsonrpc2.Message) error {