Skip to content

Commit d619a3b

Browse files
sambhavclaude
andcommitted
mcp: bidirectional custom methods, CustomMethod type, and extension registry
Builds on the custom JSON-RPC method support in #956 in three ways: 1. **Bidirectionality** — the server can now send requests to the client, and the client can register handlers for them, mirroring the existing client→server direction: - AddServerSendingCustomMethod / ServerCallCustomMethod - AddClientReceivingCustomMethod - CustomMethod.RegisterServerSending / .ServerCall - CustomMethod.RegisterClientReceiving 2. **CustomMethod[P, R, T]** — a phantom-type wrapper that captures the method name and parameter/result types once at package level, so call sites never repeat generic arguments or method-name strings: var Method = mcp.NewCustomMethod[*Params, *Result]("acme/method") result, err := Method.Call(ctx, cs, &Params{...}) 3. **Extension registry** — a mechanism for libraries to auto-wire custom methods into every Server/Client without requiring callers to do any manual setup: - RegisterExtension (global, typically called from init) - ServerOptions.Extensions / ClientOptions.Extensions (per-instance) Global extensions are applied first; per-instance extensions after (last writer wins on name collision). The example is restructured to demonstrate the pattern: a latinext sub-package is the "extension author" (defines types, registers via init(), exports a Translate() helper); main.go is the "consumer" (just imports latinext, no generics or method-name strings visible). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 65ece59 commit d619a3b

5 files changed

Lines changed: 373 additions & 70 deletions

File tree

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
// Copyright 2025 The Go MCP SDK Authors. All rights reserved.
2+
// Use of this source code is governed by an MIT-style
3+
// license that can be found in the LICENSE file.
4+
5+
// Package latinext is an example MCP extension that adds a "latin/translate"
6+
// custom JSON-RPC method. It demonstrates the extension-author pattern: types,
7+
// the CustomMethod variable, and the init() registration are all defined here
8+
// so that importers get everything wired up automatically.
9+
package latinext
10+
11+
import (
12+
"context"
13+
"fmt"
14+
"strings"
15+
16+
"github.com/modelcontextprotocol/go-sdk/mcp"
17+
)
18+
19+
// TranslateParams are the parameters for the latin/translate method.
20+
type TranslateParams struct {
21+
mcp.ParamsBase
22+
Text string `json:"text"`
23+
}
24+
25+
// TranslateResult is the result of the latin/translate method.
26+
type TranslateResult struct {
27+
mcp.ResultBase
28+
Latin string `json:"latin"`
29+
}
30+
31+
// Method captures the method name and types once. Extension consumers call
32+
// Translate() rather than this directly.
33+
var Method = mcp.NewCustomMethod[*TranslateParams, *TranslateResult]("latin/translate")
34+
35+
func init() {
36+
mcp.RegisterExtension(mcp.Extension{
37+
Server: func(s *mcp.Server) error {
38+
return Method.RegisterServerReceiving(s, DefaultHandler)
39+
},
40+
Client: func(c *mcp.Client) error {
41+
return Method.RegisterClientSending(c)
42+
},
43+
})
44+
}
45+
46+
// Translate calls the latin/translate method on the server via cs.
47+
// This is the one-liner that extension consumers use — no generics, no method
48+
// name strings.
49+
func Translate(ctx context.Context, cs *mcp.ClientSession, text string) (*TranslateResult, error) {
50+
return Method.Call(ctx, cs, &TranslateParams{Text: text})
51+
}
52+
53+
// DefaultHandler is the reference server-side implementation. It can be
54+
// overridden per-server by calling Method.RegisterServer(server, myHandler).
55+
func DefaultHandler(_ context.Context, _ *mcp.ServerSession, params *TranslateParams) (*TranslateResult, error) {
56+
key := strings.ToLower(strings.TrimSpace(params.Text))
57+
latin, ok := translations[key]
58+
if !ok {
59+
latin = fmt.Sprintf("[unknown: %q]", params.Text)
60+
}
61+
return &TranslateResult{Latin: latin}, nil
62+
}
63+
64+
var translations = map[string]string{
65+
"hello": "salve",
66+
"goodbye": "vale",
67+
"thank you": "gratias tibi ago",
68+
"how are you": "quid agis",
69+
"good morning": "bonum mane",
70+
"good night": "bonam noctem",
71+
"friend": "amicus",
72+
"water": "aqua",
73+
"love": "amor",
74+
"war": "bellum",
75+
"peace": "pax",
76+
"truth": "veritas",
77+
"light": "lux",
78+
"time": "tempus",
79+
"life": "vita",
80+
"death": "mors",
81+
"star": "stella",
82+
"earth": "terra",
83+
"sea": "mare",
84+
"the die is cast": "alea iacta est",
85+
"i came i saw i conquered": "veni vidi vici",
86+
"seize the day": "carpe diem",
87+
}

