mcp: implement a concurrency model for calls

Implement the concurrency model described in #26: notifications are
synchronous, but calls are asynchronous (except for 'initialize'). To
achieve this, implement jsonrpc2.Async(ctx) to signal asynchronous
handling. This is simpler to use than returning ErrAsyncResponse and
calling Respond, and since this is an internal detail we don't need to
worry too much about whether it's idiomatic.

Add tests that verify both features, for both client and server.

Also:
- replace req.ID.IsValid with req.IsCall
- remove the methodHandler type as we can just use MethodHandler

Fixes #26
This commit is contained in:
Rob Findley 2025-08-18 18:15:06 +00:00 • committed by Robert Findley
commit 3db848abec

View file

@ -38,12 +38,6 @@ var supportedProtocolVersions = []string{
// For notifications, both must be nil.
type MethodHandler func(ctx context.Context, method string, req Request) (result Result, err error)
// A methodHandler is a MethodHandler[Session] for some session.
// We need to give up type safety here, or we will end up with a type cycle somewhere
// else. For example, if Session.methodHandler returned a MethodHandler[Session],
// the compiler would complain.
type methodHandler any // MethodHandler[*ClientSession] | MethodHandler[*ServerSession]
// A Session is either a [ClientSession] or a [ServerSession].
type Session interface {
// ID returns the session ID, or the empty string if there is none.
@ -51,8 +45,8 @@ type Session interface {
sendingMethodInfos() map[string]methodInfo
receivingMethodInfos() map[string]methodInfo
sendingMethodHandler() methodHandler
receivingMethodHandler() methodHandler
sendingMethodHandler() MethodHandler
receivingMethodHandler() MethodHandler
getConn() *jsonrpc2.Connection
}
@ -95,13 +89,13 @@ func orZero[T any, P *U, U any](p P) T {
}
func handleNotify(ctx context.Context, method string, req Request) error {
mh := req.GetSession().sendingMethodHandler().(MethodHandler)
mh := req.GetSession().sendingMethodHandler()
_, err := mh(ctx, method, req)
return err
}
func handleSend[R Result](ctx context.Context, method string, req Request) (R, error) {
mh := req.GetSession().sendingMethodHandler().(MethodHandler)
mh := req.GetSession().sendingMethodHandler()
// mh might be user code, so ensure that it returns the right values for the jsonrpc2 protocol.
res, err := mh(ctx, method, req)
if err != nil {
@ -118,7 +112,7 @@ func defaultReceivingMethodHandler[S Session](ctx context.Context, method string
// This can be called from user code, with an arbitrary value for method.
return nil, jsonrpc2.ErrNotHandled
}
return info.handleMethod.(MethodHandler)(ctx, method, req)
return info.handleMethod(ctx, method, req)
}
func handleReceive[S Session](ctx context.Context, session S, jreq *jsonrpc.Request) (Result, error) {
@ -131,7 +125,7 @@ func handleReceive[S Session](ctx context.Context, session S, jreq *jsonrpc.Requ
return nil, fmt.Errorf("handling '%s': %w", jreq.Method, err)
}
mh := session.receivingMethodHandler().(MethodHandler)
mh := session.receivingMethodHandler()
req := info.newRequest(session, params)
// mh might be user code, so ensure that it returns the right values for the jsonrpc2 protocol.
res, err := mh(ctx, jreq.Method, req)
@ -154,10 +148,10 @@ func checkRequest(req *jsonrpc.Request, infos map[string]methodInfo) (methodInfo
if !ok {
return methodInfo{}, fmt.Errorf("%w: %q unsupported", jsonrpc2.ErrNotHandled, req.Method)
}
if info.flags&notification != 0 && req.ID.IsValid() {
if info.flags&notification != 0 && req.IsCall() {
return methodInfo{}, fmt.Errorf("%w: unexpected id for %q", jsonrpc2.ErrInvalidRequest, req.Method)
}
if info.flags&notification == 0 && !req.ID.IsValid() {
if info.flags&notification == 0 && !req.IsCall() {
return methodInfo{}, fmt.Errorf("%w: missing id for %q", jsonrpc2.ErrInvalidRequest, req.Method)
}
// missingParamsOK is checked here to catch the common case where "params" is
@ -182,7 +176,7 @@ type methodInfo struct {
newRequest func(Session, Params) Request
// Run the code when a call to the method is received.
// Used on the receive side.
handleMethod methodHandler
handleMethod MethodHandler
// Create a pointer to a Result struct.
// Used on the send side.
newResult func() Result