mcp: introduce Requests (#267)
Handlers and client/server methods take a single request argument combining the session and params. This slightly simplifies the signature for handlers, since there is no longer an often-ignored session argument. But more importantly, it opens the door to adding more information in requests, such as auth info and HTTP request headers. For #243.
This commit is contained in:
Jonathan Amsterdam
GitHub
parent
e097918484
commit
6e03217c83
1 changed files with 517 additions and 380 deletions
185
mcp/shared.go
185
mcp/shared.go
|
|
@ -36,7 +36,7 @@ var supportedProtocolVersions = []string{
|
|||
// A MethodHandler handles MCP messages.
|
||||
// For methods, exactly one of the return values must be nil.
|
||||
// For notifications, both must be nil.
|
||||
type MethodHandler[S Session] func(ctx context.Context, _ S, method string, params Params) (result Result, err error)
|
||||
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
|
||||
|
|
@ -46,7 +46,6 @@ type methodHandler any // MethodHandler[*ClientSession] | MethodHandler[*ServerS
|
|||
|
||||
// A Session is either a [ClientSession] or a [ServerSession].
|
||||
type Session interface {
|
||||
*ClientSession | *ServerSession
|
||||
// ID returns the session ID, or the empty string if there is none.
|
||||
ID() string
|
||||
|
||||
|
|
@ -58,29 +57,29 @@ type Session interface {
|
|||
}
|
||||
|
||||
// Middleware is a function from [MethodHandler] to [MethodHandler].
|
||||
type Middleware[S Session] func(MethodHandler[S]) MethodHandler[S]
|
||||
type Middleware func(MethodHandler) MethodHandler
|
||||
|
||||
// addMiddleware wraps the handler in the middleware functions.
|
||||
func addMiddleware[S Session](handlerp *MethodHandler[S], middleware []Middleware[S]) {
|
||||
func addMiddleware(handlerp *MethodHandler, middleware []Middleware) {
|
||||
for _, m := range slices.Backward(middleware) {
|
||||
*handlerp = m(*handlerp)
|
||||
}
|
||||
}
|
||||
|
||||
func defaultSendingMethodHandler[S Session](ctx context.Context, session S, method string, params Params) (Result, error) {
|
||||
info, ok := session.sendingMethodInfos()[method]
|
||||
func defaultSendingMethodHandler[S Session](ctx context.Context, method string, req Request) (Result, error) {
|
||||
info, ok := req.GetSession().sendingMethodInfos()[method]
|
||||
if !ok {
|
||||
// This can be called from user code, with an arbitrary value for method.
|
||||
return nil, jsonrpc2.ErrNotHandled
|
||||
}
|
||||
// Notifications don't have results.
|
||||
if strings.HasPrefix(method, "notifications/") {
|
||||
return nil, session.getConn().Notify(ctx, method, params)
|
||||
return nil, req.GetSession().getConn().Notify(ctx, method, req.GetParams())
|
||||
}
|
||||
// Create the result to unmarshal into.
|
||||
// The concrete type of the result is the return type of the receiving function.
|
||||
res := info.newResult()
|
||||
if err := call(ctx, session.getConn(), method, params, res); err != nil {
|
||||
if err := call(ctx, req.GetSession().getConn(), method, req.GetParams(), res); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return res, nil
|
||||
|
|
@ -95,16 +94,16 @@ func orZero[T any, P *U, U any](p P) T {
|
|||
return any(p).(T)
|
||||
}
|
||||
|
||||
func handleNotify[S Session](ctx context.Context, session S, method string, params Params) error {
|
||||
mh := session.sendingMethodHandler().(MethodHandler[S])
|
||||
_, err := mh(ctx, session, method, params)
|
||||
func handleNotify(ctx context.Context, method string, req Request) error {
|
||||
mh := req.GetSession().sendingMethodHandler().(MethodHandler)
|
||||
_, err := mh(ctx, method, req)
|
||||
return err
|
||||
}
|
||||
|
||||
func handleSend[R Result, S Session](ctx context.Context, s S, method string, params Params) (R, error) {
|
||||
mh := s.sendingMethodHandler().(MethodHandler[S])
|
||||
func handleSend[R Result](ctx context.Context, method string, req Request) (R, error) {
|
||||
mh := req.GetSession().sendingMethodHandler().(MethodHandler)
|
||||
// mh might be user code, so ensure that it returns the right values for the jsonrpc2 protocol.
|
||||
res, err := mh(ctx, s, method, params)
|
||||
res, err := mh(ctx, method, req)
|
||||
if err != nil {
|
||||
var z R
|
||||
return z, err
|
||||
|
|
@ -113,28 +112,29 @@ func handleSend[R Result, S Session](ctx context.Context, s S, method string, pa
|
|||
}
|
||||
|
||||
// defaultReceivingMethodHandler is the initial MethodHandler for servers and clients, before being wrapped by middleware.
|
||||
func defaultReceivingMethodHandler[S Session](ctx context.Context, session S, method string, params Params) (Result, error) {
|
||||
info, ok := session.receivingMethodInfos()[method]
|
||||
func defaultReceivingMethodHandler[S Session](ctx context.Context, method string, req Request) (Result, error) {
|
||||
info, ok := req.GetSession().receivingMethodInfos()[method]
|
||||
if !ok {
|
||||
// This can be called from user code, with an arbitrary value for method.
|
||||
return nil, jsonrpc2.ErrNotHandled
|
||||
}
|
||||
return info.handleMethod.(MethodHandler[S])(ctx, session, method, params)
|
||||
return info.handleMethod.(MethodHandler)(ctx, method, req)
|
||||
}
|
||||
|
||||
func handleReceive[S Session](ctx context.Context, session S, req *jsonrpc.Request) (Result, error) {
|
||||
info, err := checkRequest(req, session.receivingMethodInfos())
|
||||
func handleReceive[S Session](ctx context.Context, session S, jreq *jsonrpc.Request) (Result, error) {
|
||||
info, err := checkRequest(jreq, session.receivingMethodInfos())
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
params, err := info.unmarshalParams(req.Params)
|
||||
params, err := info.unmarshalParams(jreq.Params)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("handling '%s': %w", req.Method, err)
|
||||
return nil, fmt.Errorf("handling '%s': %w", jreq.Method, err)
|
||||
}
|
||||
|
||||
mh := session.receivingMethodHandler().(MethodHandler[S])
|
||||
mh := session.receivingMethodHandler().(MethodHandler)
|
||||
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, session, req.Method, params)
|
||||
res, err := mh(ctx, jreq.Method, req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
|
@ -179,6 +179,7 @@ type methodInfo struct {
|
|||
// Unmarshal params from the wire into a Params struct.
|
||||
// Used on the receive side.
|
||||
unmarshalParams func(json.RawMessage) (Params, error)
|
||||
newRequest func(Session, Params) Request
|
||||
// Run the code when a call to the method is received.
|
||||
// Used on the receive side.
|
||||
handleMethod methodHandler
|
||||
|
|
@ -194,7 +195,10 @@ type methodInfo struct {
|
|||
// - R: results
|
||||
|
||||
// A typedMethodHandler is like a MethodHandler, but with type information.
|
||||
type typedMethodHandler[S Session, P Params, R Result] func(context.Context, S, P) (R, error)
|
||||
type (
|
||||
typedClientMethodHandler[P Params, R Result] func(context.Context, *ClientRequest[P]) (R, error)
|
||||
typedServerMethodHandler[P Params, R Result] func(context.Context, *ServerRequest[P]) (R, error)
|
||||
)
|
||||
|
||||
type paramsPtr[T any] interface {
|
||||
*T
|
||||
|
|
@ -208,11 +212,45 @@ const (
|
|||
missingParamsOK // params may be missing or null
|
||||
)
|
||||
|
||||
func newClientMethodInfo[P paramsPtr[T], R Result, T any](d typedClientMethodHandler[P, R], flags methodFlags) methodInfo {
|
||||
mi := newMethodInfo[P, R](flags)
|
||||
mi.newRequest = func(s Session, p Params) Request {
|
||||
r := &ClientRequest[P]{Session: s.(*ClientSession)}
|
||||
if p != nil {
|
||||
r.Params = p.(P)
|
||||
}
|
||||
return r
|
||||
}
|
||||
mi.handleMethod = MethodHandler(func(ctx context.Context, _ string, req Request) (Result, error) {
|
||||
return d(ctx, req.(*ClientRequest[P]))
|
||||
})
|
||||
return mi
|
||||
}
|
||||
|
||||
func newServerMethodInfo[P paramsPtr[T], R Result, T any](d typedServerMethodHandler[P, R], flags methodFlags) methodInfo {
|
||||
mi := newMethodInfo[P, R](flags)
|
||||
mi.newRequest = func(s Session, p Params) Request {
|
||||
r := &ServerRequest[P]{Session: s.(*ServerSession)}
|
||||
if p != nil {
|
||||
r.Params = p.(P)
|
||||
}
|
||||
return r
|
||||
}
|
||||
mi.handleMethod = MethodHandler(func(ctx context.Context, _ string, req Request) (Result, error) {
|
||||
rf := &ServerRequest[P]{Session: req.GetSession().(*ServerSession)}
|
||||
if req.GetParams() != nil {
|
||||
rf.Params = req.GetParams().(P)
|
||||
}
|
||||
return d(ctx, rf)
|
||||
})
|
||||
return mi
|
||||
}
|
||||
|
||||
// newMethodInfo creates a methodInfo from a typedMethodHandler.
|
||||
//
|
||||
// If isRequest is set, the method is treated as a request rather than a
|
||||
// notification.
|
||||
func newMethodInfo[S Session, P paramsPtr[T], R Result, T any](d typedMethodHandler[S, P, R], flags methodFlags) methodInfo {
|
||||
func newMethodInfo[P paramsPtr[T], R Result, T any](flags methodFlags) methodInfo {
|
||||
return methodInfo{
|
||||
flags: flags,
|
||||
unmarshalParams: func(m json.RawMessage) (Params, error) {
|
||||
|
|
@ -234,12 +272,6 @@ func newMethodInfo[S Session, P paramsPtr[T], R Result, T any](d typedMethodHand
|
|||
}
|
||||
return orZero[Params](p), nil
|
||||
},
|
||||
handleMethod: MethodHandler[S](func(ctx context.Context, session S, _ string, params Params) (Result, error) {
|
||||
if params == nil {
|
||||
return d(ctx, session, nil)
|
||||
}
|
||||
return d(ctx, session, params.(P))
|
||||
}),
|
||||
// newResult is used on the send side, to construct the value to unmarshal the result into.
|
||||
// R is a pointer to a result struct. There is no way to "unpointer" it without reflection.
|
||||
// TODO(jba): explore generic approaches to this, perhaps by treating R in
|
||||
|
|
@ -250,26 +282,33 @@ func newMethodInfo[S Session, P paramsPtr[T], R Result, T any](d typedMethodHand
|
|||
|
||||
// serverMethod is glue for creating a typedMethodHandler from a method on Server.
|
||||
func serverMethod[P Params, R Result](
|
||||
f func(*Server, context.Context, *ServerSession, P) (R, error),
|
||||
) typedMethodHandler[*ServerSession, P, R] {
|
||||
return func(ctx context.Context, ss *ServerSession, p P) (R, error) {
|
||||
return f(ss.server, ctx, ss, p)
|
||||
f func(*Server, context.Context, *ServerRequest[P]) (R, error),
|
||||
) typedServerMethodHandler[P, R] {
|
||||
return func(ctx context.Context, req *ServerRequest[P]) (R, error) {
|
||||
return f(req.Session.server, ctx, req)
|
||||
}
|
||||
}
|
||||
|
||||
// clientMethod is glue for creating a typedMethodHandler from a method on Client.
|
||||
func clientMethod[P Params, R Result](
|
||||
f func(*Client, context.Context, *ClientSession, P) (R, error),
|
||||
) typedMethodHandler[*ClientSession, P, R] {
|
||||
return func(ctx context.Context, cs *ClientSession, p P) (R, error) {
|
||||
return f(cs.client, ctx, cs, p)
|
||||
f func(*Client, context.Context, *ClientRequest[P]) (R, error),
|
||||
) typedClientMethodHandler[P, R] {
|
||||
return func(ctx context.Context, req *ClientRequest[P]) (R, error) {
|
||||
return f(req.Session.client, ctx, req)
|
||||
}
|
||||
}
|
||||
|
||||
// sessionMethod is glue for creating a typedMethodHandler from a method on ServerSession.
|
||||
func sessionMethod[S Session, P Params, R Result](f func(S, context.Context, P) (R, error)) typedMethodHandler[S, P, R] {
|
||||
return func(ctx context.Context, sess S, p P) (R, error) {
|
||||
return f(sess, ctx, p)
|
||||
// serverSessionMethod is glue for creating a typedServerMethodHandler from a method on ServerSession.
|
||||
func serverSessionMethod[P Params, R Result](f func(*ServerSession, context.Context, P) (R, error)) typedServerMethodHandler[P, R] {
|
||||
return func(ctx context.Context, req *ServerRequest[P]) (R, error) {
|
||||
return f(req.GetSession().(*ServerSession), ctx, req.Params)
|
||||
}
|
||||
}
|
||||
|
||||
// clientSessionMethod is glue for creating a typedMethodHandler from a method on ServerSession.
|
||||
func clientSessionMethod[P Params, R Result](f func(*ClientSession, context.Context, P) (R, error)) typedClientMethodHandler[P, R] {
|
||||
return func(ctx context.Context, req *ClientRequest[P]) (R, error) {
|
||||
return f(req.GetSession().(*ClientSession), ctx, req.Params)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -282,16 +321,9 @@ const (
|
|||
CodeUnsupportedMethod = -31001
|
||||
)
|
||||
|
||||
func callNotificationHandler[S Session, P any](ctx context.Context, h func(context.Context, S, *P), sess S, params *P) (Result, error) {
|
||||
if h != nil {
|
||||
h(ctx, sess, params)
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// notifySessions calls Notify on all the sessions.
|
||||
// Should be called on a copy of the peer sessions.
|
||||
func notifySessions[S Session](sessions []S, method string, params Params) {
|
||||
func notifySessions[S Session, P Params](sessions []S, method string, params P) {
|
||||
if sessions == nil {
|
||||
return
|
||||
}
|
||||
|
|
@ -299,13 +331,25 @@ func notifySessions[S Session](sessions []S, method string, params Params) {
|
|||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
for _, s := range sessions {
|
||||
if err := handleNotify(ctx, s, method, params); err != nil {
|
||||
req := newRequest(s, params)
|
||||
if err := handleNotify(ctx, method, req); err != nil {
|
||||
// TODO(jba): surface this error better
|
||||
log.Printf("calling %s: %v", method, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func newRequest[S Session, P Params](s S, p P) Request {
|
||||
switch s := any(s).(type) {
|
||||
case *ClientSession:
|
||||
return &ClientRequest[P]{Session: s, Params: p}
|
||||
case *ServerSession:
|
||||
return &ServerRequest[P]{Session: s, Params: p}
|
||||
default:
|
||||
panic("bad session")
|
||||
}
|
||||
}
|
||||
|
||||
// Meta is additional metadata for requests, responses and other types.
|
||||
type Meta map[string]any
|
||||
|
||||
|
|
@ -335,6 +379,43 @@ func setProgressToken(p Params, pt any) {
|
|||
m[progressTokenKey] = pt
|
||||
}
|
||||
|
||||
// A Request is a method request with parameters and additional information, such as the session.
|
||||
// Request is implemented by [*ClientRequest] and [*ServerRequest].
|
||||
type Request interface {
|
||||
isRequest()
|
||||
GetSession() Session
|
||||
GetParams() Params
|
||||
}
|
||||
|
||||
// A ClientRequest is a request to a client.
|
||||
type ClientRequest[P Params] struct {
|
||||
Session *ClientSession
|
||||
Params P
|
||||
}
|
||||
|
||||
// A ServerRequest is a request to a server.
|
||||
type ServerRequest[P Params] struct {
|
||||
Session *ServerSession
|
||||
Params P
|
||||
}
|
||||
|
||||
func (*ClientRequest[P]) isRequest() {}
|
||||
func (*ServerRequest[P]) isRequest() {}
|
||||
|
||||
func (r *ClientRequest[P]) GetSession() Session { return r.Session }
|
||||
func (r *ServerRequest[P]) GetSession() Session { return r.Session }
|
||||
|
||||
func (r *ClientRequest[P]) GetParams() Params { return r.Params }
|
||||
func (r *ServerRequest[P]) GetParams() Params { return r.Params }
|
||||
|
||||
func serverRequestFor[P Params](s *ServerSession, p P) *ServerRequest[P] {
|
||||
return &ServerRequest[P]{Session: s, Params: p}
|
||||
}
|
||||
|
||||
func clientRequestFor[P Params](s *ClientSession, p P) *ClientRequest[P] {
|
||||
return &ClientRequest[P]{Session: s, Params: p}
|
||||
}
|
||||
|
||||
// Params is a parameter (input) type for an MCP call or notification.
|
||||
type Params interface {
|
||||
// isParams discourages implementation of Params outside of this package.
|
||||
|
|
|
|||
Loading…
Reference in a new issue