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
4 changed files with 67 additions and 72 deletions
Download patch file Download diff file Expand all files Collapse all files

14
mcp/client.go
View file

@ -24,7 +24,7 @@ type Client struct {
mu sync.Mutex
roots *featureSet[*Root]
sessions []*ClientSession
methodHandler_ MethodHandler[ClientSession]
methodHandler_ MethodHandler[*ClientSession]
}
// NewClient creates a new Client.
@ -37,7 +37,7 @@ func NewClient(name, version string, opts *ClientOptions) *Client {
name: name,
version: version,
roots: newFeatureSet(func(r *Root) string { return r.URI }),
methodHandler_: defaultMethodHandler[ClientSession],
methodHandler_: defaultMethodHandler[*ClientSession],
}
if opts != nil {
c.opts = *opts
@ -196,14 +196,14 @@ func (c *Client) createMessage(ctx context.Context, cs *ClientSession, params *C
//
// For example, AddMiddleware(m1, m2, m3) augments the client method handler as
// m1(m2(m3(handler))).
func (c *Client) AddMiddleware(middleware ...Middleware[ClientSession]) {
func (c *Client) AddMiddleware(middleware ...Middleware[*ClientSession]) {
c.mu.Lock()
defer c.mu.Unlock()
addMiddleware(&c.methodHandler_, middleware)
}
// clientMethodInfos maps from the RPC method name to serverMethodInfos.
var clientMethodInfos = map[string]methodInfo[ClientSession]{
var clientMethodInfos = map[string]methodInfo{
methodPing: newMethodInfo(sessionMethod((*ClientSession).ping)),
methodListRoots: newMethodInfo(clientMethod((*Client).listRoots)),
methodCreateMessage: newMethodInfo(clientMethod((*Client).createMessage)),
@ -213,9 +213,7 @@ var clientMethodInfos = map[string]methodInfo[ClientSession]{
notificationLoggingMessage: newMethodInfo(clientMethod((*Client).callLoggingHandler)),
}
var _ session[ClientSession] = (*ClientSession)(nil)
func (cs *ClientSession) methodInfos() map[string]methodInfo[ClientSession] {
func (cs *ClientSession) methodInfos() map[string]methodInfo {
return clientMethodInfos
}
@ -223,7 +221,7 @@ func (cs *ClientSession) handle(ctx context.Context, req *jsonrpc2.Request) (any
return handleRequest(ctx, req, cs)
}
func (cs *ClientSession) methodHandler() MethodHandler[ClientSession] {
func (cs *ClientSession) methodHandler() methodHandler {
cs.client.mu.Lock()
defer cs.client.mu.Unlock()
return cs.client.methodHandler_

8
mcp/mcp_test.go
View file

@ -634,10 +634,10 @@ func TestMiddleware(t *testing.T) {
// "1" is the outer middleware layer, called first; then "2" is called, and finally
// the default dispatcher.
s.AddMiddleware(traceCalls[ServerSession](&sbuf, "1"), traceCalls[ServerSession](&sbuf, "2"))
s.AddMiddleware(traceCalls[*ServerSession](&sbuf, "1"), traceCalls[*ServerSession](&sbuf, "2"))
c := NewClient("testClient", "v1.0.0", nil)
c.AddMiddleware(traceCalls[ClientSession](&cbuf, "1"), traceCalls[ClientSession](&cbuf, "2"))
c.AddMiddleware(traceCalls[*ClientSession](&cbuf, "1"), traceCalls[*ClientSession](&cbuf, "2"))
cs, err := c.Connect(ctx, ct)
if err != nil {
@ -679,9 +679,9 @@ func TestMiddleware(t *testing.T) {
// traceCalls creates a middleware function that prints the method before and after each call
// with the given prefix.
func traceCalls[S ClientSession | ServerSession](w io.Writer, prefix string) Middleware[S] {
func traceCalls[S Session](w io.Writer, prefix string) Middleware[S] {
return func(h MethodHandler[S]) MethodHandler[S] {
return func(ctx context.Context, sess *S, method string, params Params) (Result, error) {
return func(ctx context.Context, sess S, method string, params Params) (Result, error) {
fmt.Fprintf(w, "%s >%s\n", prefix, method)
defer fmt.Fprintf(w, "%s <%s\n", prefix, method)
return h(ctx, sess, method, params)

16
mcp/server.go
View file

@ -37,7 +37,7 @@ type Server struct {
tools *featureSet[*ServerTool]
resources *featureSet[*ServerResource]
sessions []*ServerSession
methodHandler_ MethodHandler[ServerSession]
methodHandler_ MethodHandler[*ServerSession]
}
// ServerOptions is used to configure behavior of the server.
@ -77,7 +77,7 @@ func NewServer(name, version string, opts *ServerOptions) *Server {
prompts: newFeatureSet(func(p *ServerPrompt) string { return p.Prompt.Name }),
tools: newFeatureSet(func(t *ServerTool) string { return t.Tool.Name }),
resources: newFeatureSet(func(r *ServerResource) string { return r.Resource.URI }),
methodHandler_: defaultMethodHandler[ServerSession],
methodHandler_: defaultMethodHandler[*ServerSession],
}
}
@ -440,14 +440,14 @@ func (ss *ServerSession) LoggingMessage(ctx context.Context, params *LoggingMess
//
// For example, AddMiddleware(m1, m2, m3) augments the server method handler as
// m1(m2(m3(handler))).
func (s *Server) AddMiddleware(middleware ...Middleware[ServerSession]) {
func (s *Server) AddMiddleware(middleware ...Middleware[*ServerSession]) {
s.mu.Lock()
defer s.mu.Unlock()
addMiddleware(&s.methodHandler_, middleware)
}
// serverMethodInfos maps from the RPC method name to serverMethodInfos.
var serverMethodInfos = map[string]methodInfo[ServerSession]{
var serverMethodInfos = map[string]methodInfo{
methodInitialize: newMethodInfo(sessionMethod((*ServerSession).initialize)),
methodPing: newMethodInfo(sessionMethod((*ServerSession).ping)),
methodListPrompts: newMethodInfo(serverMethod((*Server).listPrompts)),
@ -461,15 +461,11 @@ var serverMethodInfos = map[string]methodInfo[ServerSession]{
notificationRootsListChanged: newMethodInfo(serverMethod((*Server).callRootsListChangedHandler)),
}
// *ServerSession implements the session interface.
// See toSession for why this interface seems to be necessary.
var _ session[ServerSession] = (*ServerSession)(nil)
func (ss *ServerSession) methodInfos() map[string]methodInfo[ServerSession] {
func (ss *ServerSession) methodInfos() map[string]methodInfo {
return serverMethodInfos
}
func (ss *ServerSession) methodHandler() MethodHandler[ServerSession] {
func (ss *ServerSession) methodHandler() methodHandler {
ss.server.mu.Lock()
defer ss.server.mu.Unlock()
return ss.server.methodHandler_

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)
}