internal/mcp: implement streamable HTTP transport
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
Robert Findley
parent
11b17f5858
commit
d2f2f87bbe
1 changed files with 1570 additions and 149 deletions
737
mcp/streamable.go
Normal file
737
mcp/streamable.go
Normal 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
|
||||
}
|
||||
Loading…
Reference in a new issue