Skip to content

Commit 65ece59

Browse files
mcp: Allow registration of custom JSON-RPC methods (#956)
## Description This PR introduces support for custom (non-standard) JSON-RPC methods in the MCP Go SDK. This feature enables servers to handle arbitrary JSON-RPC methods and clients to invoke them, all while maintaining full type safety and integrating seamlessly with the SDK's existing middleware and session management. - Server-Side Registration: Added AddReceivingCustomMethod to allow servers to register handlers for custom methods. Incoming requests are automatically unmarshaled into user-defined parameter structs. - Client-Side Calling: Added AddSendingCustomMethod which returns a strongly-typed caller function that clients can use to invoke custom methods and receive strongly-typed results. - Helper Structs: Provided ParamsBase and ResultBase for users to embed in their custom structs. This allows them to easily satisfy the SDK's internal Params and Result interfaces. Fixes #954
1 parent 22d4988 commit 65ece59

6 files changed

Lines changed: 442 additions & 4 deletions

File tree

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
2+
// Use of this source code is governed by the license
3+
// that can be found in the LICENSE file.
4+
5+
// The custom-method example demonstrates registering and calling a custom
6+
// JSON-RPC method that is not part of the standard MCP spec.
7+
//
8+
// The server registers a "latin/translate" method that translates simple
9+
// English phrases into Latin. A client connects over an in-memory transport,
10+
// calls the custom method, and prints the result.
11+
package main
12+
13+
import (
14+
"context"
15+
"fmt"
16+
"log"
17+
"strings"
18+
19+
"github.com/modelcontextprotocol/go-sdk/mcp"
20+
)
21+
22+
type TranslateParams struct {
23+
mcp.ParamsBase
24+
Text string `json:"text"`
25+
}
26+
27+
type TranslateResult struct {
28+
mcp.ResultBase
29+
Latin string `json:"latin"`
30+
}
31+
32+
var translations = map[string]string{
33+
"hello": "salve",
34+
"goodbye": "vale",
35+
"thank you": "gratias tibi ago",
36+
"how are you": "quid agis",
37+
"good morning": "bonum mane",
38+
"good night": "bonam noctem",
39+
"friend": "amicus",
40+
"water": "aqua",
41+
"love": "amor",
42+
"war": "bellum",
43+
"peace": "pax",
44+
"truth": "veritas",
45+
"light": "lux",
46+
"time": "tempus",
47+
"life": "vita",
48+
"death": "mors",
49+
"star": "stella",
50+
"earth": "terra",
51+
"sea": "mare",
52+
"the die is cast": "alea iacta est",
53+
"i came i saw i conquered": "veni vidi vici",
54+
"seize the day": "carpe diem",
55+
}
56+
57+
func main() {
58+
ctx := context.Background()
59+
60+
server := mcp.NewServer(&mcp.Implementation{Name: "latin-server", Version: "v1.0.0"}, nil)
61+
62+
if err := mcp.AddReceivingCustomMethod(server, "latin/translate",
63+
func(ctx context.Context, ss *mcp.ServerSession, params *TranslateParams) (*TranslateResult, error) {
64+
key := strings.ToLower(strings.TrimSpace(params.Text))
65+
latin, ok := translations[key]
66+
if !ok {
67+
latin = fmt.Sprintf("[unknown: %q — try: %s]", params.Text, knownPhrases())
68+
}
69+
return &TranslateResult{Latin: latin}, nil
70+
}); err != nil {
71+
log.Fatal(err)
72+
}
73+
74+
ct, st := mcp.NewInMemoryTransports()
75+
76+
ss, err := server.Connect(ctx, st, nil)
77+
if err != nil {
78+
log.Fatal(err)
79+
}
80+
defer ss.Close()
81+
82+
client := mcp.NewClient(&mcp.Implementation{Name: "latin-client", Version: "v1.0.0"}, nil)
83+
if err := mcp.AddSendingCustomMethod[*TranslateParams, *TranslateResult](client, "latin/translate"); err != nil {
84+
log.Fatal(err)
85+
}
86+
87+
cs, err := client.Connect(ctx, ct, nil)
88+
if err != nil {
89+
log.Fatal(err)
90+
}
91+
defer cs.Close()
92+
93+
phrases := []string{"Hello", "Seize the day", "Peace", "Truth", "I came I saw I conquered"}
94+
for _, phrase := range phrases {
95+
result, err := mcp.CallCustomMethod[*TranslateParams, *TranslateResult](
96+
ctx, cs, "latin/translate", &TranslateParams{Text: phrase})
97+
if err != nil {
98+
log.Fatalf("translate %q: %v", phrase, err)
99+
}
100+
fmt.Printf("%-35s → %s\n", phrase, result.Latin)
101+
}
102+
}
103+
104+
func knownPhrases() string {
105+
phrases := make([]string, 0, len(translations))
106+
for k := range translations {
107+
phrases = append(phrases, fmt.Sprintf("%q", k))
108+
}
109+
return strings.Join(phrases, ", ")
110+
}

mcp/client.go

Lines changed: 79 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@ import (
1010
"fmt"
1111
"iter"
1212
"log/slog"
13+
"maps"
14+
"reflect"
1315
"slices"
1416
"strings"
1517
"sync"
@@ -32,6 +34,11 @@ type Client struct {
3234
sessions []*ClientSession
3335
sendingMethodHandler_ MethodHandler
3436
receivingMethodHandler_ MethodHandler
37+
// sendMethods is the list of methods this client may send to a
38+
// server: it always contains the standard server methods (from
39+
// serverMethodInfos) plus any custom methods registered via
40+
// [AddSendingCustomMethod].
41+
sendMethods map[string]methodInfo
3542
}
3643

3744
// NewClient creates a new [Client].
@@ -58,12 +65,16 @@ func NewClient(impl *Implementation, options *ClientOptions) *Client {
5865
opts.Logger = ensureLogger(nil)
5966
}
6067

68+
sendMethods := make(map[string]methodInfo, len(serverMethodInfos))
69+
maps.Copy(sendMethods, serverMethodInfos)
70+
6171
c := &Client{
6272
impl: impl,
6373
opts: opts,
6474
roots: newFeatureSet(func(r *Root) string { return r.URI }),
6575
sendingMethodHandler_: defaultSendingMethodHandler,
6676
receivingMethodHandler_: defaultReceivingMethodHandler[*ClientSession],
77+
sendMethods: sendMethods,
6778
}
6879
if opts.MultiRoundTrip == nil || !opts.MultiRoundTrip.Disabled {
6980
c.AddSendingMiddleware(clientMultiRoundTripMiddleware())
@@ -504,7 +515,7 @@ func injectRequestMeta[T any, P interface {
504515
Params
505516
}](cs *ClientSession, params P) P {
506517
res := cs.state.InitializeResult
507-
if params.isNil() {
518+
if params == nil {
508519
params = new(T)
509520
}
510521
m := params.GetMeta()
@@ -1153,7 +1164,9 @@ var clientMethodInfos = map[string]methodInfo{
11531164
}
11541165

11551166
func (cs *ClientSession) sendingMethodInfos() map[string]methodInfo {
1156-
return serverMethodInfos
1167+
cs.client.mu.Lock()
1168+
defer cs.client.mu.Unlock()
1169+
return cs.client.sendMethods
11571170
}
11581171

11591172
func (cs *ClientSession) receivingMethodInfos() map[string]methodInfo {
@@ -1597,3 +1610,67 @@ func paginate[P listParams, R listResult[T], T any](ctx context.Context, params
15971610
}
15981611
}
15991612
}
1613+
1614+
// AddSendingCustomMethod registers a custom JSON-RPC method
1615+
// that the client may send to the server.
1616+
//
1617+
// Registration is decoupled from invocation: extensions typically call this
1618+
// during setup, while the actual call site uses [CallCustomMethod].
1619+
//
1620+
// if err := mcp.AddSendingCustomMethod[*SearchParams, *SearchResult](c, "acme/search"); err != nil {
1621+
// return err
1622+
// }
1623+
// // ... later, anywhere a *ClientSession is available:
1624+
// result, err := mcp.CallCustomMethod[*SearchParams, *SearchResult](
1625+
// ctx, cs, "acme/search", &SearchParams{Query: "hello"})
1626+
//
1627+
// AddSendingCustomMethod returns an error if method is the name of a standard
1628+
// MCP method. Registering the same method twice replaces the previous
1629+
// registration.
1630+
func AddSendingCustomMethod[P paramsPtr[PT], R Result, PT any](
1631+
c *Client,
1632+
method string,
1633+
) error {
1634+
if _, ok := serverMethodInfos[method]; ok {
1635+
return fmt.Errorf("mcp: AddSendingCustomMethod: %q shadows a standard MCP method", method)
1636+
}
1637+
1638+
mi := methodInfo{
1639+
newResult: func() Result {
1640+
return reflect.New(reflect.TypeFor[R]().Elem()).Interface().(R)
1641+
},
1642+
}
1643+
1644+
c.mu.Lock()
1645+
defer c.mu.Unlock()
1646+
c.sendMethods[method] = mi
1647+
return nil
1648+
}
1649+
1650+
// CallCustomMethod sends a custom (non-standard) JSON-RPC method to the
1651+
// server and decodes the response into R.
1652+
//
1653+
// The method must have been registered on the session's client via
1654+
// [AddSendingCustomMethod].
1655+
func CallCustomMethod[P paramsPtr[PT], R Result, PT any](
1656+
ctx context.Context,
1657+
cs *ClientSession,
1658+
method string,
1659+
params P,
1660+
) (R, error) {
1661+
c := cs.client
1662+
c.mu.Lock()
1663+
_, ok := c.sendMethods[method]
1664+
c.mu.Unlock()
1665+
if !ok {
1666+
var zero R
1667+
return zero, fmt.Errorf("mcp: CallCustomMethod: %q is not registered; call AddSendingCustomMethod first", method)
1668+
}
1669+
if cs.usesNewProtocol() {
1670+
params = injectRequestMeta(cs, params)
1671+
}
1672+
return handleSend[R](ctx, method, &ClientRequest[P]{
1673+
Session: cs,
1674+
Params: params,
1675+
})
1676+
}

mcp/mcp_test.go

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3233,3 +3233,145 @@ func TestSubscriptionsListen_DisconnectScrubsMaps(t *testing.T) {
32333233
time.Sleep(10 * time.Millisecond)
32343234
}
32353235
}
3236+
3237+
func TestCustomMethods(t *testing.T) {
3238+
type searchParams struct {
3239+
ParamsBase
3240+
Query string `json:"query"`
3241+
Limit int `json:"limit,omitempty"`
3242+
}
3243+
3244+
type searchResult struct {
3245+
ResultBase
3246+
Hits []string `json:"hits"`
3247+
Total int `json:"total"`
3248+
}
3249+
3250+
callCustom := func(ctx context.Context, conn *jsonrpc2.Connection, method string, params, result any) error {
3251+
return conn.Call(ctx, method, params).Await(ctx, result)
3252+
}
3253+
3254+
ctx := context.Background()
3255+
s := NewServer(testImpl, nil)
3256+
3257+
if err := AddReceivingCustomMethod(s, "acme/search", func(ctx context.Context, ss *ServerSession, params *searchParams) (*searchResult, error) {
3258+
hits := []string{"result for " + params.Query}
3259+
return &searchResult{
3260+
Hits: hits,
3261+
Total: len(hits),
3262+
}, nil
3263+
}); err != nil {
3264+
t.Fatal(err)
3265+
}
3266+
3267+
ct, st := NewInMemoryTransports()
3268+
ss, err := s.Connect(ctx, st, nil)
3269+
if err != nil {
3270+
t.Fatal(err)
3271+
}
3272+
t.Cleanup(func() { _ = ss.Close() })
3273+
3274+
c := NewClient(testImpl, nil)
3275+
if err := AddSendingCustomMethod[*searchParams, *searchResult](c, "acme/search"); err != nil {
3276+
t.Fatal(err)
3277+
}
3278+
3279+
cs, err := c.Connect(ctx, ct, nil)
3280+
if err != nil {
3281+
t.Fatal(err)
3282+
}
3283+
t.Cleanup(func() { _ = cs.Close() })
3284+
3285+
// Test raw JSON-RPC call.
3286+
var result1 searchResult
3287+
if err := callCustom(ctx, cs.getConn(), "acme/search", &searchParams{Query: "hello", Limit: 10}, &result1); err != nil {
3288+
t.Fatal(err)
3289+
}
3290+
if len(result1.Hits) != 1 || result1.Hits[0] != "result for hello" {
3291+
t.Errorf("raw call: unexpected hits: %v", result1.Hits)
3292+
}
3293+
if result1.Total != 1 {
3294+
t.Errorf("raw call: unexpected total: %d", result1.Total)
3295+
}
3296+
3297+
// Test the typed CallCustomMethod helper.
3298+
result2, err := CallCustomMethod[*searchParams, *searchResult](
3299+
ctx, cs, "acme/search", &searchParams{Query: "world"})
3300+
if err != nil {
3301+
t.Fatal(err)
3302+
}
3303+
if len(result2.Hits) != 1 || result2.Hits[0] != "result for world" {
3304+
t.Errorf("CallCustomMethod: unexpected hits: %v", result2.Hits)
3305+
}
3306+
if result2.Total != 1 {
3307+
t.Errorf("CallCustomMethod: unexpected total: %d", result2.Total)
3308+
}
3309+
3310+
// CallCustomMethod must reject methods that were never registered.
3311+
if _, err := CallCustomMethod[*searchParams, *searchResult](
3312+
ctx, cs, "acme/unregistered", &searchParams{Query: "x"}); err == nil {
3313+
t.Error("CallCustomMethod: expected error for unregistered method, got nil")
3314+
}
3315+
}
3316+
3317+
func TestAddCustomMethodRejectsStandardMethods(t *testing.T) {
3318+
type params struct{ ParamsBase }
3319+
type result struct{ ResultBase }
3320+
3321+
t.Run("server", func(t *testing.T) {
3322+
s := NewServer(testImpl, nil)
3323+
err := AddReceivingCustomMethod(s, "tools/call",
3324+
func(ctx context.Context, ss *ServerSession, p *params) (*result, error) {
3325+
return &result{}, nil
3326+
})
3327+
if err == nil {
3328+
t.Fatal("AddReceivingCustomMethod: expected error when shadowing a standard method, got nil")
3329+
}
3330+
})
3331+
3332+
t.Run("client", func(t *testing.T) {
3333+
c := NewClient(testImpl, nil)
3334+
if err := AddSendingCustomMethod[*params, *result](c, "tools/call"); err == nil {
3335+
t.Fatal("AddSendingCustomMethod: expected error when shadowing a standard method, got nil")
3336+
}
3337+
})
3338+
}
3339+
3340+
// TestCallCustomMethodTypedNilParams exercises the typed-nil params path.
3341+
// User param types embed ParamsBase, so the inherited isNil forwarder would
3342+
// dereference a typed-nil outer if injectRequestMeta were called with it.
3343+
// CallCustomMethod must allocate a fresh value before the meta-injection step.
3344+
func TestCallCustomMethodTypedNilParams(t *testing.T) {
3345+
type pingParams struct{ ParamsBase }
3346+
type pingResult struct{ ResultBase }
3347+
3348+
ctx := context.Background()
3349+
s := NewServer(testImpl, nil)
3350+
if err := AddReceivingCustomMethod(s, "acme/ping",
3351+
func(ctx context.Context, ss *ServerSession, p *pingParams) (*pingResult, error) {
3352+
return &pingResult{}, nil
3353+
}); err != nil {
3354+
t.Fatal(err)
3355+
}
3356+
ct, st := NewInMemoryTransports()
3357+
ss, err := s.Connect(ctx, st, nil)
3358+
if err != nil {
3359+
t.Fatal(err)
3360+
}
3361+
t.Cleanup(func() { _ = ss.Close() })
3362+
3363+
c := NewClient(testImpl, nil)
3364+
if err := AddSendingCustomMethod[*pingParams, *pingResult](c, "acme/ping"); err != nil {
3365+
t.Fatal(err)
3366+
}
3367+
cs, err := c.Connect(ctx, ct, &ClientSessionOptions{protocolVersion: protocolVersion20260728})
3368+
if err != nil {
3369+
t.Fatal(err)
3370+
}
3371+
t.Cleanup(func() { _ = cs.Close() })
3372+
3373+
var typedNil *pingParams
3374+
if _, err := CallCustomMethod[*pingParams, *pingResult](ctx, cs, "acme/ping", typedNil); err != nil {
3375+
t.Fatalf("CallCustomMethod with typed-nil params: %v", err)
3376+
}
3377+
}

0 commit comments

Comments
 (0)