mcp: allow per request Cacheable customization (#1203)

Fixes #1094
2026-09-01 20:57:51 +02:00 · 2026-09-01 20:57:51 +02:00 · 0d3036f785
commit 0d3036f785
parent cb0de64136
5 changed files with 167 additions and 44 deletions
--- a/docs/server.md
+++ b/docs/server.md
@ -975,27 +975,33 @@ results: clients use `ttlMs` as a freshness hint to reduce polling,
 and `cacheScope` (`"public"` or `"private"`) controls whether shared
 intermediaries may cache the response.

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

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

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

 ### Completion
--- a/internal/docs/server.src.md
+++ b/internal/docs/server.src.md
@ -374,27 +374,33 @@ results: clients use `ttlMs` as a freshness hint to reduce polling,
 and `cacheScope` (`"public"` or `"private"`) controls whether shared
 intermediaries may cache the response.

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

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

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

 ### Completion
--- a/mcp/protocol.go
+++ b/mcp/protocol.go
@ -1191,9 +1191,16 @@ func (c Cacheable) GetTTLMs() int { return c.TTLMs }
 // GetCacheScope returns the cache scope.
 func (c Cacheable) GetCacheScope() string { return c.CacheScope }

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

 // The server's response to a prompts/list request from the client.
--- a/mcp/server.go
+++ b/mcp/server.go
@ -174,6 +174,14 @@ type ServerOptions struct {
 	// true, since stateless servers do not maintain sessions.
 	GetSessionID func() string

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

-func (s *Server) listPrompts(_ context.Context, req *ListPromptsRequest) (*ListPromptsResult, error) {
+func (s *Server) listPrompts(ctx context.Context, req *ListPromptsRequest) (*ListPromptsResult, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	if req.Params == nil {
@ -885,7 +893,7 @@ func (s *Server) listPrompts(_ context.Context, req *ListPromptsRequest) (*ListP
 	if err != nil {
 		return nil, err
 	}
-	res.setDefaultCacheableValues()
+	s.resolveCacheable(ctx, req, &res.Cacheable)
 	return res, nil
 }

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

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

-func (s *Server) listTools(_ context.Context, req *ListToolsRequest) (*ListToolsResult, error) {
+func (s *Server) listTools(ctx context.Context, req *ListToolsRequest) (*ListToolsResult, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	if req.Params == nil {
@ -983,7 +991,7 @@ func (s *Server) listTools(_ context.Context, req *ListToolsRequest) (*ListTools
 	if err != nil {
 		return nil, err
 	}
-	res.setDefaultCacheableValues()
+	s.resolveCacheable(ctx, req, &res.Cacheable)
 	return res, nil
 }

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

-func (s *Server) listResources(_ context.Context, req *ListResourcesRequest) (*ListResourcesResult, error) {
+func (s *Server) listResources(ctx context.Context, req *ListResourcesRequest) (*ListResourcesResult, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	if req.Params == nil {
@ -1031,11 +1039,11 @@ func (s *Server) listResources(_ context.Context, req *ListResourcesRequest) (*L
 	if err != nil {
 		return nil, err
 	}
-	res.setDefaultCacheableValues()
+	s.resolveCacheable(ctx, req, &res.Cacheable)
 	return res, nil
 }

-func (s *Server) listResourceTemplates(_ context.Context, req *ListResourceTemplatesRequest) (*ListResourceTemplatesResult, error) {
+func (s *Server) listResourceTemplates(ctx context.Context, req *ListResourceTemplatesRequest) (*ListResourceTemplatesResult, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	if req.Params == nil {
@ -1051,7 +1059,7 @@ func (s *Server) listResourceTemplates(_ context.Context, req *ListResourceTempl
 	if err != nil {
 		return nil, err
 	}
-	res.setDefaultCacheableValues()
+	s.resolveCacheable(ctx, req, &res.Cacheable)
 	return res, nil
 }

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

+// resolveCacheable settles the cache-control fields of an outgoing result.
+//
+// On entry c holds whatever the feature handler produced, which is the zero
+// value for results the SDK builds itself. [ServerOptions.SetCacheable] may
+// then rewrite it, and any cacheScope still unset is filled with the protocol
+// default.
+//
+// The list methods call this with s.mu held, so SetCacheable must not call
+// back into the server; see its documentation.
+func (s *Server) resolveCacheable(ctx context.Context, req Request, c *Cacheable) {
+	if s.opts.SetCacheable != nil {
+		s.opts.SetCacheable(ctx, req, c)
+	}
+	c.normalize()
+}
+
 // lookupResourceHandler returns the resource handler and MIME type for the resource or
 // resource template matching uri. If none, the last return value is false.
 func (s *Server) lookupResourceHandler(uri string) (ResourceHandler, string, bool) {
--- a/mcp/server_test.go
+++ b/mcp/server_test.go
@ -1692,6 +1692,86 @@ func TestServerSession_RejectsServerInitiated(t *testing.T) {
 	}
 }

+func TestServerSetCacheable(t *testing.T) {
+	ctx := context.Background()
+	// The values an opinionated resource handler sets on its own result.
+	handlerValue := Cacheable{TTLMs: 9_000, CacheScope: "private"}
+
+	for _, tc := range []struct {
+		name        string
+		callback    func(context.Context, Request, *Cacheable)
+		wantSDK     Cacheable // for a result whose fields nobody set
+		wantHandler Cacheable // for a result the handler set fields on
+	}{
+		{
+			name:        "no callback",
+			wantSDK:     Cacheable{TTLMs: 0, CacheScope: "public"},
+			wantHandler: handlerValue,
+		},
+		{
+			name: "callback overrules the handler",
+			callback: func(_ context.Context, _ Request, c *Cacheable) {
+				c.TTLMs, c.CacheScope = 60_000, "private"
+			},
+			wantSDK:     Cacheable{TTLMs: 60_000, CacheScope: "private"},
+			wantHandler: Cacheable{TTLMs: 60_000, CacheScope: "private"},
+		},
+		{
+			// The callback is handed what the handler produced, so it can
+			// single out the results that expressed no opinion. The scope it
+			// leaves empty is still filled with the protocol default.
+			name: "callback defers to the handler",
+			callback: func(_ context.Context, _ Request, c *Cacheable) {
+				if c.CacheScope == "" {
+					c.TTLMs = 1_000
+				}
+			},
+			wantSDK:     Cacheable{TTLMs: 1_000, CacheScope: "public"},
+			wantHandler: handlerValue,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			s := NewServer(testImpl, &ServerOptions{SetCacheable: tc.callback})
+			s.AddResource(&Resource{URI: "test://plain", Name: "plain"},
+				func(_ context.Context, req *ReadResourceRequest) (*ReadResourceResult, error) {
+					return &ReadResourceResult{Contents: []*ResourceContents{{URI: req.Params.URI, Text: "x"}}}, nil
+				})
+			s.AddResource(&Resource{URI: "test://opinionated", Name: "opinionated"},
+				func(_ context.Context, req *ReadResourceRequest) (*ReadResourceResult, error) {
+					return &ReadResourceResult{
+						Cacheable: handlerValue,
+						Contents:  []*ResourceContents{{URI: req.Params.URI, Text: "x"}},
+					}, nil
+				})
+			cs := mustConnect(t, s, nil)
+
+			check := func(label string, got, want Cacheable) {
+				t.Helper()
+				if got != want {
+					t.Errorf("%s Cacheable = %+v, want %+v", label, got, want)
+				}
+			}
+
+			tools, err := cs.ListTools(ctx, nil) // built by the SDK, never by a handler
+			if err != nil {
+				t.Fatal(err)
+			}
+			check("tools/list", tools.Cacheable, tc.wantSDK)
+
+			for uri, want := range map[string]Cacheable{
+				"test://plain":       tc.wantSDK,
+				"test://opinionated": tc.wantHandler,
+			} {
+				res, err := cs.ReadResource(ctx, &ReadResourceParams{URI: uri})
+				if err != nil {
+					t.Fatal(err)
+				}
+				check(uri, res.Cacheable, want)
+			}
+		})
+	}
+}
+
 // TestServerSupportedProtocolVersions verifies that
 // [ServerOptions.SupportedProtocolVersions] narrows both the versions the
 // server advertises in server/discover and the versions it negotiates,