examples/server/custom-method/main.go

Lines changed: 12 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -2,74 +2,31 @@
22
// Use of this source code is governed by the license
33
// that can be found in the LICENSE file.
44

5-
// The custom-method example demonstrates registering and calling a custom
6-
// JSON-RPC method that is not part of the standard MCP spec.
5+
// The custom-method example demonstrates the extension-author / extension-consumer
6+
// split for custom JSON-RPC methods.
77
//
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.
8+
// The latinext sub-package is the "extension author": it defines the types,
9+
// registers the method via init(), and exposes a domain-specific Translate()
10+
// helper. This file is the "extension consumer": importing latinext is all
11+
// that's needed to wire up both sides.
1112
package main
1213

1314
import (
1415
"context"
1516
"fmt"
1617
"log"
17-
"strings"
1818

19+
"github.com/modelcontextprotocol/go-sdk/examples/server/custom-method/latinext"
1920
"github.com/modelcontextprotocol/go-sdk/mcp"
2021
)
2122

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-
5723
func main() {
5824
ctx := context.Background()
5925

26+
// NewServer and NewClient automatically apply all extensions registered via
27+
// init() — including latinext's handler and sending registration.
6028
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-
}
29+
client := mcp.NewClient(&mcp.Implementation{Name: "latin-client", Version: "v1.0.0"}, nil)
7330

7431
ct, st := mcp.NewInMemoryTransports()
7532

@@ -79,32 +36,19 @@ func main() {
7936
}
8037
defer ss.Close()
8138

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-
8739
cs, err := client.Connect(ctx, ct, nil)
8840
if err != nil {
8941
log.Fatal(err)
9042
}
9143
defer cs.Close()
9244

45+
// Call the custom method — no generics, no method-name strings.
9346
phrases := []string{"Hello", "Seize the day", "Peace", "Truth", "I came I saw I conquered"}
9447
for _, phrase := range phrases {
95-
result, err := mcp.CallCustomMethod[*TranslateParams, *TranslateResult](
96-
ctx, cs, "latin/translate", &TranslateParams{Text: phrase})
48+
result, err := latinext.Translate(ctx, cs, phrase)
9749
if err != nil {
9850
log.Fatalf("translate %q: %v", phrase, err)
9951
}
10052
fmt.Printf("%-35s → %s\n", phrase, result.Latin)
10153
}
10254
}
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: 50 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,11 @@ type Client struct {
3939
// serverMethodInfos) plus any custom methods registered via
4040
// [AddSendingCustomMethod].
4141
sendMethods map[string]methodInfo
42+
// receiveMethods is the merged map of methods this client may receive from
43+
// a server: it always contains the standard client methods (from
44+
// clientMethodInfos) plus any custom methods registered via
45+
// [AddClientReceivingCustomMethod].
46+
receiveMethods map[string]methodInfo
4247
}
4348

