internal/mcp: dispatcher middleware

Implement the dispatcher middleware framework, allowing users to
wrap method dispatching with their own code.

The Dispatcher functions have an untyped params arg, but expect its
value to be a typed, unmarshaled XXXParams struct. That required
splitting the elegant dispatch[P, R] function into two parts, one
to do the unmarshaling and one to call the method. The nicest way
I could find to do that was to map from method names to a pair
of functions, one to unmarshal and one to call. We could also write
it as two switches on method name, which would require less machinery
but would duplicate the list of method names, already a tad fragile
because they are literal strings.

I also exported the initialize params and result, so user Dispatchers
could access them. All Params and Result types will have to be exported
for that reason.

Change-Id: Ie2c95cc17b4ef3d3d095dccdadb294e5f1590796
Reviewed-on: https://go-review.googlesource.com/c/tools/+/673515
Reviewed-by: Robert Findley <rfindley@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
This commit is contained in:
Jonathan Amsterdam 2025-05-15 16:35:46 -04:00
commit 1323fb5216

View file

@ -409,20 +409,20 @@ func (*ClientSession) ResourceTemplates(context.Context, *ListResourceTemplatesP
We provide a mechanism to add MCP-level middleware on the server side, which runs after the request has been parsed but before any normal handling.
```go
// A Dispatcher dispatches an MCP message to the appropriate handler.
// A ServerMethodHandler dispatches an MCP message to the appropriate handler.
// The params argument will be an XXXParams struct pointer, such as *GetPromptParams.
// The response if err is non-nil should be an XXXResult struct pointer.
type Dispatcher func(ctx context.Context, s *ServerSession, method string, params any) (result any, err error)
type ServerMethodHandler func(ctx context.Context, s *ServerSession, method string, params any) (result any, err error)
// AddDispatchers calls each function from right to left on the previous result, beginning
// AddMiddlewarecalls each function from right to left on the previous result, beginning
// with the server's current dispatcher, and installs the result as the new dispatcher.
func (*Server) AddDispatchers(middleware ...func(Dispatcher) Dispatcher))
func (*Server) AddMiddleware(middleware ...func(ServerMethodHandler) ServerMethodHandler)
```
As an example, this code adds server-side logging:
```go
func withLogging(h mcp.Dispatcher) mcp.Dispatcher {
func withLogging(h mcp.ServerMethodHandler) mcp.ServerMethodHandler{
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) }()
@ -430,9 +430,11 @@ func withLogging(h mcp.Dispatcher) mcp.Dispatcher {
}
}
server.AddDispatchers(withLogging)
server.AddMiddleware(withLogging)
```
We will provide the same functionality on the client side as well.
**Differences from mcp-go**: Version 0.26.0 of mcp-go defines 24 server hooks. Each hook consists of a field in the `Hooks` struct, a `Hooks.Add` method, and a type for the hook function. These are rarely used. The most common is `OnError`, which occurs fewer than ten times in open-source code.
### Errors