internal/mcp: rename Stream->Connection
After implementing the Streamable transport, it is apparent to me that the name 'Stream' is confusing, as it can refer either to a stream in the streamable transport, or to a JSON-RPC stream. But the name 'Connection' is available, and a natural name for the result of the Transport.Connect interface method. Therefore, rename Stream to Connection, and update the streamable implementation to refer to its HTTP connections as 'streams', consistent with other SDKs. Change-Id: Ica55f84caabcbfa5a2fe029cc0c2db57da574875 Reviewed-on: https://go-review.googlesource.com/c/tools/+/682376 LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com> Reviewed-by: Jonathan Amsterdam <jba@google.com>
This commit is contained in:
Rob Findley
Robert Findley
parent
9453d6bf6c
commit
1fead0ba64
7 changed files with 127 additions and 142 deletions
|
|
@ -29,7 +29,7 @@ func NewCommandTransport(cmd *exec.Cmd) *CommandTransport {
|
|||
}
|
||||
|
||||
// Connect starts the command, and connects to it over stdin/stdout.
|
||||
func (t *CommandTransport) Connect(ctx context.Context) (Stream, error) {
|
||||
func (t *CommandTransport) Connect(ctx context.Context) (Connection, error) {
|
||||
stdout, err := t.cmd.StdoutPipe()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
|
@ -42,7 +42,7 @@ func (t *CommandTransport) Connect(ctx context.Context) (Stream, error) {
|
|||
if err := t.cmd.Start(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return newIOStream(&pipeRWC{t.cmd, stdout, stdin}), nil
|
||||
return newIOConn(&pipeRWC{t.cmd, stdout, stdin}), nil
|
||||
}
|
||||
|
||||
// A pipeRWC is an io.ReadWriteCloser that communicates with a subprocess over
|
||||
|
|
|
|||
|
|
@ -24,7 +24,7 @@
|
|||
// instances using the [Server.Connect] method, which creates a
|
||||
// [ServerSession].
|
||||
//
|
||||
// A [Transport] connects a bidirectional [Stream] of jsonrpc2 messages. In
|
||||
// A [Transport] connects a bidirectional [Connection] of jsonrpc2 messages. In
|
||||
// practice, transports in the MCP spec are are either client transports or
|
||||
// server transports. For example, the [StdIOTransport] is a server transport
|
||||
// that communicates over stdin/stdout, and its counterpart is a
|
||||
|
|
|
|||
|
|
@ -597,8 +597,8 @@ func TestBatching(t *testing.T) {
|
|||
c := NewClient("testClient", "v1.0.0", nil)
|
||||
// TODO: this test is broken, because increasing the batch size here causes
|
||||
// 'initialize' to block. Therefore, we can only test with a size of 1.
|
||||
// Since batching is being removed, we can probably just delete this.
|
||||
const batchSize = 1
|
||||
BatchSize(ct, batchSize)
|
||||
cs, err := c.Connect(ctx, ct)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
|
|
|
|||
38
mcp/sse.go
38
mcp/sse.go
|
|
@ -104,7 +104,7 @@ func NewSSEHandler(getServer func(request *http.Request) *Server) *SSEHandler {
|
|||
// A SSEServerTransport is a logical SSE session created through a hanging GET
|
||||
// request.
|
||||
//
|
||||
// When connected, it returns the following [Stream] implementation:
|
||||
// When connected, it returns the following [Connection] implementation:
|
||||
// - Writes are SSE 'message' events to the GET response.
|
||||
// - Reads are received from POSTs to the session endpoint, via
|
||||
// [SSEServerTransport.ServeHTTP].
|
||||
|
|
@ -120,7 +120,7 @@ type SSEServerTransport struct {
|
|||
mu sync.Mutex
|
||||
w http.ResponseWriter // the hanging response body
|
||||
closed bool // set when the stream is closed
|
||||
done chan struct{} // closed when the stream is closed
|
||||
done chan struct{} // closed when the connection is closed
|
||||
}
|
||||
|
||||
// NewSSEServerTransport creates a new SSE transport for the given messages
|
||||
|
|
@ -168,8 +168,8 @@ func (t *SSEServerTransport) ServeHTTP(w http.ResponseWriter, req *http.Request)
|
|||
}
|
||||
|
||||
// Connect sends the 'endpoint' event to the client.
|
||||
// See [SSEServerTransport] for more details on the [Stream] implementation.
|
||||
func (t *SSEServerTransport) Connect(context.Context) (Stream, error) {
|
||||
// See [SSEServerTransport] for more details on the [Connection] implementation.
|
||||
func (t *SSEServerTransport) Connect(context.Context) (Connection, error) {
|
||||
t.mu.Lock()
|
||||
_, err := writeEvent(t.w, event{
|
||||
name: "endpoint",
|
||||
|
|
@ -179,7 +179,7 @@ func (t *SSEServerTransport) Connect(context.Context) (Stream, error) {
|
|||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return sseServerStream{t}, nil
|
||||
return sseServerConn{t}, nil
|
||||
}
|
||||
|
||||
func (h *SSEHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
|
||||
|
|
@ -257,14 +257,14 @@ func (h *SSEHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
|
|||
}
|
||||
}
|
||||
|
||||
// sseServerStream implements the Stream interface for a single [SSEServerTransport].
|
||||
// It hides the Stream interface from the SSEServerTransport API.
|
||||
type sseServerStream struct {
|
||||
// sseServerConn implements the [Connection] interface for a single [SSEServerTransport].
|
||||
// It hides the Connection interface from the SSEServerTransport API.
|
||||
type sseServerConn struct {
|
||||
t *SSEServerTransport
|
||||
}
|
||||
|
||||
// Read implements jsonrpc2.Reader.
|
||||
func (s sseServerStream) Read(ctx context.Context) (jsonrpc2.Message, error) {
|
||||
func (s sseServerConn) Read(ctx context.Context) (jsonrpc2.Message, error) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
|
|
@ -276,7 +276,7 @@ func (s sseServerStream) Read(ctx context.Context) (jsonrpc2.Message, error) {
|
|||
}
|
||||
|
||||
// Write implements jsonrpc2.Writer.
|
||||
func (s sseServerStream) Write(ctx context.Context, msg jsonrpc2.Message) error {
|
||||
func (s sseServerConn) Write(ctx context.Context, msg jsonrpc2.Message) error {
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
|
|
@ -305,7 +305,7 @@ func (s sseServerStream) Write(ctx context.Context, msg jsonrpc2.Message) error
|
|||
// It must be safe to call Close more than once, as the close may
|
||||
// asynchronously be initiated by either the server closing its connection, or
|
||||
// by the hanging GET exiting.
|
||||
func (s sseServerStream) Close() error {
|
||||
func (s sseServerConn) Close() error {
|
||||
s.t.mu.Lock()
|
||||
defer s.t.mu.Unlock()
|
||||
if !s.t.closed {
|
||||
|
|
@ -339,7 +339,7 @@ func NewSSEClientTransport(baseURL string) *SSEClientTransport {
|
|||
}
|
||||
|
||||
// Connect connects through the client endpoint.
|
||||
func (c *SSEClientTransport) Connect(ctx context.Context) (Stream, error) {
|
||||
func (c *SSEClientTransport) Connect(ctx context.Context) (Connection, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, "GET", c.sseEndpoint.String(), nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
|
@ -370,7 +370,7 @@ func (c *SSEClientTransport) Connect(ctx context.Context) (Stream, error) {
|
|||
}
|
||||
|
||||
// From here on, the stream takes ownership of resp.Body.
|
||||
s := &sseClientStream{
|
||||
s := &sseClientConn{
|
||||
sseEndpoint: c.sseEndpoint,
|
||||
msgEndpoint: msgEndpoint,
|
||||
incoming: make(chan []byte, 100),
|
||||
|
|
@ -485,12 +485,12 @@ func scanEvents(r io.Reader) iter.Seq2[event, error] {
|
|||
}
|
||||
}
|
||||
|
||||
// An sseClientStream is a logical jsonrpc2 stream that implements the client
|
||||
// An sseClientConn is a logical jsonrpc2 connection that implements the client
|
||||
// half of the SSE protocol:
|
||||
// - Writes are POSTS to the session endpoint.
|
||||
// - Reads are SSE 'message' events, and pushes them onto a buffered channel.
|
||||
// - Close terminates the GET request.
|
||||
type sseClientStream struct {
|
||||
type sseClientConn struct {
|
||||
sseEndpoint *url.URL // SSE endpoint for the GET
|
||||
msgEndpoint *url.URL // session endpoint for POSTs
|
||||
incoming chan []byte // queue of incoming messages
|
||||
|
|
@ -501,13 +501,13 @@ type sseClientStream struct {
|
|||
done chan struct{} // closed when the stream is closed
|
||||
}
|
||||
|
||||
func (c *sseClientStream) isDone() bool {
|
||||
func (c *sseClientConn) isDone() bool {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return c.closed
|
||||
}
|
||||
|
||||
func (c *sseClientStream) Read(ctx context.Context) (jsonrpc2.Message, error) {
|
||||
func (c *sseClientConn) Read(ctx context.Context) (jsonrpc2.Message, error) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
|
|
@ -528,7 +528,7 @@ func (c *sseClientStream) Read(ctx context.Context) (jsonrpc2.Message, error) {
|
|||
}
|
||||
}
|
||||
|
||||
func (c *sseClientStream) Write(ctx context.Context, msg jsonrpc2.Message) error {
|
||||
func (c *sseClientConn) Write(ctx context.Context, msg jsonrpc2.Message) error {
|
||||
data, err := jsonrpc2.EncodeMessage(msg)
|
||||
if err != nil {
|
||||
return err
|
||||
|
|
@ -552,7 +552,7 @@ func (c *sseClientStream) Write(ctx context.Context, msg jsonrpc2.Message) error
|
|||
return nil
|
||||
}
|
||||
|
||||
func (c *sseClientStream) Close() error {
|
||||
func (c *sseClientConn) Close() error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if !c.closed {
|
||||
|
|
|
|||
|
|
@ -150,20 +150,20 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque
|
|||
// 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{}),
|
||||
id: sessionID,
|
||||
incoming: make(chan jsonrpc2.Message, 10),
|
||||
done: make(chan struct{}),
|
||||
outgoingMessages: make(map[streamID][]*streamableMsg),
|
||||
signals: make(map[streamID]chan struct{}),
|
||||
requestStreams: make(map[JSONRPCID]streamID),
|
||||
streamRequests: make(map[streamID]map[JSONRPCID]struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// A StreamableServerTransport implements the [Transport] interface for a
|
||||
// single session.
|
||||
type StreamableServerTransport struct {
|
||||
nextConnID atomic.Int64 // incrementing next connection ID
|
||||
nextStreamID atomic.Int64 // incrementing next stream ID
|
||||
|
||||
id string
|
||||
incoming chan jsonrpc2.Message // messages from the client to the server
|
||||
|
|
@ -189,48 +189,48 @@ type StreamableServerTransport struct {
|
|||
// TODO(rfindley): simplify.
|
||||
|
||||
// outgoingMessages is the collection of outgoingMessages messages, keyed by the logical
|
||||
// connection ID where they should be delivered.
|
||||
// stream ID where they should be delivered.
|
||||
//
|
||||
// connID 0 is used for messages that don't correlate with an incoming
|
||||
// streamID 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
|
||||
outgoingMessages map[streamID][]*streamableMsg
|
||||
|
||||
// signals maps a logical connection ID to a 1-buffered channel, owned by an
|
||||
// signals maps a logical stream 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
|
||||
// response can receive messages for a logical stream. After claiming
|
||||
// the stream, 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{}
|
||||
// request for the given streamID.
|
||||
signals map[streamID]chan struct{}
|
||||
|
||||
// requestConns maps incoming requests to their logical connection ID.
|
||||
// requestStreams maps incoming requests to their logical stream ID.
|
||||
//
|
||||
// Lifecycle: requestConns persists for the duration of the session.
|
||||
// Lifecycle: requestStreams persists for the duration of the session.
|
||||
//
|
||||
// TODO(rfindley): clean up once requests are handled.
|
||||
requestConns map[JSONRPCID]connID
|
||||
requestStreams map[JSONRPCID]streamID
|
||||
|
||||
// outstandingRequests tracks the set of unanswered incoming RPCs for each logical
|
||||
// connection.
|
||||
// stream.
|
||||
//
|
||||
// When the server has responded to each request, the connection should be
|
||||
// When the server has responded to each request, the stream 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{}
|
||||
streamRequests map[streamID]map[JSONRPCID]struct{}
|
||||
}
|
||||
|
||||
type connID int64
|
||||
type streamID int64
|
||||
|
||||
// a streamableMsg is an SSE event with an index into its logical stream.
|
||||
type streamableMsg struct {
|
||||
|
|
@ -241,7 +241,7 @@ type streamableMsg struct {
|
|||
// Connect implements the [Transport] interface.
|
||||
//
|
||||
// TODO(rfindley): Connect should return a new object.
|
||||
func (s *StreamableServerTransport) Connect(context.Context) (Stream, error) {
|
||||
func (s *StreamableServerTransport) Connect(context.Context) (Connection, error) {
|
||||
return s, nil
|
||||
}
|
||||
|
||||
|
|
@ -284,7 +284,7 @@ func (t *StreamableServerTransport) ServeHTTP(w http.ResponseWriter, req *http.R
|
|||
|
||||
func (t *StreamableServerTransport) serveGET(w http.ResponseWriter, req *http.Request) {
|
||||
// connID 0 corresponds to the default GET request.
|
||||
id, nextIdx := connID(0), 0
|
||||
id, nextIdx := streamID(0), 0
|
||||
if len(req.Header.Values("Last-Event-ID")) > 0 {
|
||||
eid := req.Header.Get("Last-Event-ID")
|
||||
var ok bool
|
||||
|
|
@ -338,15 +338,15 @@ func (t *StreamableServerTransport) servePOST(w http.ResponseWriter, req *http.R
|
|||
}
|
||||
|
||||
// Update accounting for this request.
|
||||
id := connID(t.nextConnID.Add(1))
|
||||
id := streamID(t.nextStreamID.Add(1))
|
||||
signal := make(chan struct{}, 1)
|
||||
t.mu.Lock()
|
||||
if len(requests) > 0 {
|
||||
t.outstandingRequests[id] = make(map[JSONRPCID]struct{})
|
||||
t.streamRequests[id] = make(map[JSONRPCID]struct{})
|
||||
}
|
||||
for reqID := range requests {
|
||||
t.requestConns[reqID] = id
|
||||
t.outstandingRequests[id][reqID] = struct{}{}
|
||||
t.requestStreams[reqID] = id
|
||||
t.streamRequests[id][reqID] = struct{}{}
|
||||
}
|
||||
t.signals[id] = signal
|
||||
t.mu.Unlock()
|
||||
|
|
@ -362,7 +362,7 @@ func (t *StreamableServerTransport) servePOST(w http.ResponseWriter, req *http.R
|
|||
t.streamResponse(w, req, id, 0, signal)
|
||||
}
|
||||
|
||||
func (t *StreamableServerTransport) streamResponse(w http.ResponseWriter, req *http.Request, id connID, nextIndex int, signal chan struct{}) {
|
||||
func (t *StreamableServerTransport) streamResponse(w http.ResponseWriter, req *http.Request, id streamID, nextIndex int, signal chan struct{}) {
|
||||
defer func() {
|
||||
t.mu.Lock()
|
||||
delete(t.signals, id)
|
||||
|
|
@ -404,7 +404,7 @@ stream:
|
|||
}
|
||||
|
||||
t.mu.Lock()
|
||||
nOutstanding := len(t.outstandingRequests[id])
|
||||
nOutstanding := len(t.streamRequests[id])
|
||||
nOutgoing := len(t.outgoingMessages[id])
|
||||
t.mu.Unlock()
|
||||
// If all requests have been handled and replied to, we can terminate this
|
||||
|
|
@ -443,21 +443,21 @@ stream:
|
|||
}
|
||||
|
||||
// Event IDs: encode both the logical connection ID and the index, as
|
||||
// <connID>_<idx>, to be consistent with the typescript implementation.
|
||||
// <streamID>_<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.
|
||||
// streamID and message index idx.
|
||||
//
|
||||
// See also [parseEventID].
|
||||
func formatEventID(id connID, idx int) string {
|
||||
func formatEventID(id streamID, idx int) string {
|
||||
return fmt.Sprintf("%d_%d", id, idx)
|
||||
}
|
||||
|
||||
// parseEventID parses a Last-Event-ID value into a logical connection id and
|
||||
// parseEventID parses a Last-Event-ID value into a logical stream id and
|
||||
// index.
|
||||
//
|
||||
// See also [formatEventID].
|
||||
func parseEventID(eventID string) (conn connID, idx int, ok bool) {
|
||||
func parseEventID(eventID string) (conn streamID, idx int, ok bool) {
|
||||
_, err := fmt.Sscanf(eventID, "%d_%d", &conn, &idx)
|
||||
if err != nil || conn < 0 || idx < 0 {
|
||||
return 0, 0, false
|
||||
|
|
@ -465,7 +465,7 @@ func parseEventID(eventID string) (conn connID, idx int, ok bool) {
|
|||
return conn, idx, true
|
||||
}
|
||||
|
||||
// Read implements the [Stream] interface.
|
||||
// Read implements the [Connection] interface.
|
||||
func (t *StreamableServerTransport) Read(ctx context.Context) (jsonrpc2.Message, error) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
|
|
@ -480,7 +480,7 @@ func (t *StreamableServerTransport) Read(ctx context.Context) (jsonrpc2.Message,
|
|||
}
|
||||
}
|
||||
|
||||
// Write implements the [Stream] interface.
|
||||
// Write implements the [Connection] 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
|
||||
|
|
@ -501,10 +501,10 @@ func (t *StreamableServerTransport) Write(ctx context.Context, msg JSONRPCMessag
|
|||
//
|
||||
// For messages sent outside of a request context, this is the default
|
||||
// connection 0.
|
||||
var forConn connID
|
||||
var forConn streamID
|
||||
if forRequest.IsValid() {
|
||||
t.mu.Lock()
|
||||
forConn = t.requestConns[forRequest]
|
||||
forConn = t.requestStreams[forRequest]
|
||||
t.mu.Unlock()
|
||||
}
|
||||
|
||||
|
|
@ -519,7 +519,7 @@ func (t *StreamableServerTransport) Write(ctx context.Context, msg JSONRPCMessag
|
|||
return fmt.Errorf("session is closed") // TODO: should this be EOF?
|
||||
}
|
||||
|
||||
if _, ok := t.outstandingRequests[forConn]; !ok && forConn != 0 {
|
||||
if _, ok := t.streamRequests[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
|
||||
|
|
@ -538,9 +538,9 @@ func (t *StreamableServerTransport) Write(ctx context.Context, msg JSONRPCMessag
|
|||
})
|
||||
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)
|
||||
delete(t.streamRequests[forConn], replyTo)
|
||||
if len(t.streamRequests[forConn]) == 0 {
|
||||
delete(t.streamRequests, forConn)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -554,7 +554,7 @@ func (t *StreamableServerTransport) Write(ctx context.Context, msg JSONRPCMessag
|
|||
return nil
|
||||
}
|
||||
|
||||
// Close implements the [Stream] interface.
|
||||
// Close implements the [Connection] interface.
|
||||
func (t *StreamableServerTransport) Close() error {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
|
|
@ -582,14 +582,14 @@ func NewStreamableClientTransport(url string) *StreamableClientTransport {
|
|||
|
||||
// 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.
|
||||
// The resulting [Connection] 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
|
||||
// When closed, the connection issues a DELETE request to terminate the logical
|
||||
// session.
|
||||
func (t *StreamableClientTransport) Connect(ctx context.Context) (Stream, error) {
|
||||
return &streamableClientStream{
|
||||
func (t *StreamableClientTransport) Connect(ctx context.Context) (Connection, error) {
|
||||
return &streamableClientConn{
|
||||
url: t.url,
|
||||
client: http.DefaultClient,
|
||||
incoming: make(chan []byte, 100),
|
||||
|
|
@ -597,7 +597,7 @@ func (t *StreamableClientTransport) Connect(ctx context.Context) (Stream, error)
|
|||
}, nil
|
||||
}
|
||||
|
||||
type streamableClientStream struct {
|
||||
type streamableClientConn struct {
|
||||
url string
|
||||
sessionID string
|
||||
client *http.Client
|
||||
|
|
@ -612,8 +612,8 @@ type streamableClientStream struct {
|
|||
err error
|
||||
}
|
||||
|
||||
// Read implements the [Stream] interface.
|
||||
func (s *streamableClientStream) Read(ctx context.Context) (jsonrpc2.Message, error) {
|
||||
// Read implements the [Connection] interface.
|
||||
func (s *streamableClientConn) Read(ctx context.Context) (jsonrpc2.Message, error) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
|
|
@ -624,8 +624,8 @@ func (s *streamableClientStream) Read(ctx context.Context) (jsonrpc2.Message, er
|
|||
}
|
||||
}
|
||||
|
||||
// Write implements the [Stream] interface.
|
||||
func (s *streamableClientStream) Write(ctx context.Context, msg jsonrpc2.Message) error {
|
||||
// Write implements the [Connection] interface.
|
||||
func (s *streamableClientConn) Write(ctx context.Context, msg jsonrpc2.Message) error {
|
||||
s.mu.Lock()
|
||||
if s.err != nil {
|
||||
s.mu.Unlock()
|
||||
|
|
@ -661,7 +661,7 @@ func (s *streamableClientStream) Write(ctx context.Context, msg jsonrpc2.Message
|
|||
return nil
|
||||
}
|
||||
|
||||
func (s *streamableClientStream) postMessage(ctx context.Context, sessionID string, msg JSONRPCMessage) (string, error) {
|
||||
func (s *streamableClientConn) postMessage(ctx context.Context, sessionID string, msg JSONRPCMessage) (string, error) {
|
||||
data, err := jsonrpc2.EncodeMessage(msg)
|
||||
if err != nil {
|
||||
return "", err
|
||||
|
|
@ -697,7 +697,7 @@ func (s *streamableClientStream) postMessage(ctx context.Context, sessionID stri
|
|||
return sessionID, nil
|
||||
}
|
||||
|
||||
func (s *streamableClientStream) handleSSE(resp *http.Response) {
|
||||
func (s *streamableClientConn) handleSSE(resp *http.Response) {
|
||||
defer resp.Body.Close()
|
||||
|
||||
done := make(chan struct{})
|
||||
|
|
@ -718,8 +718,8 @@ func (s *streamableClientStream) handleSSE(resp *http.Response) {
|
|||
}
|
||||
}
|
||||
|
||||
// Close implements the [Stream] interface.
|
||||
func (s *streamableClientStream) Close() error {
|
||||
// Close implements the [Connection] interface.
|
||||
func (s *streamableClientConn) Close() error {
|
||||
s.closeOnce.Do(func() {
|
||||
close(s.done)
|
||||
|
||||
|
|
|
|||
|
|
@ -28,10 +28,10 @@ var ErrConnectionClosed = errors.New("connection closed")
|
|||
// Transports should be used for at most one call to [Server.Connect] or
|
||||
// [Client.Start].
|
||||
type Transport interface {
|
||||
// Connect returns the logical stream.
|
||||
// Connect returns the logical JSON-RPC connection..
|
||||
//
|
||||
// It is called exactly once by [Server.Connect] or [Client.Connect].
|
||||
Connect(ctx context.Context) (Stream, error)
|
||||
Connect(ctx context.Context) (Connection, error)
|
||||
}
|
||||
|
||||
type (
|
||||
|
|
@ -47,8 +47,8 @@ type (
|
|||
JSONRPCResponse = jsonrpc2.Response
|
||||
)
|
||||
|
||||
// A Stream is a bidirectional jsonrpc2 Stream.
|
||||
type Stream interface {
|
||||
// A Connection is a logical bidirectional JSON-RPC connection.
|
||||
type Connection interface {
|
||||
Read(context.Context) (JSONRPCMessage, error)
|
||||
Write(context.Context, JSONRPCMessage) error
|
||||
Close() error // may be called concurrently by both peers
|
||||
|
|
@ -66,8 +66,8 @@ type ioTransport struct {
|
|||
rwc io.ReadWriteCloser
|
||||
}
|
||||
|
||||
func (t *ioTransport) Connect(context.Context) (Stream, error) {
|
||||
return newIOStream(t.rwc), nil
|
||||
func (t *ioTransport) Connect(context.Context) (Connection, error) {
|
||||
return newIOConn(t.rwc), nil
|
||||
}
|
||||
|
||||
// NewStdIOTransport constructs a transport that communicates over
|
||||
|
|
@ -100,12 +100,12 @@ type handler interface {
|
|||
|
||||
func connect[H handler](ctx context.Context, t Transport, b binder[H]) (H, error) {
|
||||
var zero H
|
||||
stream, err := t.Connect(ctx)
|
||||
conn, err := t.Connect(ctx)
|
||||
if err != nil {
|
||||
return zero, err
|
||||
}
|
||||
// If logging is configured, write message logs.
|
||||
reader, writer := jsonrpc2.Reader(stream), jsonrpc2.Writer(stream)
|
||||
reader, writer := jsonrpc2.Reader(conn), jsonrpc2.Writer(conn)
|
||||
var (
|
||||
h H
|
||||
preempter canceller
|
||||
|
|
@ -118,7 +118,7 @@ func connect[H handler](ctx context.Context, t Transport, b binder[H]) (H, error
|
|||
_ = jsonrpc2.NewConnection(ctx, jsonrpc2.ConnectionConfig{
|
||||
Reader: reader,
|
||||
Writer: writer,
|
||||
Closer: stream,
|
||||
Closer: conn,
|
||||
Bind: bind,
|
||||
Preempter: &preempter,
|
||||
OnDone: func() {
|
||||
|
|
@ -187,23 +187,23 @@ func NewLoggingTransport(delegate Transport, w io.Writer) *LoggingTransport {
|
|||
return &LoggingTransport{delegate, w}
|
||||
}
|
||||
|
||||
// Connect connects the underlying transport, returning a [Stream] that writes
|
||||
// Connect connects the underlying transport, returning a [Connection] that writes
|
||||
// logs to the configured destination.
|
||||
func (t *LoggingTransport) Connect(ctx context.Context) (Stream, error) {
|
||||
func (t *LoggingTransport) Connect(ctx context.Context) (Connection, error) {
|
||||
delegate, err := t.delegate.Connect(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &loggingStream{delegate, t.w}, nil
|
||||
return &loggingConn{delegate, t.w}, nil
|
||||
}
|
||||
|
||||
type loggingStream struct {
|
||||
delegate Stream
|
||||
type loggingConn struct {
|
||||
delegate Connection
|
||||
w io.Writer
|
||||
}
|
||||
|
||||
// loggingReader is a stream middleware that logs incoming messages.
|
||||
func (s *loggingStream) Read(ctx context.Context) (jsonrpc2.Message, error) {
|
||||
func (s *loggingConn) Read(ctx context.Context) (jsonrpc2.Message, error) {
|
||||
msg, err := s.delegate.Read(ctx)
|
||||
if err != nil {
|
||||
fmt.Fprintf(s.w, "read error: %v", err)
|
||||
|
|
@ -218,7 +218,7 @@ func (s *loggingStream) Read(ctx context.Context) (jsonrpc2.Message, error) {
|
|||
}
|
||||
|
||||
// loggingWriter is a stream middleware that logs outgoing messages.
|
||||
func (s *loggingStream) Write(ctx context.Context, msg jsonrpc2.Message) error {
|
||||
func (s *loggingConn) Write(ctx context.Context, msg jsonrpc2.Message) error {
|
||||
err := s.delegate.Write(ctx, msg)
|
||||
if err != nil {
|
||||
fmt.Fprintf(s.w, "write error: %v", err)
|
||||
|
|
@ -232,7 +232,7 @@ func (s *loggingStream) Write(ctx context.Context, msg jsonrpc2.Message) error {
|
|||
return err
|
||||
}
|
||||
|
||||
func (s *loggingStream) Close() error {
|
||||
func (s *loggingConn) Close() error {
|
||||
return s.delegate.Close()
|
||||
}
|
||||
|
||||
|
|
@ -255,14 +255,14 @@ func (r rwc) Close() error {
|
|||
return errors.Join(r.rc.Close(), r.wc.Close())
|
||||
}
|
||||
|
||||
// An ioStream is a transport that delimits messages with newlines across
|
||||
// An ioConn is a transport that delimits messages with newlines across
|
||||
// a bidirectional stream, and supports JSONRPC2 message batching.
|
||||
//
|
||||
// See https://github.com/ndjson/ndjson-spec for discussion of newline
|
||||
// delimited JSON.
|
||||
//
|
||||
// See [msgBatch] for more discussion of message batching.
|
||||
type ioStream struct {
|
||||
type ioConn struct {
|
||||
rwc io.ReadWriteCloser // the underlying stream
|
||||
in *json.Decoder // a decoder bound to rwc
|
||||
|
||||
|
|
@ -280,23 +280,18 @@ type ioStream struct {
|
|||
batches map[jsonrpc2.ID]*msgBatch // lazily allocated
|
||||
}
|
||||
|
||||
func newIOStream(rwc io.ReadWriteCloser) *ioStream {
|
||||
return &ioStream{
|
||||
func newIOConn(rwc io.ReadWriteCloser) *ioConn {
|
||||
return &ioConn{
|
||||
rwc: rwc,
|
||||
in: json.NewDecoder(rwc),
|
||||
}
|
||||
}
|
||||
|
||||
// Connect returns the receiver, as a streamTransport is a logical stream.
|
||||
func (t *ioStream) Connect(ctx context.Context) (Stream, error) {
|
||||
return t, nil
|
||||
}
|
||||
|
||||
// addBatch records a msgBatch for an incoming batch payload.
|
||||
// It returns an error if batch is malformed, containing previously seen IDs.
|
||||
//
|
||||
// See [msgBatch] for more.
|
||||
func (t *ioStream) addBatch(batch *msgBatch) error {
|
||||
func (t *ioConn) addBatch(batch *msgBatch) error {
|
||||
t.batchMu.Lock()
|
||||
defer t.batchMu.Unlock()
|
||||
for id := range batch.unresolved {
|
||||
|
|
@ -319,7 +314,7 @@ func (t *ioStream) addBatch(batch *msgBatch) error {
|
|||
// The second result reports whether resp was part of a batch. If this is true,
|
||||
// the first result is nil if the batch is still incomplete, or the full set of
|
||||
// batch responses if resp completed the batch.
|
||||
func (t *ioStream) updateBatch(resp *jsonrpc2.Response) ([]*jsonrpc2.Response, bool) {
|
||||
func (t *ioConn) updateBatch(resp *jsonrpc2.Response) ([]*jsonrpc2.Response, bool) {
|
||||
t.batchMu.Lock()
|
||||
defer t.batchMu.Unlock()
|
||||
|
||||
|
|
@ -357,11 +352,11 @@ type msgBatch struct {
|
|||
responses []*jsonrpc2.Response
|
||||
}
|
||||
|
||||
func (t *ioStream) Read(ctx context.Context) (jsonrpc2.Message, error) {
|
||||
func (t *ioConn) Read(ctx context.Context) (jsonrpc2.Message, error) {
|
||||
return t.read(ctx, t.in)
|
||||
}
|
||||
|
||||
func (t *ioStream) read(ctx context.Context, in *json.Decoder) (jsonrpc2.Message, error) {
|
||||
func (t *ioConn) read(ctx context.Context, in *json.Decoder) (jsonrpc2.Message, error) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
|
|
@ -432,7 +427,7 @@ func readBatch(data []byte) (msgs []jsonrpc2.Message, isBatch bool, _ error) {
|
|||
return []jsonrpc2.Message{msg}, false, err
|
||||
}
|
||||
|
||||
func (t *ioStream) Write(ctx context.Context, msg jsonrpc2.Message) error {
|
||||
func (t *ioConn) Write(ctx context.Context, msg jsonrpc2.Message) error {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
|
|
@ -479,7 +474,7 @@ func (t *ioStream) Write(ctx context.Context, msg jsonrpc2.Message) error {
|
|||
return err
|
||||
}
|
||||
|
||||
func (t *ioStream) Close() error {
|
||||
func (t *ioConn) Close() error {
|
||||
return t.rwc.Close()
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -12,16 +12,6 @@ import (
|
|||
"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
|
||||
)
|
||||
|
||||
// BatchSize causes a transport to collect n requests or notifications before
|
||||
// sending a batch on the wire (responses are always sent in isolation).
|
||||
//
|
||||
// Exported for testing in the mcp_test package.
|
||||
func BatchSize(t Transport, n int) {
|
||||
if st, ok := t.(*ioStream); ok {
|
||||
st.outgoingBatch = make([]jsonrpc2.Message, 0, n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchFraming(t *testing.T) {
|
||||
// This test checks that the ndjsonFramer can read and write JSON batches.
|
||||
//
|
||||
|
|
@ -31,7 +21,7 @@ func TestBatchFraming(t *testing.T) {
|
|||
ctx := context.Background()
|
||||
|
||||
r, w := io.Pipe()
|
||||
tport := newIOStream(rwc{r, w})
|
||||
tport := newIOConn(rwc{r, w})
|
||||
tport.outgoingBatch = make([]jsonrpc2.Message, 0, 2)
|
||||
|
||||
// Read the two messages into a channel, for easy testing later.
|
||||
|
|
|
|||
Loading…
Reference in a new issue