Skip to content

Commit 0d3036f

Browse files
guglielmo-sanguglielmo-san
authored
mcp: allow per request Cacheable customization (#1203)
Fixes #1094
1 parent cb0de64 commit 0d3036f

5 files changed

Lines changed: 167 additions & 44 deletions

File tree

‎docs/server.md‎

Lines changed: 21 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -975,27 +975,33 @@ results: clients use `ttlMs` as a freshness hint to reduce polling,
975975
and `cacheScope` (`"public"` or `"private"`) controls whether shared
976976
intermediaries may cache the response.
977977

978-
**Server-side defaults.** After paginating, the SDK sets `CacheScope = "public"`
979-
and leaves `TTLMs = 0` (which the client cache treats as "immediately
980-
stale"). A handler that wants its responses cached must set `TTLMs`
981-
explicitly on the returned result.
978+
**Server-side defaults.** The SDK fills in `CacheScope = "public"` whenever a
979+
result leaves it empty, since the field is required on the wire. Everything
980+
else is left as produced: a `ResourceHandler` that sets `TTLMs` or `CacheScope`
981+
on its `ReadResourceResult` keeps those values. The list results and
982+
`server/discover` are built by the SDK itself, so they start out with
983+
`TTLMs = 0`, which a client cache treats as "immediately stale".
984+
985+
**Setting a policy.** `ServerOptions.SetCacheable` decides the fields for every
986+
result that carries them. It runs after the corresponding handler, receives
987+
what that handler produced, and may keep, adjust, or replace it. The request is
988+
passed in so that policy can vary by method, session, or authorization context.
982989

983990
```go
984-
mcp.AddTool(server, &mcp.Tool{Name: "expensive"}, func(...) (*mcp.CallToolResult, any, error) {
985-
// ...
986-
})
987-
// Override the default for tools/list:
988-
server.AddSendingMiddleware(func(next mcp.MethodHandler) mcp.MethodHandler {
989-
return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) {
990-
res, err := next(ctx, method, req)
991-
if lr, ok := res.(*mcp.ListToolsResult); ok {
992-
lr.TTLMs = 30_000 // 30s cache freshness hint
991+
server := mcp.NewServer(impl, &mcp.ServerOptions{
992+
SetCacheable: func(_ context.Context, req mcp.Request, c *mcp.Cacheable) {
993+
// A 30s freshness hint, except where the handler chose its own.
994+
if c.TTLMs == 0 {
995+
c.TTLMs = 30_000
993996
}
994-
return res, err
995-
}
997+
},
996998
})
997999
```
9981000

1001+
`SetCacheable` may run with the server's lock held, so it must not call back
1002+
into the `Server`: adding or removing a feature, or ranging over
1003+
`Server.Sessions`, deadlocks.
1004+
9991005
## Utilities
10001006

10011007
### Completion

‎internal/docs/server.src.md‎

Lines changed: 21 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -374,27 +374,33 @@ results: clients use `ttlMs` as a freshness hint to reduce polling,
374374
and `cacheScope` (`"public"` or `"private"`) controls whether shared
375375
intermediaries may cache the response.
376376

377-
**Server-side defaults.** After paginating, the SDK sets `CacheScope = "public"`
378-
and leaves `TTLMs = 0` (which the client cache treats as "immediately
379-
stale"). A handler that wants its responses cached must set `TTLMs`
380-
explicitly on the returned result.
377+
**Server-side defaults.** The SDK fills in `CacheScope = "public"` whenever a
378+
result leaves it empty, since the field is required on the wire. Everything
379+
else is left as produced: a `ResourceHandler` that sets `TTLMs` or `CacheScope`
380+
on its `ReadResourceResult` keeps those values. The list results and
381+
`server/discover` are built by the SDK itself, so they start out with
382+
`TTLMs = 0`, which a client cache treats as "immediately stale".
383+
384+
**Setting a policy.** `ServerOptions.SetCacheable` decides the fields for every
385+
result that carries them. It runs after the corresponding handler, receives
386+
what that handler produced, and may keep, adjust, or replace it. The request is
387+
passed in so that policy can vary by method, session, or authorization context.
381388

382389
```go
383-
mcp.AddTool(server, &mcp.Tool{Name: "expensive"}, func(...) (*mcp.CallToolResult, any, error) {
384-
// ...
385-
})
386-
// Override the default for tools/list:
387-
server.AddSendingMiddleware(func(next mcp.MethodHandler) mcp.MethodHandler {
388-
return func(ctx context.Context, method string, req mcp.Request) (mcp.Result, error) {
389-
res, err := next(ctx, method, req)
390-
if lr, ok := res.(*mcp.ListToolsResult); ok {
391-
lr.TTLMs = 30_000 // 30s cache freshness hint
390+
server := mcp.NewServer(impl, &mcp.ServerOptions{
391+
SetCacheable: func(_ context.Context, req mcp.Request, c *mcp.Cacheable) {
392+
// A 30s freshness hint, except where the handler chose its own.
393+
if c.TTLMs == 0 {
394+
c.TTLMs = 30_000
392395
}
393-
return res, err
394-
}
396+
},
395397
})
396398
```
397399

400+
`SetCacheable` may run with the server's lock held, so it must not call back
401+
into the `Server`: adding or removing a feature, or ranging over
402+
`Server.Sessions`, deadlocks.
403+
398404
## Utilities
399405

400406
### Completion

‎mcp/protocol.go‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1191,9 +1191,16 @@ func (c Cacheable) GetTTLMs() int { return c.TTLMs }
11911191
// GetCacheScope returns the cache scope.
11921192
func (c Cacheable) GetCacheScope() string { return c.CacheScope }
11931193

1194-
// setDefaultCacheableValues sets the default values for the cacheable fields.
1195-
func (c *Cacheable) setDefaultCacheableValues() {
1196-
c.CacheScope = "public"
1194+
// normalize fills in the protocol default for any cache field left
1195+
// unset. An absent cacheScope means "public", but the field is required on the
1196+
// wire, so the default is materialized here rather than sent empty.
1197+
//
1198+
// Values already present are preserved: this must not undo a decision made by
1199+
// a resource handler or by [ServerOptions.SetCacheable].
1200+
func (c *Cacheable) normalize() {
1201+
if c.CacheScope == "" {
1202+
c.CacheScope = "public"
1203+
}
11971204
}
11981205

11991206
// The server's response to a prompts/list request from the client.

‎mcp/server.go‎

Lines changed: 35 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,14 @@ type ServerOptions struct {
174174
// true, since stateless servers do not maintain sessions.
175175
GetSessionID func() string
176176

177+
// SetCacheable, if non-nil, decides the cache-control fields (ttlMs and
178+
// cacheScope) of every result that carries them: server/discover, the
179+
// four list methods, and resources/read.
180+
//
181+
// It is called once per such result, after the corresponding handler has
182+
// returned, with c holding the values that handler produced.
183+
SetCacheable func(ctx context.Context, req Request, c *Cacheable)
184+
177185
// SupportedProtocolVersions, if non-empty, restricts the protocol versions
178186
// this server advertises. If empty, every version returned by
179187
// [SupportedProtocolVersions] is used.
@@ -870,7 +878,7 @@ func (s *Server) Sessions() iter.Seq[*ServerSession] {
870878
return slices.Values(clients)
871879
}
872880

873-
func (s *Server) listPrompts(_ context.Context, req *ListPromptsRequest) (*ListPromptsResult, error) {
881+
func (s *Server) listPrompts(ctx context.Context, req *ListPromptsRequest) (*ListPromptsResult, error) {
874882
s.mu.Lock()
875883
defer s.mu.Unlock()
876884
if req.Params == nil {
@@ -885,7 +893,7 @@ func (s *Server) listPrompts(_ context.Context, req *ListPromptsRequest) (*ListP
885893
if err != nil {
886894
return nil, err
887895
}
888-
res.setDefaultCacheableValues()
896+
s.resolveCacheable(ctx, req, &res.Cacheable)
889897
return res, nil
890898
}
891899

@@ -915,7 +923,7 @@ func (s *Server) getPrompt(ctx context.Context, req *GetPromptRequest) (*GetProm
915923
// the server's capabilities, the server's identity, and the server's
916924
// instructions, allowing clients to negotiate without performing the legacy
917925
// initialize handshake.
918-
func (s *Server) discover(_ context.Context, req *ServerRequest[*DiscoverParams]) (*DiscoverResult, error) {
926+
func (s *Server) discover(ctx context.Context, req *ServerRequest[*DiscoverParams]) (*DiscoverResult, error) {
919927
req.Session.mu.Lock()
920928
versions := req.Session.supportedVersions
921929
req.Session.mu.Unlock()
@@ -947,7 +955,7 @@ func (s *Server) discover(_ context.Context, req *ServerRequest[*DiscoverParams]
947955
Capabilities: s.capabilities(),
948956
Instructions: s.opts.Instructions,
949957
}
950-
res.setDefaultCacheableValues()
958+
s.resolveCacheable(ctx, req, &res.Cacheable)
951959
return res, nil
952960
}
953961

@@ -968,7 +976,7 @@ func filterSupportedVersions(t Transport, versions []string) []string {
968976
return out
969977
}
970978

971-
func (s *Server) listTools(_ context.Context, req *ListToolsRequest) (*ListToolsResult, error) {
979+
func (s *Server) listTools(ctx context.Context, req *ListToolsRequest) (*ListToolsResult, error) {
972980
s.mu.Lock()
973981
defer s.mu.Unlock()
974982
if req.Params == nil {
@@ -983,7 +991,7 @@ func (s *Server) listTools(_ context.Context, req *ListToolsRequest) (*ListTools
983991
if err != nil {
984992
return nil, err
985993
}
986-
res.setDefaultCacheableValues()
994+
s.resolveCacheable(ctx, req, &res.Cacheable)
987995
return res, nil
988996
}
989997

@@ -1016,7 +1024,7 @@ func (s *Server) callTool(ctx context.Context, req *CallToolRequest) (*CallToolR
10161024
return res, err
10171025
}
10181026

1019-
func (s *Server) listResources(_ context.Context, req *ListResourcesRequest) (*ListResourcesResult, error) {
1027+
func (s *Server) listResources(ctx context.Context, req *ListResourcesRequest) (*ListResourcesResult, error) {
10201028
s.mu.Lock()
10211029
defer s.mu.Unlock()
10221030
if req.Params == nil {
@@ -1031,11 +1039,11 @@ func (s *Server) listResources(_ context.Context, req *ListResourcesRequest) (*L
10311039
if err != nil {
10321040
return nil, err
10331041
}
1034-
res.setDefaultCacheableValues()
1042+
s.resolveCacheable(ctx, req, &res.Cacheable)
10351043
return res, nil
10361044
}
10371045

1038-
func (s *Server) listResourceTemplates(_ context.Context, req *ListResourceTemplatesRequest) (*ListResourceTemplatesResult, error) {
1046+
func (s *Server) listResourceTemplates(ctx context.Context, req *ListResourceTemplatesRequest) (*ListResourceTemplatesResult, error) {
10391047
s.mu.Lock()
10401048
defer s.mu.Unlock()
10411049
if req.Params == nil {
@@ -1051,7 +1059,7 @@ func (s *Server) listResourceTemplates(_ context.Context, req *ListResourceTempl
10511059
if err != nil {
10521060
return nil, err
10531061
}
1054-
res.setDefaultCacheableValues()
1062+
s.resolveCacheable(ctx, req, &res.Cacheable)
10551063
return res, nil
10561064
}
10571065

@@ -1075,7 +1083,7 @@ func (s *Server) readResource(ctx context.Context, req *ReadResourceRequest) (*R
10751083
if err := handleMultiRoundTripResult(req.Session, s.opts.Logger, res); err != nil {
10761084
return nil, err
10771085
}
1078-
res.setDefaultCacheableValues()
1086+
s.resolveCacheable(ctx, req, &res.Cacheable)
10791087
if res.resultType == resultTypeInputRequired {
10801088
return res, nil
10811089
}
@@ -1094,6 +1102,22 @@ func (s *Server) readResource(ctx context.Context, req *ReadResourceRequest) (*R
10941102
return res, nil
10951103
}
10961104

1105+
// resolveCacheable settles the cache-control fields of an outgoing result.
1106+
//
1107+
// On entry c holds whatever the feature handler produced, which is the zero
1108+
// value for results the SDK builds itself. [ServerOptions.SetCacheable] may
1109+
// then rewrite it, and any cacheScope still unset is filled with the protocol
1110+
// default.
1111+
//
1112+
// The list methods call this with s.mu held, so SetCacheable must not call
1113+
// back into the server; see its documentation.
1114+
func (s *Server) resolveCacheable(ctx context.Context, req Request, c *Cacheable) {
1115+
if s.opts.SetCacheable != nil {
1116+
s.opts.SetCacheable(ctx, req, c)
1117+
}
1118+
c.normalize()
1119+
}
1120+
10971121
// lookupResourceHandler returns the resource handler and MIME type for the resource or
10981122
// resource template matching uri. If none, the last return value is false.
10991123
func (s *Server) lookupResourceHandler(uri string) (ResourceHandler, string, bool) {

‎mcp/server_test.go‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1692,6 +1692,86 @@ func TestServerSession_RejectsServerInitiated(t *testing.T) {
16921692
}
16931693
}
16941694

1695+
func TestServerSetCacheable(t *testing.T) {
1696+
ctx := context.Background()
1697+
// The values an opinionated resource handler sets on its own result.
1698+
handlerValue := Cacheable{TTLMs: 9_000, CacheScope: "private"}
1699+
1700+
for _, tc := range []struct {
1701+
name string
1702+
callback func(context.Context, Request, *Cacheable)
1703+
wantSDK Cacheable // for a result whose fields nobody set
1704+
wantHandler Cacheable // for a result the handler set fields on
1705+
}{
1706+
{
1707+
name: "no callback",
1708+
wantSDK: Cacheable{TTLMs: 0, CacheScope: "public"},
1709+
wantHandler: handlerValue,
1710+
},
1711+
{
1712+
name: "callback overrules the handler",
1713+
callback: func(_ context.Context, _ Request, c *Cacheable) {
1714+
c.TTLMs, c.CacheScope = 60_000, "private"
1715+
},
1716+
wantSDK: Cacheable{TTLMs: 60_000, CacheScope: "private"},
1717+
wantHandler: Cacheable{TTLMs: 60_000, CacheScope: "private"},
1718+
},
1719+
{
1720+
// The callback is handed what the handler produced, so it can
1721+
// single out the results that expressed no opinion. The scope it
1722+
// leaves empty is still filled with the protocol default.
1723+
name: "callback defers to the handler",
1724+
callback: func(_ context.Context, _ Request, c *Cacheable) {
1725+
if c.CacheScope == "" {
1726+
c.TTLMs = 1_000
1727+
}
1728+
},
1729+
wantSDK: Cacheable{TTLMs: 1_000, CacheScope: "public"},
1730+
wantHandler: handlerValue,
1731+
},
1732+
} {
1733+
t.Run(tc.name, func(t *testing.T) {
1734+
s := NewServer(testImpl, &ServerOptions{SetCacheable: tc.callback})
1735+
s.AddResource(&Resource{URI: "test://plain", Name: "plain"},
1736+
func(_ context.Context, req *ReadResourceRequest) (*ReadResourceResult, error) {
1737+
return &ReadResourceResult{Contents: []*ResourceContents{{URI: req.Params.URI, Text: "x"}}}, nil
1738+
})
1739+
s.AddResource(&Resource{URI: "test://opinionated", Name: "opinionated"},
1740+
func(_ context.Context, req *ReadResourceRequest) (*ReadResourceResult, error) {
1741+
return &ReadResourceResult{
1742+
Cacheable: handlerValue,
1743+
Contents: []*ResourceContents{{URI: req.Params.URI, Text: "x"}},
1744+
}, nil
1745+
})
1746+
cs := mustConnect(t, s, nil)
1747+
1748+
check := func(label string, got, want Cacheable) {
1749+
t.Helper()
1750+
if got != want {
1751+
t.Errorf("%s Cacheable = %+v, want %+v", label, got, want)
1752+
}
1753+
}
1754+
1755+
tools, err := cs.ListTools(ctx, nil) // built by the SDK, never by a handler
1756+
if err != nil {
1757+
t.Fatal(err)
1758+
}
1759+
check("tools/list", tools.Cacheable, tc.wantSDK)
1760+
1761+
for uri, want := range map[string]Cacheable{
1762+
"test://plain": tc.wantSDK,
1763+
"test://opinionated": tc.wantHandler,
1764+
} {
1765+
res, err := cs.ReadResource(ctx, &ReadResourceParams{URI: uri})
1766+
if err != nil {
1767+
t.Fatal(err)
1768+
}
1769+
check(uri, res.Cacheable, want)
1770+
}
1771+
})
1772+
}
1773+
}
1774+
16951775
// TestServerSupportedProtocolVersions verifies that
16961776
// [ServerOptions.SupportedProtocolVersions] narrows both the versions the
16971777
// server advertises in server/discover and the versions it negotiates,

0 commit comments

Comments
 (0)