diff --git a/design/design.md b/design/design.md index 6fcd033..0219567 100644 --- a/design/design.md +++ b/design/design.md @@ -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) diff --git a/examples/hello/main.go b/examples/hello/main.go index b8c02b8..0491d16 100644 --- a/examples/hello/main.go +++ b/examples/hello/main.go @@ -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) } } } diff --git a/mcp/prompt.go b/mcp/prompt.go index 5171a04..fce2e5e 100644 --- a/mcp/prompt.go +++ b/mcp/prompt.go @@ -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 } diff --git a/mcp/server.go b/mcp/server.go index 13f1cba..15e7cdc 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -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) } diff --git a/mcp/sse.go b/mcp/sse.go index 31dbb34..578605f 100644 --- a/mcp/sse.go +++ b/mcp/sse.go @@ -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 } diff --git a/mcp/sse_test.go b/mcp/sse_test.go index 5ec06cd..efbadcc 100644 --- a/mcp/sse_test.go +++ b/mcp/sse_test.go @@ -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) + } + } + }) + } +} diff --git a/mcp/streamable.go b/mcp/streamable.go new file mode 100644 index 0000000..4e6d386 --- /dev/null +++ b/mcp/streamable.go @@ -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 +// _, 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 +} diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go new file mode 100644 index 0000000..0b13256 --- /dev/null +++ b/mcp/streamable_test.go @@ -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 +} diff --git a/mcp/tool.go b/mcp/tool.go index ced7896..116eac0 100644 --- a/mcp/tool.go +++ b/mcp/tool.go @@ -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]{ diff --git a/mcp/transport.go b/mcp/transport.go index 91fcf10..3e31c59 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -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 {