diff --git a/mcp/client.go b/mcp/client.go index aa72545..17875c3 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -251,7 +251,7 @@ func (cs *ClientSession) receivingMethodInfos() map[string]methodInfo { return clientMethodInfos } -func (cs *ClientSession) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) { +func (cs *ClientSession) handle(ctx context.Context, req *JSONRPCRequest) (any, error) { return handleReceive(ctx, cs, req) } diff --git a/mcp/conformance_test.go b/mcp/conformance_test.go index b793e3e..3be54a3 100644 --- a/mcp/conformance_test.go +++ b/mcp/conformance_test.go @@ -46,12 +46,12 @@ var update = flag.Bool("update", false, "if set, update conformance test data") // with -update to have the test runner update the expected output, which may // be client or server depending on the perspective of the test. type conformanceTest struct { - name string // test name - path string // path to test file - archive *txtar.Archive // raw archive, for updating - tools, prompts, resources []string // named features to include - client []jsonrpc2.Message // client messages - server []jsonrpc2.Message // server messages + name string // test name + path string // path to test file + archive *txtar.Archive // raw archive, for updating + tools, prompts, resources []string // named features to include + client []JSONRPCMessage // client messages + server []JSONRPCMessage // server messages } // TODO(rfindley): add client conformance tests. @@ -117,24 +117,24 @@ func runServerTest(t *testing.T, test *conformanceTest) { t.Fatal(err) } - writeMsg := func(msg jsonrpc2.Message) { + writeMsg := func(msg JSONRPCMessage) { if err := cStream.Write(ctx, msg); err != nil { t.Fatalf("Write failed: %v", err) } } var ( - serverMessages []jsonrpc2.Message - outRequests []*jsonrpc2.Request - outResponses []*jsonrpc2.Response + serverMessages []JSONRPCMessage + outRequests []*JSONRPCRequest + outResponses []*JSONRPCResponse ) // Separate client requests and responses; we use them differently. for _, msg := range test.client { switch msg := msg.(type) { - case *jsonrpc2.Request: + case *JSONRPCRequest: outRequests = append(outRequests, msg) - case *jsonrpc2.Response: + case *JSONRPCResponse: outResponses = append(outResponses, msg) default: t.Fatalf("bad message type %T", msg) @@ -143,7 +143,7 @@ func runServerTest(t *testing.T, test *conformanceTest) { // nextResponse handles incoming requests and notifications, and returns the // next incoming response. - nextResponse := func() (*jsonrpc2.Response, error, bool) { + nextResponse := func() (*JSONRPCResponse, error, bool) { for { msg, err := cStream.Read(ctx) if err != nil { @@ -156,7 +156,7 @@ func runServerTest(t *testing.T, test *conformanceTest) { return nil, err, false } serverMessages = append(serverMessages, msg) - if req, ok := msg.(*jsonrpc2.Request); ok && req.ID.IsValid() { + if req, ok := msg.(*JSONRPCRequest); ok && req.ID.IsValid() { // Pair up the next outgoing response with this request. // We assume requests arrive in the same order every time. if len(outResponses) == 0 { @@ -167,7 +167,7 @@ func runServerTest(t *testing.T, test *conformanceTest) { outResponses = outResponses[1:] continue } - return msg.(*jsonrpc2.Response), nil, true + return msg.(*JSONRPCResponse), nil, true } } @@ -191,7 +191,7 @@ func runServerTest(t *testing.T, test *conformanceTest) { // There might be more notifications or requests, but there shouldn't be more // responses. // Run this in a goroutine so the current thread can wait for it. - var extra *jsonrpc2.Response + var extra *JSONRPCResponse go func() { extra, err, _ = nextResponse() }() @@ -240,8 +240,8 @@ func runServerTest(t *testing.T, test *conformanceTest) { t.Fatalf("os.WriteFile(%q) failed: %v", test.path, err) } } else { - // jsonrpc2.Messages are not comparable, so we instead compare lines of JSON. - transform := cmpopts.AcyclicTransformer("toJSON", func(msg jsonrpc2.Message) []string { + // JSONRPCMessages are not comparable, so we instead compare lines of JSON. + transform := cmpopts.AcyclicTransformer("toJSON", func(msg JSONRPCMessage) []string { encoded, err := jsonrpc2.EncodeIndent(msg, "", "\t") if err != nil { t.Fatal(err) @@ -271,9 +271,9 @@ func loadConformanceTest(dir, path string) (*conformanceTest, error) { } // decodeMessages loads JSON-RPC messages from the archive file. - decodeMessages := func(data []byte) ([]jsonrpc2.Message, error) { + decodeMessages := func(data []byte) ([]JSONRPCMessage, error) { dec := json.NewDecoder(bytes.NewReader(data)) - var res []jsonrpc2.Message + var res []JSONRPCMessage for dec.More() { var raw json.RawMessage if err := dec.Decode(&raw); err != nil { diff --git a/mcp/server.go b/mcp/server.go index 3df5b68..87007eb 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -597,7 +597,7 @@ func (ss *ServerSession) receivingMethodHandler() methodHandler { func (ss *ServerSession) getConn() *jsonrpc2.Connection { return ss.conn } // handle invokes the method described by the given JSON RPC request. -func (ss *ServerSession) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) { +func (ss *ServerSession) handle(ctx context.Context, req *JSONRPCRequest) (any, error) { ss.mu.Lock() initialized := ss.initialized ss.mu.Unlock() diff --git a/mcp/shared.go b/mcp/shared.go index 14939a5..942e0d6 100644 --- a/mcp/shared.go +++ b/mcp/shared.go @@ -108,7 +108,7 @@ func defaultReceivingMethodHandler[S Session](ctx context.Context, session S, me return info.handleMethod.(MethodHandler[S])(ctx, session, method, params) } -func handleReceive[S Session](ctx context.Context, session S, req *jsonrpc2.Request) (Result, error) { +func handleReceive[S Session](ctx context.Context, session S, req *JSONRPCRequest) (Result, error) { info, ok := session.receivingMethodInfos()[req.Method] if !ok { return nil, jsonrpc2.ErrNotHandled diff --git a/mcp/sse.go b/mcp/sse.go index 187b14e..2ee23d0 100644 --- a/mcp/sse.go +++ b/mcp/sse.go @@ -111,7 +111,7 @@ func NewSSEHandler(getServer func(request *http.Request) *Server) *SSEHandler { // - Close terminates the hanging GET. type SSEServerTransport struct { endpoint string - incoming chan jsonrpc2.Message // queue of incoming messages; never closed + incoming chan JSONRPCMessage // queue of incoming messages; never closed // We must guard both pushes to the incoming queue and writes to the response // writer, because incoming POST requests are arbitrarily concurrent and we @@ -138,7 +138,7 @@ func NewSSEServerTransport(endpoint string, w http.ResponseWriter) *SSEServerTra return &SSEServerTransport{ endpoint: endpoint, w: w, - incoming: make(chan jsonrpc2.Message, 100), + incoming: make(chan JSONRPCMessage, 100), done: make(chan struct{}), } } @@ -264,7 +264,7 @@ type sseServerConn struct { } // Read implements jsonrpc2.Reader. -func (s sseServerConn) Read(ctx context.Context) (jsonrpc2.Message, error) { +func (s sseServerConn) Read(ctx context.Context) (JSONRPCMessage, error) { select { case <-ctx.Done(): return nil, ctx.Err() @@ -276,7 +276,7 @@ func (s sseServerConn) Read(ctx context.Context) (jsonrpc2.Message, error) { } // Write implements jsonrpc2.Writer. -func (s sseServerConn) Write(ctx context.Context, msg jsonrpc2.Message) error { +func (s sseServerConn) Write(ctx context.Context, msg JSONRPCMessage) error { if ctx.Err() != nil { return ctx.Err() } @@ -507,7 +507,7 @@ func (c *sseClientConn) isDone() bool { return c.closed } -func (c *sseClientConn) Read(ctx context.Context) (jsonrpc2.Message, error) { +func (c *sseClientConn) Read(ctx context.Context) (JSONRPCMessage, error) { select { case <-ctx.Done(): return nil, ctx.Err() @@ -528,7 +528,7 @@ func (c *sseClientConn) Read(ctx context.Context) (jsonrpc2.Message, error) { } } -func (c *sseClientConn) Write(ctx context.Context, msg jsonrpc2.Message) error { +func (c *sseClientConn) Write(ctx context.Context, msg JSONRPCMessage) error { data, err := jsonrpc2.EncodeMessage(msg) if err != nil { return err diff --git a/mcp/streamable.go b/mcp/streamable.go index c70d6d7..3b860b9 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -151,7 +151,7 @@ func (h *StreamableHTTPHandler) ServeHTTP(w http.ResponseWriter, req *http.Reque func NewStreamableServerTransport(sessionID string) *StreamableServerTransport { return &StreamableServerTransport{ id: sessionID, - incoming: make(chan jsonrpc2.Message, 10), + incoming: make(chan JSONRPCMessage, 10), done: make(chan struct{}), outgoingMessages: make(map[streamID][]*streamableMsg), signals: make(map[streamID]chan struct{}), @@ -166,7 +166,7 @@ type StreamableServerTransport struct { nextStreamID atomic.Int64 // incrementing next stream ID id string - incoming chan jsonrpc2.Message // messages from the client to the server + incoming chan JSONRPCMessage // messages from the client to the server mu sync.Mutex @@ -466,7 +466,7 @@ func parseEventID(eventID string) (conn streamID, idx int, ok bool) { } // Read implements the [Connection] interface. -func (t *StreamableServerTransport) Read(ctx context.Context) (jsonrpc2.Message, error) { +func (t *StreamableServerTransport) Read(ctx context.Context) (JSONRPCMessage, error) { select { case <-ctx.Done(): return nil, ctx.Err() @@ -613,7 +613,7 @@ type streamableClientConn struct { } // Read implements the [Connection] interface. -func (s *streamableClientConn) Read(ctx context.Context) (jsonrpc2.Message, error) { +func (s *streamableClientConn) Read(ctx context.Context) (JSONRPCMessage, error) { select { case <-ctx.Done(): return nil, ctx.Err() @@ -625,7 +625,7 @@ func (s *streamableClientConn) Read(ctx context.Context) (jsonrpc2.Message, erro } // Write implements the [Connection] interface. -func (s *streamableClientConn) Write(ctx context.Context, msg jsonrpc2.Message) error { +func (s *streamableClientConn) Write(ctx context.Context, msg JSONRPCMessage) error { s.mu.Lock() if s.err != nil { s.mu.Unlock() diff --git a/mcp/transport.go b/mcp/transport.go index 008fc3d..dafd1e7 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -37,8 +37,6 @@ type Transport interface { 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. @@ -95,7 +93,7 @@ type binder[T handler] interface { } type handler interface { - handle(ctx context.Context, req *jsonrpc2.Request) (any, error) + handle(ctx context.Context, req *JSONRPCRequest) (any, error) } func connect[H handler](ctx context.Context, t Transport, b binder[H]) (H, error) { @@ -136,7 +134,7 @@ type canceller struct { } // Preempt implements jsonrpc2.Preempter. -func (c *canceller) Preempt(ctx context.Context, req *jsonrpc2.Request) (result any, err error) { +func (c *canceller) Preempt(ctx context.Context, req *JSONRPCRequest) (result any, err error) { if req.Method == "notifications/cancelled" { var params CancelledParams if err := json.Unmarshal(req.Params, ¶ms); err != nil { @@ -203,7 +201,7 @@ type loggingConn struct { } // loggingReader is a stream middleware that logs incoming messages. -func (s *loggingConn) Read(ctx context.Context) (jsonrpc2.Message, error) { +func (s *loggingConn) Read(ctx context.Context) (JSONRPCMessage, error) { msg, err := s.delegate.Read(ctx) if err != nil { fmt.Fprintf(s.w, "read error: %v", err) @@ -218,7 +216,7 @@ func (s *loggingConn) Read(ctx context.Context) (jsonrpc2.Message, error) { } // loggingWriter is a stream middleware that logs outgoing messages. -func (s *loggingConn) Write(ctx context.Context, msg jsonrpc2.Message) error { +func (s *loggingConn) Write(ctx context.Context, msg JSONRPCMessage) error { err := s.delegate.Write(ctx, msg) if err != nil { fmt.Fprintf(s.w, "write error: %v", err) @@ -268,11 +266,11 @@ type ioConn struct { // If outgoiBatch has a positive capacity, it will be used to batch requests // and notifications before sending. - outgoingBatch []jsonrpc2.Message + outgoingBatch []JSONRPCMessage // Unread messages in the last batch. Since reads are serialized, there is no // need to guard here. - queue []jsonrpc2.Message + queue []JSONRPCMessage // batches correlate incoming requests to the batch in which they arrived. // Since writes may be concurrent to reads, we need to guard this with a mutex. @@ -314,7 +312,7 @@ func (t *ioConn) 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 *ioConn) updateBatch(resp *jsonrpc2.Response) ([]*jsonrpc2.Response, bool) { +func (t *ioConn) updateBatch(resp *JSONRPCResponse) ([]*JSONRPCResponse, bool) { t.batchMu.Lock() defer t.batchMu.Unlock() @@ -349,14 +347,14 @@ func (t *ioConn) updateBatch(resp *jsonrpc2.Response) ([]*jsonrpc2.Response, boo // When there are no unresolved calls, the response payload is sent. type msgBatch struct { unresolved map[jsonrpc2.ID]int - responses []*jsonrpc2.Response + responses []*JSONRPCResponse } -func (t *ioConn) Read(ctx context.Context) (jsonrpc2.Message, error) { +func (t *ioConn) Read(ctx context.Context) (JSONRPCMessage, error) { return t.read(ctx, t.in) } -func (t *ioConn) read(ctx context.Context, in *json.Decoder) (jsonrpc2.Message, error) { +func (t *ioConn) read(ctx context.Context, in *json.Decoder) (JSONRPCMessage, error) { select { case <-ctx.Done(): return nil, ctx.Err() @@ -381,7 +379,7 @@ func (t *ioConn) read(ctx context.Context, in *json.Decoder) (jsonrpc2.Message, if batch { var respBatch *msgBatch // track incoming requests in the batch for _, msg := range msgs { - if req, ok := msg.(*jsonrpc2.Request); ok { + if req, ok := msg.(*JSONRPCRequest); ok { if respBatch == nil { respBatch = &msgBatch{ unresolved: make(map[jsonrpc2.ID]int), @@ -406,7 +404,7 @@ func (t *ioConn) read(ctx context.Context, in *json.Decoder) (jsonrpc2.Message, // 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) { +func readBatch(data []byte) (msgs []JSONRPCMessage, isBatch bool, _ error) { // Try to read an array of messages first. var rawBatch []json.RawMessage if err := json.Unmarshal(data, &rawBatch); err == nil { @@ -424,10 +422,10 @@ func readBatch(data []byte) (msgs []jsonrpc2.Message, isBatch bool, _ error) { } // Try again with a single message. msg, err := jsonrpc2.DecodeMessage(data) - return []jsonrpc2.Message{msg}, false, err + return []JSONRPCMessage{msg}, false, err } -func (t *ioConn) Write(ctx context.Context, msg jsonrpc2.Message) error { +func (t *ioConn) Write(ctx context.Context, msg JSONRPCMessage) error { select { case <-ctx.Done(): return ctx.Err() @@ -438,7 +436,7 @@ func (t *ioConn) Write(ctx context.Context, msg jsonrpc2.Message) error { // check that first. Otherwise, it is a request or notification, and we may // want to collect it into a batch before sending, if we're configured to use // outgoing batches. - if resp, ok := msg.(*jsonrpc2.Response); ok { + if resp, ok := msg.(*JSONRPCResponse); ok { if batch, ok := t.updateBatch(resp); ok { if len(batch) > 0 { data, err := marshalMessages(batch) @@ -478,7 +476,7 @@ func (t *ioConn) Close() error { return t.rwc.Close() } -func marshalMessages[T jsonrpc2.Message](msgs []T) ([]byte, error) { +func marshalMessages[T JSONRPCMessage](msgs []T) ([]byte, error) { var rawMsgs []json.RawMessage for _, msg := range msgs { raw, err := jsonrpc2.EncodeMessage(msg) diff --git a/mcp/transport_test.go b/mcp/transport_test.go index 83c6d7a..db18a35 100644 --- a/mcp/transport_test.go +++ b/mcp/transport_test.go @@ -22,10 +22,10 @@ func TestBatchFraming(t *testing.T) { r, w := io.Pipe() tport := newIOConn(rwc{r, w}) - tport.outgoingBatch = make([]jsonrpc2.Message, 0, 2) + tport.outgoingBatch = make([]JSONRPCMessage, 0, 2) // Read the two messages into a channel, for easy testing later. - read := make(chan jsonrpc2.Message) + read := make(chan JSONRPCMessage) go func() { for range 2 { msg, _ := tport.Read(ctx) @@ -34,7 +34,7 @@ func TestBatchFraming(t *testing.T) { }() // The first write should not yet be observed by the reader. - tport.Write(ctx, &jsonrpc2.Request{ID: jsonrpc2.Int64ID(1), Method: "test"}) + tport.Write(ctx, &JSONRPCRequest{ID: jsonrpc2.Int64ID(1), Method: "test"}) select { case got := <-read: t.Fatalf("after one write, got message %v", got) @@ -42,10 +42,10 @@ func TestBatchFraming(t *testing.T) { } // ...but the second write causes both messages to be observed. - tport.Write(ctx, &jsonrpc2.Request{ID: jsonrpc2.Int64ID(2), Method: "test"}) + tport.Write(ctx, &JSONRPCRequest{ID: jsonrpc2.Int64ID(2), Method: "test"}) for _, want := range []int64{1, 2} { got := <-read - if got := got.(*jsonrpc2.Request).ID.Raw(); got != want { + if got := got.(*JSONRPCRequest).ID.Raw(); got != want { t.Errorf("got message #%d, want #%d", got, want) } }