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
169
mcp/sse.go
169
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
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue