diff --git a/design/design.md b/design/design.md index abaf900..9670587 100644 --- a/design/design.md +++ b/design/design.md @@ -416,11 +416,11 @@ We provide a mechanism to add MCP-level middleware on the both the client and se // For methods, a MethodHandler must return either an XXResult struct pointer and a nil error, or // nil with a non-nil error. // For notifications, a MethodHandler must return nil, nil. -type MethodHandler[S ClientSession | ServerSession] func( - ctx context.Context, _ *S, method string, params any) (result any, err error) +type MethodHandler[S Session] func( + ctx context.Context, _ *S, method string, params Params) (result Result, err error) // 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 client/server's current method handler using the provided // middleware. Middleware is applied from right to left, so that the first one @@ -428,14 +428,14 @@ type Middleware[S ClientSession | ServerSession] func(MethodHandler[S]) MethodHa // // For example, AddMiddleware(m1, m2, m3) augments the server method handler as // m1(m2(m3(handler))). -func (c *Client) AddMiddleware(middleware ...Middleware[ClientSession]) -func (s *Server) AddMiddleware(middleware ...Middleware[ServerSession]) +func (c *Client) AddMiddleware(middleware ...Middleware[*ClientSession]) +func (s *Server) AddMiddleware(middleware ...Middleware[*ServerSession]) ``` As an example, this code adds server-side logging: ```go -func withLogging(h mcp.MethodHandler[ServerSession]) mcp.MethodHandler[ServerSession]{ +func withLogging(h mcp.MethodHandler[*ServerSession]) mcp.MethodHandler[*ServerSession]{ return func(ctx context.Context, s *mcp.ServerSession, method string, params any) (res any, err error) { log.Printf("request: %s %v", method, params) defer func() { log.Printf("response: %v, %v", res, err) }() diff --git a/mcp/client.go b/mcp/client.go index 0b1e583..9d712df 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -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_ diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 1dea894..198538b 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -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) diff --git a/mcp/server.go b/mcp/server.go index 94c146f..b275c40 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -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_ diff --git a/mcp/shared.go b/mcp/shared.go index d33ec6d..5582ac0 100644 --- a/mcp/shared.go +++ b/mcp/shared.go @@ -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) }