Route go interceptor methods through receiving middleware (#10)

<!-- Provide a brief summary of your changes -->

Route `interceptors/list` and `interceptor/invoke` through the receiving
middleware chain so that middleware installed (audit, etc.) fires for
interceptor protocol methods the same way it does for standard MCP
methods.

## Motivation and Context
<!-- Why is this change needed? What problem does it solve? -->

Before this change, custom JSON-RPC methods were dispatched via a
separate code path that bypassed receiving middleware entirely. This
meant any middleware an operator installed would silently not apply to
the interceptor protocol endpoints.

The fix adds `AddReceivingCustomMethod[P, R, T]` as a package-level
generic function to the vendored go-sdk/mcp package. Custom methods
registered through it are stored in a per-server customMethods map and
merged into `receivingMethodInfos()`, so `defaultReceivingMethodHandler`
dispatches them exactly like built-in methods. External types satisfy
the sealed `Params`/`Result` interfaces by embedding the
`ParamsBase`/`ResultBase` types exported from the mcp package.

## How Has This Been Tested?
<!-- Have you tested this in a real application? Which scenarios were
tested? -->

Existing test suites and new ones for verifying middleware paths.

## Breaking Changes
<!-- Will users need to update their code or configurations? -->

no.

## Types of changes
<!-- What types of changes does your code introduce? Put an `x` in all
the boxes that apply: -->
- [x] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update

## Checklist
<!-- Go over all the following points, and put an `x` in all the boxes
that apply. -->
- [x] I have read the [MCP
Documentation](https://modelcontextprotocol.io)
- [x] My code follows the repository's style guidelines
- [x] New and existing tests pass locally
- [x] I have added appropriate error handling
- [x] I have added or updated documentation as needed

## Additional context
<!-- Add any other context, implementation notes, or design decisions
-->

The changes to `vendor/go-sdk/mcp/server.go` are intentionally minimal
and carry a `NOTE` comment acknowledging they are additions to the
upstream SDK needed for this reference implementation.

---------

Signed-off-by: Ukjae Jeong <jeongukjae@gmail.com>
This commit is contained in:
Ukjae Jeong 2026-04-30 17:14:55 +01:00 • committed by GitHub
commit 5bd7ab4d2a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 170 additions and 90 deletions

View file

@ -57,20 +57,18 @@ func setupChainWithInterceptors(t *testing.T, is ...interceptors.Interceptor) *c
// registerInterceptorMethods adds interceptors/list and interceptor/invoke
// custom methods to the server, backed by the given interceptor list.
func registerInterceptorMethods(server *mcp.Server, is []interceptors.Interceptor) {
server.AddCustomMethod(interceptors.MethodList,
func(_ context.Context, _ *mcp.ServerSession, raw json.RawMessage) (any, error) {
var params interceptors.ListParams
if len(raw) > 0 {
if err := json.Unmarshal(raw, &params); err != nil {
return nil, err
}
mcp.AddReceivingCustomMethod(server, interceptors.MethodList,
func(_ context.Context, req *mcp.ServerRequest[*interceptors.ListParams]) (*interceptors.ListResult, error) {
var event string
if req.Params != nil {
event = req.Params.Event
}
infos := make([]interceptors.InterceptorInfo, 0, len(is))
for _, i := range is {
if params.Event != "" {
if event != "" {
match := false
for _, e := range i.GetMetadata().Hook.Events {
if e == params.Event {
if e == event {
match = true
break
}
@ -85,12 +83,12 @@ func registerInterceptorMethods(server *mcp.Server, is []interceptors.Intercepto
},
)
server.AddCustomMethod(interceptors.MethodInvoke,
func(ctx context.Context, _ *mcp.ServerSession, raw json.RawMessage) (any, error) {
var params interceptors.InvokeParams
if err := json.Unmarshal(raw, &params); err != nil {
return nil, err
mcp.AddReceivingCustomMethod(server, interceptors.MethodInvoke,
func(ctx context.Context, req *mcp.ServerRequest[*interceptors.InvokeParams]) (*interceptors.InvokeResult, error) {
if req.Params == nil {
return nil, fmt.Errorf("params required")
}
params := req.Params
var target interceptors.Interceptor
for _, i := range is {
if i.GetMetadata().Name == params.Name {

View file

@ -6,7 +6,6 @@ package extension
import (
"context"
"encoding/json"
"fmt"
"slices"
"time"
@ -19,18 +18,16 @@ import (
// handleList implements the "interceptors/list" JSON-RPC method.
// It returns all registered interceptors, optionally filtered by event.
func (e *Extension) handleList(_ context.Context, _ *mcp.ServerSession, raw json.RawMessage) (any, error) {
var params interceptors.ListParams
if len(raw) > 0 {
if err := json.Unmarshal(raw, &params); err != nil {
return nil, err
}
func (e *Extension) handleList(_ context.Context, req *mcp.ServerRequest[*interceptors.ListParams]) (*interceptors.ListResult, error) {
var event string
if req.Params != nil {
event = req.Params.Event
}
all := e.getInterceptors()
infos := make([]interceptors.InterceptorInfo, 0, len(all))
for _, ri := range all {
if params.Event != "" && !slices.Contains(ri.GetMetadata().Hook.Events, params.Event) {
if event != "" && !slices.Contains(ri.GetMetadata().Hook.Events, event) {
continue
}
infos = append(infos, interceptors.InfoFromInterceptor(ri))
@ -41,17 +38,14 @@ func (e *Extension) handleList(_ context.Context, _ *mcp.ServerSession, raw json
// handleInvoke implements the "interceptor/invoke" JSON-RPC method.
// It invokes a single interceptor by name and returns its result.
func (e *Extension) handleInvoke(ctx context.Context, _ *mcp.ServerSession, raw json.RawMessage) (any, error) {
var params interceptors.InvokeParams
if err := json.Unmarshal(raw, &params); err != nil {
return nil, err
}
if params.Name == "" {
func (e *Extension) handleInvoke(ctx context.Context, req *mcp.ServerRequest[*interceptors.InvokeParams]) (*interceptors.InvokeResult, error) {
if req.Params == nil || req.Params.Name == "" {
return nil, &jsonrpc.Error{
Code: jsonrpc.CodeInvalidParams,
Message: "name is required",
}
}
params := req.Params
// Look up the interceptor by name.
i := e.findByName(params.Name)

View file

@ -60,8 +60,9 @@ func New(opts ...Option) *Extension {
// initialize response with interceptor capabilities on the given
// server. It can be called on multiple servers.
func (e *Extension) Install(server *mcp.Server) {
server.AddCustomMethod(interceptors.MethodList, e.handleList)
server.AddCustomMethod(interceptors.MethodInvoke, e.handleInvoke)
// Register JSON-RPC methods for interceptor discovery and invocation.
mcp.AddReceivingCustomMethod(server, interceptors.MethodList, e.handleList)
mcp.AddReceivingCustomMethod(server, interceptors.MethodInvoke, e.handleInvoke)
server.AddReceivingMiddleware(e.initMiddleware())
}

View file

@ -610,3 +610,74 @@ func TestServer_CombinedValidatorsAndMutators(t *testing.T) {
assert.Equal(t, int32(3), respMutCount.Load(), "expected 3 response mutators to run")
assert.Equal(t, int32(3), respValCount.Load(), "expected 3 response validators to run")
}
// TestCustomMethodsFlowThroughMiddleware verifies that interceptors/list and
// interceptor/invoke are routed through AddReceivingMiddleware, just like
// standard MCP methods.
func TestCustomMethodsFlowThroughMiddleware(t *testing.T) {
t.Parallel()
var count atomic.Int64
counting := func(next mcp.MethodHandler) mcp.MethodHandler {
return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) {
count.Add(1)
return next(ctx, method, req)
}
}
mcpServer := buildServer(t, allowAllValidator("v1"))
mcpServer.AddReceivingMiddleware(counting)
cs := connectHTTPClient(t, mcpServer)
count.Store(0) // reset after initialization handshake
_, err := cs.CallTool(context.Background(), &mcp.CallToolParams{
Name: "echo",
Arguments: map[string]any{"text": "hello"},
})
require.NoError(t, err)
assert.Greater(t, count.Swap(0), int64(0), "middleware must fire for tools/call")
var listResult interceptors.ListResult
err = cs.CallCustom(context.Background(), interceptors.MethodList, nil, &listResult)
require.NoError(t, err)
assert.Len(t, listResult.Interceptors, 1)
assert.Greater(t, count.Swap(0), int64(0), "middleware must fire for interceptors/list")
payload, _ := json.Marshal(map[string]any{"name": "echo", "arguments": map[string]any{}})
var invokeResult interceptors.InvokeResult
err = cs.CallCustom(context.Background(), interceptors.MethodInvoke, &interceptors.InvokeParams{
Name: "v1",
Event: interceptors.EventToolsCall,
Phase: interceptors.PhaseRequest,
Payload: payload,
}, &invokeResult)
require.NoError(t, err)
assert.Greater(t, count.Swap(0), int64(0), "middleware must fire for interceptor/invoke")
}
// TestMiddlewareCanRejectCustomMethods verifies that middleware installed via
// AddReceivingMiddleware can block interceptors/list and interceptor/invoke,
// enabling auth guards over the interceptor protocol.
func TestMiddlewareCanRejectCustomMethods(t *testing.T) {
t.Parallel()
var rejected atomic.Bool
blockingMiddleware := func(next mcp.MethodHandler) mcp.MethodHandler {
return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) {
if method != "initialize" {
rejected.Store(true)
return nil, fmt.Errorf("unauthorized")
}
return next(ctx, method, req)
}
}
mcpServer := buildServer(t, allowAllValidator("v1"))
mcpServer.AddReceivingMiddleware(blockingMiddleware)
cs := connectHTTPClient(t, mcpServer)
var listResult interceptors.ListResult
err := cs.CallCustom(context.Background(), interceptors.MethodList, nil, &listResult)
assert.Error(t, err, "blockingMiddleware should reject interceptors/list")
assert.True(t, rejected.Load())
}

View file

@ -133,53 +133,53 @@ func resultText(t *testing.T, result *mcp.CallToolResult) string {
return tc.Text
}
// --- RPC server setup ---
// --- Low-level server primitives ---
// setupRPCServer creates an interceptor extension with the given interceptors
// and returns a connected client session. The client can call custom methods
// via raw JSON-RPC.
func setupRPCServer(t *testing.T, is ...interceptors.Interceptor) *mcp.ClientSession {
// buildServer creates an interceptor server with the echo tool and the given
// interceptors. It does not start serving or connect a client, so callers can
// install middleware before calling connectHTTPClient.
func buildServer(t *testing.T, is ...interceptors.Interceptor) *mcp.Server {
t.Helper()
mcpServer := mcp.NewServer(&mcp.Implementation{
Name: "test-server",
Version: "0.1.0",
}, nil)
// A minimal tool so the server has something to initialize with.
mcpServer.AddTool(&mcp.Tool{
Name: "echo",
Description: "echoes input",
InputSchema: map[string]any{"type": "object"},
}, func(_ context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
return &mcp.CallToolResult{
Content: []mcp.Content{&mcp.TextContent{Text: fmt.Sprintf("echo: %s", req.Params.Arguments)}},
}, nil
})
mcpServer := mcp.NewServer(&mcp.Implementation{Name: "test-server", Version: "0.1.0"}, nil)
for _, tt := range defaultTools() {
mcpServer.AddTool(tt.tool, tt.handler)
}
ext := extension.New()
for _, i := range is {
ext.AddInterceptor(i)
}
ext.Install(mcpServer)
return mcpServer
}
// connectHTTPClient starts an HTTP test server backed by srv and returns a
// connected client session. Cleanup is registered via t.Cleanup.
func connectHTTPClient(t *testing.T, srv *mcp.Server) *mcp.ClientSession {
t.Helper()
handler := mcp.NewStreamableHTTPHandler(
func(r *http.Request) *mcp.Server { return mcpServer },
func(r *http.Request) *mcp.Server { return srv },
nil,
)
httpServer := httptest.NewServer(handler)
t.Cleanup(httpServer.Close)
client := mcp.NewClient(&mcp.Implementation{Name: "test-client", Version: "0.1.0"}, nil)
cs, err := client.Connect(context.Background(), &mcp.StreamableClientTransport{
Endpoint: httpServer.URL,
}, nil)
require.NoError(t, err)
t.Cleanup(func() { cs.Close() })
return cs
}
// setupRPCServer creates an interceptor server with the given interceptors
// and returns a connected client session for raw JSON-RPC method calls.
func setupRPCServer(t *testing.T, is ...interceptors.Interceptor) *mcp.ClientSession {
t.Helper()
srv := buildServer(t, is...)
return connectHTTPClient(t, srv)
}
// --- Interceptor builders ---
// All handlers work with json.RawMessage payloads since that's what
// interceptor/invoke delivers.

View file

@ -4,7 +4,11 @@
package interceptors
import "encoding/json"
import (
"encoding/json"
"github.com/modelcontextprotocol/go-sdk/mcp"
)
// JSON-RPC method names for the interceptor protocol.
const (
@ -28,6 +32,7 @@ const (
// ListParams are the optional parameters for the interceptors/list method.
type ListParams struct {
mcp.ParamsBase
Event string `json:"event,omitempty"` // Filter by event name; empty = all
}
@ -64,6 +69,7 @@ func InfoFromInterceptor(i Interceptor) InterceptorInfo {
// ListResult is the result of the interceptors/list method.
type ListResult struct {
mcp.ResultBase
Interceptors []InterceptorInfo `json:"interceptors"`
}
@ -72,6 +78,7 @@ type ListResult struct {
// InvokeParams are the parameters for the interceptor/invoke method.
// The caller specifies which interceptor to invoke by name.
type InvokeParams struct {
mcp.ParamsBase
Name string `json:"name"`
Event string `json:"event"`
Phase InterceptionPhase `json:"phase"`
@ -84,6 +91,7 @@ type InvokeParams struct {
// InvokeResult is the result of the interceptor/invoke method.
// Fields are type-specific: validators populate Validation, mutators populate Mutation + Payload.
type InvokeResult struct {
mcp.ResultBase
Interceptor string `json:"interceptor"`
Type InterceptorType `json:"type"`
Phase InterceptionPhase `json:"phase"`

View file

@ -54,27 +54,46 @@ type Server struct {
receivingMethodHandler_ MethodHandler
resourceSubscriptions map[string]map[*ServerSession]bool // uri -> session -> bool
pendingNotifications map[string]*time.Timer // notification name -> timer for pending notification send
customMethods map[string]CustomMethodHandler
customMethods map[string]methodInfo
}
// CustomMethodHandler handles a custom JSON-RPC method registered via
// [Server.AddCustomMethod]. It receives the raw JSON params and returns
// a result that will be marshaled to JSON.
type CustomMethodHandler func(ctx context.Context, ss *ServerSession, params json.RawMessage) (any, error)
// NOTE: ParamsBase, ResultBase, and AddReceivingCustomMethod are additions to
// the upstream go-sdk to route custom JSON-RPC methods through receiving
// handler calls like any built-in MCP method, as needed by the MCP interceptors
// extension (SEP). This is a reference implementation, so there's some
// inefficiency; receivingMethodInfos() merges maps on every request instead of
// caching.
// AddCustomMethod registers a custom JSON-RPC method on the server.
// Custom methods bypass the typed Params/Result pipeline and operate
// on raw JSON. They are dispatched before the standard method handlers.
//
// This is useful for extensions that define their own JSON-RPC methods
// (e.g. interceptors/list, interceptor/executeChain).
func (s *Server) AddCustomMethod(method string, handler CustomMethodHandler) {
// ParamsBase can be embedded by types outside this package to satisfy the
// mcp.Params interface, including its unexported marker method.
type ParamsBase struct {
Meta `json:"_meta,omitempty"`
}
func (x *ParamsBase) isParams() {}
func (x *ParamsBase) GetProgressToken() any { return getProgressToken(x) }
func (x *ParamsBase) SetProgressToken(t any) { setProgressToken(x, t) }
// ResultBase can be embedded by types outside this package to satisfy the
// mcp.Result interface, including its unexported marker method.
type ResultBase struct {
Meta `json:"_meta,omitempty"`
}
func (x *ResultBase) isResult() {}
// AddReceivingCustomMethod registers a typed custom JSON-RPC method that flows
// through receiving handler calls like any built-in MCP method. P must embed
// ParamsBase and R must embed ResultBase to satisfy the interface constraints.
// Go does not support generic methods, so this is a package-level function.
func AddReceivingCustomMethod[P paramsPtr[T], R Result, T any](
s *Server, method string,
handler typedServerMethodHandler[P, R],
) {
s.mu.Lock()
defer s.mu.Unlock()
if s.customMethods == nil {
s.customMethods = make(map[string]CustomMethodHandler)
}
s.customMethods[method] = handler
// missingParamsOK: custom methods may legally omit params (e.g. interceptors/list).
s.customMethods[method] = newServerMethodInfo(handler, missingParamsOK)
}
// hasCustomMethod reports whether the given method is registered as a custom method.
@ -85,17 +104,6 @@ func (s *Server) hasCustomMethod(method string) bool {
return ok
}
// callCustomMethod calls the custom method handler for the given method.
func (s *Server) callCustomMethod(ctx context.Context, ss *ServerSession, method string, params json.RawMessage) (any, error) {
s.mu.Lock()
h := s.customMethods[method]
s.mu.Unlock()
if h == nil {
return nil, fmt.Errorf("unknown custom method %q", method)
}
return h(ctx, ss, params)
}
// ServerOptions is used to configure behavior of the server.
type ServerOptions struct {
// Optional instructions for connected clients.
@ -235,6 +243,7 @@ func NewServer(impl *Implementation, options *ServerOptions) *Server {
receivingMethodHandler_: defaultReceivingMethodHandler[*ServerSession],
resourceSubscriptions: make(map[string]map[*ServerSession]bool),
pendingNotifications: make(map[string]*time.Timer),
customMethods: make(map[string]methodInfo),
}
}
@ -1426,7 +1435,12 @@ func initializeMethodInfo() methodInfo {
func (ss *ServerSession) sendingMethodInfos() map[string]methodInfo { return clientMethodInfos }
func (ss *ServerSession) receivingMethodInfos() map[string]methodInfo { return serverMethodInfos }
func (ss *ServerSession) receivingMethodInfos() map[string]methodInfo {
infos := make(map[string]methodInfo)
maps.Copy(infos, serverMethodInfos)
maps.Copy(infos, ss.server.customMethods)
return infos
}
func (ss *ServerSession) sendingMethodHandler() MethodHandler {
s := ss.server
@ -1474,12 +1488,6 @@ func (ss *ServerSession) handle(ctx context.Context, req *jsonrpc.Request) (any,
// server->client calls and notifications to the incoming request from which
// they originated. See [idContextKey] for details.
ctx = context.WithValue(ctx, idContextKey{}, req.ID)
// Dispatch custom methods before the standard typed pipeline.
if ss.server.hasCustomMethod(req.Method) {
return ss.server.callCustomMethod(ctx, ss, req.Method, req.Params)
}
return handleReceive(ctx, ss, req)
}