4449
// NewClient creates a new [Client].
@@ -68,17 +73,23 @@ func NewClient(impl *Implementation, options *ClientOptions) *Client {
6873
sendMethods := make(map[string]methodInfo, len(serverMethodInfos))
6974
maps.Copy(sendMethods, serverMethodInfos)
7075

76+
receiveMethods := make(map[string]methodInfo, len(clientMethodInfos))
77+
maps.Copy(receiveMethods, clientMethodInfos)
78+
7179
c := &Client{
7280
impl: impl,
7381
opts: opts,
7482
roots: newFeatureSet(func(r *Root) string { return r.URI }),
7583
sendingMethodHandler_: defaultSendingMethodHandler,
7684
receivingMethodHandler_: defaultReceivingMethodHandler[*ClientSession],
7785
sendMethods: sendMethods,
86+
receiveMethods: receiveMethods,
7887
}
7988
if opts.MultiRoundTrip == nil || !opts.MultiRoundTrip.Disabled {
8089
c.AddSendingMiddleware(clientMultiRoundTripMiddleware())
8190
}
91+
applyExtensionsToClient(c)
92+
runExtensions(opts.Extensions, func(e Extension) func(*Client) error { return e.Client }, c)
8293
return c
8394
}
8495

@@ -204,6 +215,10 @@ type ClientOptions struct {
204215
// reset" guidance, letting a transient miss pass without tearing down an
205216
// otherwise live session. Has no effect unless KeepAlive is non-zero.
206217
KeepAliveFailureThreshold int
218+
// Extensions are applied to the client during [NewClient], after any
219+
// globally registered extensions (see [RegisterExtension]). Per-client
220+
// extensions override global ones for the same method names.
221+
Extensions []Extension
207222
}
208223

209224
// toolContextKeyType is the context key type for passing tool definitions
@@ -1170,7 +1185,9 @@ func (cs *ClientSession) sendingMethodInfos() map[string]methodInfo {
11701185
}
11711186

11721187
func (cs *ClientSession) receivingMethodInfos() map[string]methodInfo {
1173-
return clientMethodInfos
1188+
cs.client.mu.Lock()
1189+
defer cs.client.mu.Unlock()
1190+
return cs.client.receiveMethods
11741191
}
11751192

11761193
func (cs *ClientSession) handle(ctx context.Context, req *jsonrpc.Request) (any, error) {
@@ -1674,3 +1691,35 @@ func CallCustomMethod[P paramsPtr[PT], R Result, PT any](
16741691
Params: params,
16751692
})
16761693
}
1694+
1695+
// AddClientReceivingCustomMethod registers a handler for a custom
1696+
// (non-standard) JSON-RPC method that the client may receive from a server.
1697+
//
1698+
// When a server sends a request with the given method name, the params will be
1699+
// unmarshaled into P, the handler will be called, and the returned R will be
1700+
// marshaled as the JSON-RPC result.
1701+
//
1702+
// P and R must implement [Params] and [Result] respectively, which is most
1703+
// easily done by embedding [ParamsBase] and [ResultBase].
1704+
//
1705+
// AddClientReceivingCustomMethod returns an error if method is the name of a
1706+
// standard MCP method. Registering the same custom method twice replaces the
1707+
// previous handler.
1708+
func AddClientReceivingCustomMethod[P paramsPtr[T], R Result, T any](
1709+
c *Client,
1710+
method string,
1711+
handler func(ctx context.Context, cs *ClientSession, params P) (R, error),
1712+
) error {
1713+
if _, ok := clientMethodInfos[method]; ok {
1714+
return fmt.Errorf("mcp: AddClientReceivingCustomMethod: %q shadows a standard MCP method", method)
1715+
}
1716+
1717+
typed := typedClientMethodHandler[P, R](func(ctx context.Context, req *ClientRequest[P]) (R, error) {
1718+
return handler(ctx, req.Session, req.Params)
1719+
})
1720+
1721+
c.mu.Lock()
1722+
defer c.mu.Unlock()
1723+
c.receiveMethods[method] = newClientMethodInfo(typed, missingParamsOK)
1724+
return nil
1725+
}

0 commit comments

Comments
 (0)