modelcontextprotocol/go-sdk
Watch
1
0
Fork
You've already forked go-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

internal/mcp: simplify session type params

Browse source 
Write a separate constraint for ClientSession | ServerSession,
and change the pointerness of the type parameter.

There is still a dynamic cast, but overall I think this is a bit
simpler:

- The exported types have the more readable constraint "Session"
  instead of "ClientSession | ServerSession".

- We can put the unexported methods in Session instead of needing a
  separate interface (previously "session").

- This CL enables a simplified treatment of sending middleware, in
  the next CL.

Change-Id: I41184fe93ea201d1da29dbfba26ac59f93771277
Reviewed-on: https://go-review.googlesource.com/c/tools/+/676175
Reviewed-by: Sam Thanawalla <samthanawalla@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Robert Findley <rfindley@google.com>
This commit is contained in:
Jonathan Amsterdam 2025-05-25 07:35:56 -04:00
commit 19d58e6955
1 changed files with 67 additions and 72 deletions
Download patch file Download diff file Expand all files Collapse all files

89
mcp/shared.go
View file

@ -24,48 +24,45 @@ import (
// 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 ClientSession | ServerSession] func(
ctx context.Context, _ *S, method string, params Params) (result Result, err error)
type MethodHandler[S Session] func(
ctx context.Context, _ S, method string, params Params) (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 {
*ClientSession | *ServerSession
methodHandler() methodHandler
methodInfos() map[string]methodInfo
getConn() *jsonrpc2.Connection
}
// Middleware is a function from MethodHandlers to MethodHandlers.
type Middleware[S ClientSession | ServerSession] func(MethodHandler[S]) MethodHandler[S]
type Middleware[S Session] func(MethodHandler[S]) MethodHandler[S]
// addMiddleware wraps the handler in the middleware functions.
func addMiddleware[S ClientSession | ServerSession](handlerp *MethodHandler[S], middleware []Middleware[S]) {
func addMiddleware[S Session](handlerp *MethodHandler[S], middleware []Middleware[S]) {
for _, m := range slices.Backward(middleware) {
*handlerp = m(*handlerp)
}
}
// session has methods common to both ClientSession and ServerSession.
type session[S ClientSession | ServerSession] interface {
methodHandler() MethodHandler[S]
methodInfos() map[string]methodInfo[S]
getConn() *jsonrpc2.Connection
}
// toSession[S] converts its argument to a session[S].
// Note that since S is constrained to ClientSession | ServerSession, and pointers to those
// types both implement session[S] already, this should be a no-op.
// That it is not, is due (I believe) to a deficency in generics, possibly related to core types.
// TODO(jba): revisit in Go 1.26; perhaps the change in spec due to the removal of core types
// will have resulted by then in a more generous implementation.
func toSession[S ClientSession | ServerSession](sess *S) session[S] {
return any(sess).(session[S])
}
// defaultMethodHandler is the initial MethodHandler for servers and clients, before being wrapped by middleware.
func defaultMethodHandler[S ClientSession | ServerSession](ctx context.Context, sess *S, method string, params Params) (Result, error) {
info, ok := toSession(sess).methodInfos()[method]
func defaultMethodHandler[S Session](ctx context.Context, session S, method string, params Params) (Result, error) {
info, ok := session.methodInfos()[method]
if !ok {
// This can be called from user code, with an arbitrary value for method.
return nil, jsonrpc2.ErrNotHandled
}
return info.handleMethod(ctx, sess, method, params)
return info.handleMethod.(MethodHandler[S])(ctx, session, method, params)
}
func handleRequest[S ClientSession | ServerSession](ctx context.Context, req *jsonrpc2.Request, sess *S) (any, error) {
info, ok := toSession(sess).methodInfos()[req.Method]
func handleRequest[S Session](ctx context.Context, req *jsonrpc2.Request, session S) (any, error) {
info, ok := session.methodInfos()[req.Method]
if !ok {
return nil, jsonrpc2.ErrNotHandled
}
@ -75,8 +72,8 @@ func handleRequest[S ClientSession | ServerSession](ctx context.Context, req *js
}
// mh might be user code, so ensure that it returns the right values for the jsonrpc2 protocol.
mh := toSession(sess).methodHandler()
res, err := mh(ctx, sess, req.Method, params)
mh := session.methodHandler().(MethodHandler[S])
res, err := mh(ctx, session, req.Method, params)
if err != nil {
return nil, err
}
@ -84,11 +81,11 @@ func handleRequest[S ClientSession | ServerSession](ctx context.Context, req *js
}
// methodInfo is information about invoking a method.
type methodInfo[TSession ClientSession | ServerSession] struct {
type methodInfo struct {
// unmarshal params from the wire into an XXXParams struct
unmarshalParams func(json.RawMessage) (Params, error)
// run the code for the method
handleMethod MethodHandler[TSession]
handleMethod methodHandler
}
// The following definitions support converting from typed to untyped method handlers.
@ -98,11 +95,11 @@ type methodInfo[TSession ClientSession | ServerSession] struct {
// - R: results
// A typedMethodHandler is like a MethodHandler, but with type information.
type typedMethodHandler[S any, P Params, R Result] func(context.Context, *S, P) (R, error)
type typedMethodHandler[S Session, P Params, R Result] func(context.Context, S, P) (R, error)
// newMethodInfo creates a methodInfo from a typedMethodHandler.
func newMethodInfo[S ClientSession | ServerSession, P Params, R Result](d typedMethodHandler[S, P, R]) methodInfo[S] {
return methodInfo[S]{
func newMethodInfo[S Session, P Params, R Result](d typedMethodHandler[S, P, R]) methodInfo {
return methodInfo{
unmarshalParams: func(m json.RawMessage) (Params, error) {
var p P
if m != nil {
@ -112,29 +109,33 @@ func newMethodInfo[S ClientSession | ServerSession, P Params, R Result](d typedM
}
return p, nil
},
handleMethod: func(ctx context.Context, ss *S, _ string, params Params) (Result, error) {
return d(ctx, ss, params.(P))
},
handleMethod: MethodHandler[S](func(ctx context.Context, session S, _ string, params Params) (Result, error) {
return d(ctx, session, params.(P))
}),
}
}
// 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] {
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)
}
}
// clientMethod is glue for creating a typedMethodHandler from a method on Server.
func clientMethod[P Params, R Result](f func(*Client, context.Context, *ClientSession, P) (R, error)) typedMethodHandler[ClientSession, P, R] {
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)
}
}
// sessionMethod is glue for creating a typedMethodHandler from a method on ServerSession.
func sessionMethod[S ClientSession | ServerSession, 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) {
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)
}
}
@ -152,7 +153,7 @@ const (
CodeUnsupportedMethod = -31001
)
func callNotificationHandler[S ClientSession | ServerSession, P any](ctx context.Context, h func(context.Context, *S, *P), sess *S, params *P) (Result, error) {
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)
}
@ -161,15 +162,15 @@ func callNotificationHandler[S ClientSession | ServerSession, P any](ctx context
// notifySessions calls Notify on all the sessions.
// Should be called on a copy of the peer sessions.
func notifySessions[S ClientSession | ServerSession](sessions []*S, method string, params Params) {
func notifySessions[S Session](sessions []S, method string, params Params) {
if sessions == nil {
return
}
// TODO: make this timeout configurable, or call Notify asynchronously.
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
for _, ss := range sessions {
if err := toSession(ss).getConn().Notify(ctx, method, params); err != nil {
for _, s := range sessions {
if err := s.getConn().Notify(ctx, method, params); err != nil {
// TODO(jba): surface this error better
log.Printf("calling %s: %v", method, err)
}