mcp: allow per request Cacheable customization

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` runs without any server lock held, so it may call back into the
+server. Receiving middleware still runs after it, which remains the escape
+hatch for policy that has to inspect the result body.
+
 ## Utilities
 
 ### Completion

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` runs without any server lock held, so it may call back into the
+server. Receiving middleware still runs after it, which remains the escape
+hatch for policy that has to inspect the result body.
+
 ## Utilities
 
 ### Completion

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"
+// normalizeCacheable 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) normalizeCacheable() {
+	if c.CacheScope == "" {
+		c.CacheScope = "public"
+	}
 }
 
 // The server's response to a prompts/list request from the client.

mcp/server.go

@@ -169,6 +169,14 @@ type ServerOptions struct {
 	// GetSessionID is not consulted when [StreamableHTTPOptions.Stateless] is
 	// 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)
 }
 
 // NewServer creates a new MCP server. The resulting server has no features:
@@ -836,22 +844,22 @@ func (s *Server) Sessions() iter.Seq[*ServerSession] {
 	return slices.Values(clients)
 }
 
-func (s *Server) listPrompts(_ context.Context, req *ListPromptsRequest) (*ListPromptsResult, error) {
-	s.mu.Lock()
-	defer s.mu.Unlock()
+func (s *Server) listPrompts(ctx context.Context, req *ListPromptsRequest) (*ListPromptsResult, error) {
 	if req.Params == nil {
 		req.Params = &ListPromptsParams{}
 	}
+	s.mu.Lock()
 	res, err := paginateList(s.prompts, s.opts.PageSize, req.Params, &ListPromptsResult{}, func(res *ListPromptsResult, prompts []*serverPrompt) {
 		res.Prompts = []*Prompt{} // avoid JSON null
 		for _, p := range prompts {
 			res.Prompts = append(res.Prompts, p.prompt)
 		}
 	})
+	s.mu.Unlock()
 	if err != nil {
 		return nil, err
 	}
-	res.setDefaultCacheableValues()
+	s.resolveCacheable(ctx, req, &res.Cacheable)
 	return res, nil
 }
 
@@ -881,7 +889,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()
@@ -913,7 +921,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
 }
 
@@ -934,22 +942,22 @@ func filterSupportedVersions(t Transport) []string {
 	return out
 }
 
-func (s *Server) listTools(_ context.Context, req *ListToolsRequest) (*ListToolsResult, error) {
-	s.mu.Lock()
-	defer s.mu.Unlock()
+func (s *Server) listTools(ctx context.Context, req *ListToolsRequest) (*ListToolsResult, error) {
 	if req.Params == nil {
 		req.Params = &ListToolsParams{}
 	}
+	s.mu.Lock()
 	res, err := paginateList(s.tools, s.opts.PageSize, req.Params, &ListToolsResult{}, func(res *ListToolsResult, tools []*serverTool) {
 		res.Tools = []*Tool{} // avoid JSON null
 		for _, t := range tools {
 			res.Tools = append(res.Tools, t.tool)
 		}
 	})
+	s.mu.Unlock()
 	if err != nil {
 		return nil, err
 	}
-	res.setDefaultCacheableValues()
+	s.resolveCacheable(ctx, req, &res.Cacheable)
 	return res, nil
 }
 
@@ -982,42 +990,42 @@ func (s *Server) callTool(ctx context.Context, req *CallToolRequest) (*CallToolR
 	return res, err
 }
 
-func (s *Server) listResources(_ context.Context, req *ListResourcesRequest) (*ListResourcesResult, error) {
-	s.mu.Lock()
-	defer s.mu.Unlock()
+func (s *Server) listResources(ctx context.Context, req *ListResourcesRequest) (*ListResourcesResult, error) {
 	if req.Params == nil {
 		req.Params = &ListResourcesParams{}
 	}
+	s.mu.Lock()
 	res, err := paginateList(s.resources, s.opts.PageSize, req.Params, &ListResourcesResult{}, func(res *ListResourcesResult, resources []*serverResource) {
 		res.Resources = []*Resource{} // avoid JSON null
 		for _, r := range resources {
 			res.Resources = append(res.Resources, r.resource)
 		}
 	})
+	s.mu.Unlock()
 	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) {
-	s.mu.Lock()
-	defer s.mu.Unlock()
+func (s *Server) listResourceTemplates(ctx context.Context, req *ListResourceTemplatesRequest) (*ListResourceTemplatesResult, error) {
 	if req.Params == nil {
 		req.Params = &ListResourceTemplatesParams{}
 	}
+	s.mu.Lock()
 	res, err := paginateList(s.resourceTemplates, s.opts.PageSize, req.Params, &ListResourceTemplatesResult{},
 		func(res *ListResourceTemplatesResult, rts []*serverResourceTemplate) {
 			res.ResourceTemplates = []*ResourceTemplate{} // avoid JSON null
 			for _, rt := range rts {
 				res.ResourceTemplates = append(res.ResourceTemplates, rt.resourceTemplate)
 			}
 		})
+	s.mu.Unlock()
 	if err != nil {
 		return nil, err
 	}
-	res.setDefaultCacheableValues()
+	s.resolveCacheable(ctx, req, &res.Cacheable)
 	return res, nil
 }
 
@@ -1041,7 +1049,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
 	}
@@ -1060,6 +1068,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.
+//
+// Callers must not hold s.mu: SetCacheable is user code and may call back into
+// the server.
+func (s *Server) resolveCacheable(ctx context.Context, req Request, c *Cacheable) {
+	if s.opts.SetCacheable != nil {
+		s.opts.SetCacheable(ctx, req, c)
+	}
+	c.normalizeCacheable()
+}
+
 // 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) {

mcp/server_test.go

@@ -1691,3 +1691,83 @@ 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)
+			}
+		})
+	}
+}